"""Adaptive learning system for continuous improvement of LlamaAgent.

This module implements sophisticated learning capabilities including:
- Multi-environment data collection
- Supervised fine-tuning (SFT)
- Reinforcement learning from human feedback (RLHF)
- Performance tracking and metrics
- Adaptive learning strategies

Author: Nik Jois <nikjois@llamasearch.ai>
"""

import asyncio
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional

try:
    import numpy as np
except ImportError:
    np = None

logger = logging.getLogger(__name__)
class LearningStrategy(Enum):
    """Learning strategy types."""
    SUPERVISED_FINE_TUNING = "supervised_fine_tuning"
    REINFORCEMENT_LEARNING = "reinforcement_learning"
    SELF_SUPERVISED = "self_supervised"
    MULTI_MODAL = "multi_modal"
    ACTIVE_LEARNING = "active_learning"


class DataEnvironment(Enum):
    """Data collection environments."""
    CODE_REPOSITORY = "code_repository"
    WEB_BROWSER = "web_browser"
    COMPUTER_DESKTOP = "computer_desktop"
    CHAT_CONVERSATIONS = "chat_conversations"
    DOCUMENTATION = "documentation"
    API_INTERACTIONS = "api_interactions"


class TaskStatus(Enum):
    """Task execution status."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class TrainingExample:
    """Individual training example with metadata."""
    input_data: str
    expected_output: str
    context: Optional[Dict[str, Any]] = None
    environment: Optional[DataEnvironment] = None
    quality_score: float = 1.0
    feedback: Optional[str] = None
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    metadata: Dict[str, Any] = field(default_factory=dict)
@dataclass
class PerformanceMetrics:
    """Performance metrics for tracking improvement."""
    accuracy: float = 0.0
    precision: float = 0.0
    recall: float = 0.0
    f1_score: float = 0.0
    task_completion_rate: float = 0.0
    average_response_time: float = 0.0
    user_satisfaction: float = 0.0
    error_rate: float = 0.0
    improvement_rate: float = 0.0
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
@dataclass
class LearningSession:
    """Individual learning session with specific objectives."""
    session_id: str
    strategy: LearningStrategy
    environment: DataEnvironment
    objective: str
    training_examples: List[TrainingExample] = field(default_factory=list)
    performance_before: Optional[PerformanceMetrics] = None
    performance_after: Optional[PerformanceMetrics] = None
    duration: float = 0.0
    success: bool = False
    metadata: Dict[str, Any] = field(default_factory=dict)
class DataCollectionAgent:
    """Agent for collecting training data from multiple environments."""
    
    def __init__(self):
        self.collected_examples: List[TrainingExample] = []
        self.collection_stats: Dict[str, Any] = {
            "total_collected": 0,
            "environments_covered": set(),
            "quality_distribution": {"high": 0, "medium": 0, "low": 0}
        }
    
    async def collect_from_code_repositories(self, repo_paths: List[str]) -> List[TrainingExample]:
        """Collect training data from code repositories."""
        examples: List[TrainingExample] = []
        
        for repo_path in repo_paths:
            try:
                repo_dir = Path(repo_path)
                if not repo_dir.exists():
                    continue
                
                # Collect Python files
                py_files = list(repo_dir.rglob("*.py"))
                for py_file in py_files[:10]:  # Limit for demo
                    try:
                        with open(py_file, 'r', encoding='utf-8') as f:
                            content = f.read()
                        
                        if len(content) > 100:  # Skip very small files
                            example = TrainingExample(
                                input_data=f"Analyze this Python code:\n{content[:500]}...",
                                expected_output=f"Code analysis: {py_file.name}",
                                context={"file_path": str(py_file), "file_type": "python"},
                                environment=DataEnvironment.CODE_REPOSITORY,
                                quality_score=0.8
                            )
                            examples.append(example)
                    except Exception as e:
                        logger.warning(f"Failed to process {py_file}: {e}")
            except Exception as e:
                logger.error(f"Failed to process repository {repo_path}: {e}")
        self.collected_examples.extend(examples)
        self.collection_stats["total_collected"] += len(examples)
        self.collection_stats["environments_covered"].add(DataEnvironment.CODE_REPOSITORY)
        return examples
    
    async def collect_from_web_interactions(self, urls: List[str]) -> List[TrainingExample]:
        """Collect training data from web interactions."""
        examples: List[TrainingExample] = []
        
        # Simulate web interaction tasks
        web_tasks = [
            ("Navigate to documentation", "Successfully accessed documentation"),
            ("Search for information", "Found relevant search results"),
            ("Extract key information", "Extracted important details"),
            ("Summarize content", "Created concise summary")
        ]
        
        for url in urls[:5]:  # Limit for demo
            for task, expected in web_tasks:
                example = TrainingExample(
                    input_data=f"Web task: {task} on {url}",
                    expected_output=expected,
                    context={"url": url, "task_type": "web_interaction"},
                    environment=DataEnvironment.WEB_BROWSER,
                    quality_score=0.7
                )
                examples.append(example)
        self.collected_examples.extend(examples)
        self.collection_stats["total_collected"] += len(examples)
        self.collection_stats["environments_covered"].add(DataEnvironment.WEB_BROWSER)
        return examples
    
    async def collect_from_computer_interactions(self) -> List[TrainingExample]:
        """Collect training data from computer desktop interactions."""
        examples: List[TrainingExample] = []
        
        # Simulate computer interaction tasks
        computer_tasks = [
            ("Open a text editor and create a new document", "Successfully created new document"),
            ("Navigate file system to find specific files", "Located target files efficiently"),
            ("Use command line tools for system tasks", "Executed commands successfully"),
            ("Manage application windows and workspaces", "Organized workspace effectively")
        ]
        
        for task, expected in computer_tasks:
            example = TrainingExample(
                input_data=f"Computer task: {task}",
                expected_output=expected,
                context={"task_type": "computer_interaction"},
                environment=DataEnvironment.COMPUTER_DESKTOP,
                quality_score=0.8
            )
            examples.append(example)
        self.collected_examples.extend(examples)
        self.collection_stats["total_collected"] += len(examples)
        self.collection_stats["environments_covered"].add(DataEnvironment.COMPUTER_DESKTOP)
        return examples


class SupervisedFineTuningEngine:
    """Engine for supervised fine-tuning of agent models."""
    
    def __init__(self):
        self.training_history: List[LearningSession] = []
        self.model_checkpoints: Dict[str, Any] = {}
    
    async def train_model(self, training_data: List[TrainingExample], epochs: int = 3) -> LearningSession:
        """Perform supervised fine-tuning."""
        
        session_id = f"sft_{int(time.time()}"
        session = LearningSession(
            session_id=session_id,
            strategy=LearningStrategy.SUPERVISED_FINE_TUNING,
            environment=DataEnvironment.CHAT_CONVERSATIONS,
            objective="Improve model performance through supervised fine-tuning",
            training_examples=training_data
        )
        start_time = time.time()
        
        try:
            # Simulate fine-tuning process
            logger.info(f"Starting SFT with {len(training_data)} examples")
            
            # Record baseline performance
            session.performance_before = await self._evaluate_model()
            
            # Simulate training epochs
            for epoch in range(epochs):
                logger.info(f"Training epoch {epoch + 1}/{epochs}")
                # Simulate batch processing
                batch_size = min(32, len(training_data)
                for i in range(0, len(training_data), batch_size):
                    batch = training_data[i:i + batch_size]
                    await self._process_training_batch(batch)
                # Simulate validation
                await asyncio.sleep(0.1)  # Simulate processing time
            
            # Record final performance
            session.performance_after = await self._evaluate_model()
            session.duration = time.time() - start_time
            session.success = True
            
            logger.info(f"SFT completed in {session.duration:.2f}s")
        except Exception as e:
            logger.error(f"SFT failed: {e}")
            session.success = False
            session.metadata["error"] = str(e)
        finally:
            self.training_history.append(session)
        return session
    
    async def _process_training_batch(self, batch: List[TrainingExample]) -> None:
        """Process a batch of training examples."""
        # Simulate batch processing
        await asyncio.sleep(0.05)
    async def _evaluate_model(self) -> PerformanceMetrics:
        """Evaluate current model performance."""
        # Simulate model evaluation
        if np:
            accuracy = np.random.uniform(0.7, 0.95)
            precision = np.random.uniform(0.65, 0.9)
            recall = np.random.uniform(0.6, 0.88)
        else:
            import random
            accuracy = random.uniform(0.7, 0.95)
            precision = random.uniform(0.65, 0.9)
            recall = random.uniform(0.6, 0.88)
        f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
        
        return PerformanceMetrics(
            accuracy=accuracy,
            precision=precision,
            recall=recall,
            f1_score=f1_score,
            task_completion_rate=accuracy,
            average_response_time=1.2,
            user_satisfaction=accuracy * 0.9,
            error_rate=1.0 - accuracy
        )
class ReinforcementLearningEngine:
    """Engine for reinforcement learning and RLHF."""
    
    def __init__(self):
        self.reward_model = RewardModel()
        self.policy_optimizer = PolicyOptimizer()
        self.training_history: List[LearningSession] = []
    
    async def train_with_feedback(self, interactions: List[Dict[str, Any]], num_iterations: int = 5) -> LearningSession:
        """Train using reinforcement learning from human feedback."""
        
        session_id = f"rlhf_{int(time.time()}"
        session = LearningSession(
            session_id=session_id,
            strategy=LearningStrategy.REINFORCEMENT_LEARNING,
            environment=DataEnvironment.CHAT_CONVERSATIONS,
            objective="Improve model behavior through human feedback",
            training_examples=[]
        )
        start_time = time.time()
        
        try:
            logger.info(f"Starting RLHF with {len(interactions)} feedback examples")
            
            # Train reward model
            await self.reward_model.train(interactions)
            # Policy optimization
            for iteration in range(num_iterations):
                logger.info(f"RLHF iteration {iteration + 1}/{num_iterations}")
                # Generate responses
                responses = await self._generate_responses()
                
                # Get rewards from reward model
                rewards = await self.reward_model.score_responses(responses)
                # Update policy
                await self.policy_optimizer.update(responses, rewards)
                await asyncio.sleep(0.1)  # Simulate processing time
            
            session.performance_after = await self._evaluate_rl_performance()
            session.duration = time.time() - start_time
            session.success = True
            
            logger.info(f"RLHF completed in {session.duration:.2f}s")
        except Exception as e:
            logger.error(f"RLHF failed: {e}")
            session.success = False
            session.metadata["error"] = str(e)
        finally:
            self.training_history.append(session)
        return session
    
    async def _generate_responses(self) -> List[str]:
        """Generate responses for policy optimization."""
        # Simulate response generation
        return [f"Response {i}" for i in range(10)]
    
    async def _evaluate_rl_performance(self) -> PerformanceMetrics:
        """Evaluate reinforcement learning performance."""
        # Simulate RL evaluation
        if np:
            reward_score = np.random.uniform(0.6, 0.9)
        else:
            import random
            reward_score = random.uniform(0.6, 0.9)
        return PerformanceMetrics(
            accuracy=reward_score,
            task_completion_rate=reward_score,
            user_satisfaction=reward_score,
            improvement_rate=0.1
        )
class RewardModel:
    """Reward model for RLHF training."""
    
    async def train(self, interactions: List[Dict[str, Any]]) -> None:
        """Train the reward model on human feedback."""
        logger.info(f"Training reward model with {len(interactions)} interactions")
        await asyncio.sleep(0.2)  # Simulate training
    
    async def score_responses(self, responses: List[str]) -> List[float]:
        """Score responses using the reward model."""
        # Simulate scoring
        if np:
            return [np.random.uniform(0.3, 1.0) for _ in responses]
        else:
            import random
            return [random.uniform(0.3, 1.0) for _ in responses]


class PolicyOptimizer:
    """Policy optimizer for RL training."""
    
    async def update(self, responses: List[str], rewards: List[float]) -> None:
        """Update policy based on responses and rewards."""
        logger.info(f"Updating policy with {len(responses)} responses")
        await asyncio.sleep(0.1)  # Simulate policy update


class PerformanceTracker:
    """Tracks performance metrics over time."""
    
    def __init__(self):
        self.metrics_history: List[PerformanceMetrics] = []
        self.baseline_metrics: Optional[PerformanceMetrics] = None
    
    async def record_metrics(self, metrics: PerformanceMetrics) -> None:
        """Record performance metrics."""
        # Calculate improvement rate
        if self.baseline_metrics:
            if self.baseline_metrics.accuracy > 0:
                metrics.improvement_rate = (
                    metrics.accuracy - self.baseline_metrics.accuracy
                ) / self.baseline_metrics.accuracy
        
        self.metrics_history.append(metrics)
        if not self.baseline_metrics:
            self.baseline_metrics = metrics
    
    def get_improvement_trends(self) -> Dict[str, float]:
        """Get improvement trends over time."""
        if len(self.metrics_history) < 2:
            return {}
        
        recent_metrics = self.metrics_history[-5:]  # Last 5 measurements
        
        if np:
            avg_accuracy = np.mean([m.accuracy for m in recent_metrics])
            avg_completion_rate = np.mean([m.task_completion_rate for m in recent_metrics])
            avg_satisfaction = np.mean([m.user_satisfaction for m in recent_metrics])
        else:
            avg_accuracy = sum(m.accuracy for m in recent_metrics) / len(recent_metrics)
            avg_completion_rate = sum(m.task_completion_rate for m in recent_metrics) / len(recent_metrics)
            avg_satisfaction = sum(m.user_satisfaction for m in recent_metrics) / len(recent_metrics)
        return {
            "average_accuracy": float(avg_accuracy),
            "average_completion_rate": float(avg_completion_rate),
            "average_satisfaction": float(avg_satisfaction),
            "total_sessions": len(self.metrics_history)
        }


class AdaptiveLearningSystem:
    """Main adaptive learning system coordinating all components."""
    
    def __init__(self):
        self.data_collector = DataCollectionAgent()
        self.sft_engine = SupervisedFineTuningEngine()
        self.rl_engine = ReinforcementLearningEngine()
        self.performance_tracker = PerformanceTracker()
        self.learning_pipeline_active = False
    
    async def run_comprehensive_learning(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """Run comprehensive learning pipeline with multi-environment data collection."""
        
        if self.learning_pipeline_active:
            raise RuntimeError("Learning pipeline already active")
        self.learning_pipeline_active = True
        pipeline_start = time.time()
        
        try:
            logger.info("Starting comprehensive adaptive learning pipeline")
            # Phase 1: Multi-environment data collection
            logger.info("Phase 1: Multi-environment data collection")
            all_examples: List[TrainingExample] = []
            
            # Collect from different environments
            if config.get("collect_code_data", True):
                repo_paths = config.get("repo_paths", [])
                if repo_paths:
                    code_examples = await self.data_collector.collect_from_code_repositories(repo_paths)
                    all_examples.extend(code_examples)
            if config.get("collect_web_data", True):
                urls = config.get("web_urls", [])
                if urls:
                    web_examples = await self.data_collector.collect_from_web_interactions(urls)
                    all_examples.extend(web_examples)
            if config.get("collect_computer_data", True):
                computer_examples = await self.data_collector.collect_from_computer_interactions()
                all_examples.extend(computer_examples)
            # Phase 2: Data filtering and quality assessment
            logger.info("Phase 2: Data filtering and quality assessment")
            high_quality_examples = [ex for ex in all_examples if ex.quality_score >= 0.8]
            
            # Phase 3: Supervised fine-tuning
            if config.get("enable_sft", True) and high_quality_examples:
                logger.info("Phase 3: Supervised fine-tuning")
                sft_session = await self.sft_engine.train_model(
                    high_quality_examples, 
                    epochs=config.get("sft_epochs", 3)
                )
                if sft_session.performance_after:
                    await self.performance_tracker.record_metrics(sft_session.performance_after)
            # Phase 4: Reinforcement learning (if enabled)
            if config.get("enable_rlhf", False):
                logger.info("Phase 4: Reinforcement learning from human feedback")
                mock_interactions = [
                    {"prompt": "Plan a project", "response": "I'll help you plan", "reward": 0.8},
                    {"prompt": "Write a Python function", "response": "def example():", "reward": 0.9}
                ]
                rl_session = await self.rl_engine.train_with_feedback(mock_interactions)
                if rl_session.performance_after:
                    await self.performance_tracker.record_metrics(rl_session.performance_after)
            # Phase 5: Performance analysis and insights
            logger.info("Phase 5: Performance analysis and insights")
            insights = await self.get_learning_insights()
            
            total_time = time.time() - pipeline_start
            
            return {
                "success": True,
                "total_examples_collected": len(all_examples),
                "high_quality_examples": len(high_quality_examples),
                "environments_covered": len(self.data_collector.collection_stats["environments_covered"]),
                "total_time": total_time,
                "insights": insights,
                "collection_stats": self.data_collector.collection_stats
            }
            
        except Exception as e:
            logger.error(f"Learning pipeline failed: {e}")
            return {
                "success": False,
                "error": str(e),
                "total_time": time.time() - pipeline_start
            }
        
        finally:
            self.learning_pipeline_active = False
    
    async def get_learning_insights(self) -> Dict[str, Any]:
        """Get insights from learning history and performance trends."""
        trends = self.performance_tracker.get_improvement_trends()
        
        insights = {
            "performance_trends": trends,
            "total_learning_sessions": len(self.sft_engine.training_history) + len(self.rl_engine.training_history),
            "data_collection_coverage": len(self.data_collector.collection_stats["environments_covered"]),
            "recommendations": [
                "Continue collecting diverse training examples",
                "Focus on high-quality data sources",
                "Monitor performance trends for optimal learning strategies"
            ]
        }
        
        return insights
    
    def get_system_status(self) -> Dict[str, Any]:
        """Get current system status."""
        return {
            "pipeline_active": self.learning_pipeline_active,
            "total_examples_collected": len(self.data_collector.collected_examples),
            "performance_history_length": len(self.performance_tracker.metrics_history),
            "sft_sessions": len(self.sft_engine.training_history),
            "rl_sessions": len(self.rl_engine.training_history),
            "baseline_established": self.performance_tracker.baseline_metrics is not None
        }


# Factory function
def create_adaptive_learning_system() -> AdaptiveLearningSystem:
    """Create and return an adaptive learning system instance."""
    return AdaptiveLearningSystem()