"""
Advanced Agentic System
----------------------
A sophisticated multi-agent system with:

Core Components:
1. Agent Management
2. Task Execution
3. Learning & Adaptation
4. Communication
5. Resource Management

Advanced Features:
1. Self-Improvement
2. Multi-Agent Coordination
3. Dynamic Role Assignment
4. Emergent Behavior
"""

import logging
from typing import Dict, Any, List, Optional, Union, TypeVar
from dataclasses import dataclass, field
from enum import Enum
import json
import asyncio
from datetime import datetime
import uuid
from concurrent.futures import ThreadPoolExecutor
import numpy as np
from collections import defaultdict

from orchestrator import (
    AgentOrchestrator,
    AgentRole,
    AgentState,
    TaskPriority,
    Task
)
from reasoning import UnifiedReasoningEngine as ReasoningEngine, StrategyType as ReasoningMode
from reasoning.meta_learning import MetaLearningStrategy

class AgentCapability(Enum):
    """Core capabilities of agents."""
    REASONING = "reasoning"
    LEARNING = "learning"
    EXECUTION = "execution"
    COORDINATION = "coordination"
    MONITORING = "monitoring"

class AgentPersonality(Enum):
    """Different personality types for agents."""
    ANALYTICAL = "analytical"
    CREATIVE = "creative"
    CAUTIOUS = "cautious"
    PROACTIVE = "proactive"
    ADAPTIVE = "adaptive"

@dataclass
class AgentProfile:
    """Profile defining an agent's characteristics."""
    id: str
    name: str
    role: AgentRole
    capabilities: List[AgentCapability]
    personality: AgentPersonality
    expertise_areas: List[str]
    learning_rate: float
    risk_tolerance: float
    created_at: datetime
    metadata: Dict[str, Any]

class Agent:
    """Advanced autonomous agent with learning capabilities."""
    
    def __init__(
        self,
        profile: AgentProfile,
        reasoning_engine: ReasoningEngine,
        meta_learning: MetaLearningStrategy,
        config: Dict[str, Any] = None
    ):
        self.profile = profile
        self.reasoning_engine = reasoning_engine
        self.meta_learning = meta_learning
        self.config = config or {}
        
        # State management
        self.state = AgentState.IDLE
        self.current_task: Optional[Task] = None
        self.task_history: List[Task] = []
        
        # Learning and adaptation
        self.knowledge_base: Dict[str, Any] = {}
        self.learned_patterns: List[Dict[str, Any]] = []
        self.adaptation_history: List[Dict[str, Any]] = []
        
        # Performance metrics
        self.metrics: Dict[str, List[float]] = defaultdict(list)
        self.performance_history: List[Dict[str, float]] = []
        
        # Communication
        self.message_queue = asyncio.Queue()
        self.response_queue = asyncio.Queue()
        
        # Resource management
        self.resource_usage: Dict[str, float] = {}
        self.resource_limits: Dict[str, float] = {}
        
        # Async support
        self.executor = ThreadPoolExecutor(max_workers=2)
        self.lock = asyncio.Lock()
        
        # Logging
        self.logger = logging.getLogger(f"Agent-{profile.id}")
        
        # Initialize components
        self._init_components()

    def _init_components(self):
        """Initialize agent components."""
        # Set up knowledge base
        self.knowledge_base = {
            "expertise": {area: 0.5 for area in self.profile.expertise_areas},
            "learned_skills": set(),
            "interaction_patterns": defaultdict(int),
            "success_patterns": defaultdict(float)
        }
        
        # Set up resource limits
        self.resource_limits = {
            "cpu": 1.0,
            "memory": 1000,
            "api_calls": 100,
            "learning_capacity": 0.8
        }

    async def process_task(self, task: Task) -> Dict[str, Any]:
        """Process an assigned task."""
        try:
            self.current_task = task
            self.state = AgentState.BUSY
            
            # Analyze task
            analysis = await self._analyze_task(task)
            
            # Plan execution
            plan = await self._plan_execution(analysis)
            
            # Execute plan
            result = await self._execute_plan(plan)
            
            # Learn from execution
            await self._learn_from_execution(task, result)
            
            # Update metrics
            self._update_metrics(task, result)
            
            return {
                "success": True,
                "task_id": task.id,
                "result": result,
                "metrics": self._get_execution_metrics()
            }
            
        except Exception as e:
            self.logger.error(f"Error processing task: {e}")
            self.state = AgentState.ERROR
            return {
                "success": False,
                "task_id": task.id,
                "error": str(e)
            }
        finally:
            self.state = AgentState.IDLE
            self.current_task = None

    async def _analyze_task(self, task: Task) -> Dict[str, Any]:
        """Analyze task requirements and constraints."""
        # Use reasoning engine for analysis
        analysis = await self.reasoning_engine.reason(
            query=task.description,
            context={
                "agent_profile": self.profile.__dict__,
                "task_history": self.task_history,
                "knowledge_base": self.knowledge_base
            },
            mode=ReasoningMode.ANALYTICAL
        )
        
        return {
            "requirements": analysis.get("requirements", []),
            "constraints": analysis.get("constraints", []),
            "complexity": analysis.get("complexity", 0.5),
            "estimated_duration": analysis.get("estimated_duration", 3600),
            "required_capabilities": analysis.get("required_capabilities", [])
        }

    async def _plan_execution(self, analysis: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Plan task execution based on analysis."""
        # Use reasoning engine for planning
        plan = await self.reasoning_engine.reason(
            query="Plan execution steps",
            context={
                "analysis": analysis,
                "agent_capabilities": self.profile.capabilities,
                "resource_limits": self.resource_limits
            },
            mode=ReasoningMode.FOCUSED
        )
        
        return plan.get("steps", [])

    async def _execute_plan(self, plan: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Execute the planned steps."""
        results = []
        
        for step in plan:
            try:
                # Check resources
                if not self._check_resources(step):
                    raise RuntimeError("Insufficient resources for step execution")
                
                # Execute step
                step_result = await self._execute_step(step)
                results.append(step_result)
                
                # Update resource usage
                self._update_resource_usage(step)
                
                # Learn from step execution
                await self._learn_from_step(step, step_result)
                
            except Exception as e:
                self.logger.error(f"Error executing step: {e}")
                results.append({"error": str(e)})
        
        return {
            "success": all(r.get("success", False) for r in results),
            "results": results
        }

    async def _execute_step(self, step: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a single step of the plan."""
        step_type = step.get("type", "unknown")
        
        if step_type == "reasoning":
            return await self._execute_reasoning_step(step)
        elif step_type == "learning":
            return await self._execute_learning_step(step)
        elif step_type == "action":
            return await self._execute_action_step(step)
        else:
            raise ValueError(f"Unknown step type: {step_type}")

    async def _execute_reasoning_step(self, step: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a reasoning step."""
        result = await self.reasoning_engine.reason(
            query=step["query"],
            context=step.get("context", {}),
            mode=ReasoningMode.ANALYTICAL
        )
        
        return {
            "success": result.get("success", False),
            "reasoning_result": result
        }

    async def _execute_learning_step(self, step: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a learning step."""
        result = await self.meta_learning.learn(
            data=step["data"],
            context=step.get("context", {})
        )
        
        return {
            "success": result.get("success", False),
            "learning_result": result
        }

    async def _execute_action_step(self, step: Dict[str, Any]) -> Dict[str, Any]:
        """Execute an action step."""
        action_type = step.get("action_type")
        
        if action_type == "api_call":
            return await self._make_api_call(step)
        elif action_type == "data_processing":
            return await self._process_data(step)
        elif action_type == "coordination":
            return await self._coordinate_action(step)
        else:
            raise ValueError(f"Unknown action type: {action_type}")

    def _check_resources(self, step: Dict[str, Any]) -> bool:
        """Check if sufficient resources are available."""
        required_resources = step.get("required_resources", {})
        
        for resource, amount in required_resources.items():
            if self.resource_usage.get(resource, 0) + amount > self.resource_limits.get(resource, float('inf')):
                return False
        
        return True

    def _update_resource_usage(self, step: Dict[str, Any]):
        """Update resource usage after step execution."""
        used_resources = step.get("used_resources", {})
        
        for resource, amount in used_resources.items():
            self.resource_usage[resource] = self.resource_usage.get(resource, 0) + amount

    async def _learn_from_execution(self, task: Task, result: Dict[str, Any]):
        """Learn from task execution experience."""
        # Prepare learning data
        learning_data = {
            "task": task.__dict__,
            "result": result,
            "context": {
                "agent_state": self.state,
                "resource_usage": self.resource_usage,
                "performance_metrics": self._get_execution_metrics()
            }
        }
        
        # Learn patterns
        patterns = await self.meta_learning.learn(
            data=learning_data,
            context=self.knowledge_base
        )
        
        # Update knowledge base
        self._update_knowledge_base(patterns)
        
        # Record adaptation
        self.adaptation_history.append({
            "timestamp": datetime.now(),
            "patterns": patterns,
            "metrics": self._get_execution_metrics()
        })

    async def _learn_from_step(self, step: Dict[str, Any], result: Dict[str, Any]):
        """Learn from individual step execution."""
        if result.get("success", False):
            # Update success patterns
            pattern_key = f"{step['type']}:{step.get('action_type', 'none')}"
            self.knowledge_base["success_patterns"][pattern_key] += 1
            
            # Learn from successful execution
            await self.meta_learning.learn(
                data={
                    "step": step,
                    "result": result
                },
                context={"pattern_key": pattern_key}
            )

    def _update_knowledge_base(self, patterns: Dict[str, Any]):
        """Update knowledge base with new patterns."""
        # Update expertise levels
        for area, pattern in patterns.get("expertise_patterns", {}).items():
            if area in self.knowledge_base["expertise"]:
                current = self.knowledge_base["expertise"][area]
                self.knowledge_base["expertise"][area] = current * 0.9 + pattern * 0.1
        
        # Add new learned skills
        new_skills = patterns.get("learned_skills", set())
        self.knowledge_base["learned_skills"].update(new_skills)
        
        # Update interaction patterns
        for pattern, count in patterns.get("interaction_patterns", {}).items():
            self.knowledge_base["interaction_patterns"][pattern] += count

    def _update_metrics(self, task: Task, result: Dict[str, Any]):
        """Update performance metrics."""
        metrics = {
            "success": float(result.get("success", False)),
            "duration": (datetime.now() - task.created_at).total_seconds(),
            "resource_efficiency": self._calculate_resource_efficiency(),
            "learning_progress": self._calculate_learning_progress()
        }
        
        for key, value in metrics.items():
            self.metrics[key].append(value)
        
        self.performance_history.append({
            "timestamp": datetime.now(),
            "metrics": metrics
        })

    def _calculate_resource_efficiency(self) -> float:
        """Calculate resource usage efficiency."""
        if not self.resource_limits:
            return 1.0
            
        efficiencies = []
        for resource, usage in self.resource_usage.items():
            limit = self.resource_limits.get(resource, float('inf'))
            if limit > 0:
                efficiencies.append(1 - (usage / limit))
                
        return sum(efficiencies) / len(efficiencies) if efficiencies else 1.0

    def _calculate_learning_progress(self) -> float:
        """Calculate learning progress."""
        if not self.knowledge_base["expertise"]:
            return 0.0
            
        return sum(self.knowledge_base["expertise"].values()) / len(self.knowledge_base["expertise"])

    def _get_execution_metrics(self) -> Dict[str, float]:
        """Get current execution metrics."""
        return {
            key: sum(values[-10:]) / len(values[-10:])
            for key, values in self.metrics.items()
            if values
        }

class AgenticSystem:
    """Advanced multi-agent system with orchestration."""
    
    def __init__(self, config: Dict[str, Any] = None):
        self.config = config or {}
        
        # Initialize orchestrator
        self.orchestrator = AgentOrchestrator(config)
        
        # Initialize components
        self.agents: Dict[str, Agent] = {}
        self.reasoning_engine = ReasoningEngine(
            min_confidence=self.config.get('min_confidence', 0.7),
            parallel_threshold=self.config.get('parallel_threshold', 3),
            learning_rate=self.config.get('learning_rate', 0.1),
            strategy_weights=self.config.get('strategy_weights', {
                "LOCAL_LLM": 0.8,
                "CHAIN_OF_THOUGHT": 0.6,
                "TREE_OF_THOUGHTS": 0.5,
                "META_LEARNING": 0.4
            })
        )
        self.meta_learning = MetaLearningStrategy(config)
        
        # System state
        self.state = "initialized"
        self.metrics: Dict[str, List[float]] = defaultdict(list)
        
        # Async support
        self.executor = ThreadPoolExecutor(max_workers=4)
        self.lock = asyncio.Lock()
        
        # Logging
        self.logger = logging.getLogger("AgenticSystem")

    async def create_agent(
        self,
        name: str,
        role: AgentRole,
        capabilities: List[AgentCapability],
        personality: AgentPersonality,
        expertise_areas: List[str]
    ) -> str:
        """Create a new agent."""
        # Create agent profile
        profile = AgentProfile(
            id=str(uuid.uuid4()),
            name=name,
            role=role,
            capabilities=capabilities,
            personality=personality,
            expertise_areas=expertise_areas,
            learning_rate=0.1,
            risk_tolerance=0.5,
            created_at=datetime.now(),
            metadata={}
        )
        
        # Create agent instance
        agent = Agent(
            profile=profile,
            reasoning_engine=self.reasoning_engine,
            meta_learning=self.meta_learning,
            config=self.config.get("agent_config", {})
        )
        
        # Register with orchestrator
        agent_id = await self.orchestrator.register_agent(
            role=role,
            capabilities=[c.value for c in capabilities]
        )
        
        # Store agent
        async with self.lock:
            self.agents[agent_id] = agent
        
        return agent_id

    async def submit_task(
        self,
        description: str,
        priority: TaskPriority = TaskPriority.MEDIUM,
        deadline: Optional[datetime] = None
    ) -> str:
        """Submit a task to the system."""
        return await self.orchestrator.submit_task(
            description=description,
            priority=priority,
            deadline=deadline
        )

    async def get_task_status(self, task_id: str) -> Dict[str, Any]:
        """Get status of a task."""
        return await self.orchestrator.get_task_status(task_id)

    async def get_agent_status(self, agent_id: str) -> Dict[str, Any]:
        """Get status of an agent."""
        agent = self.agents.get(agent_id)
        if not agent:
            raise ValueError(f"Unknown agent: {agent_id}")
            
        return {
            "profile": agent.profile.__dict__,
            "state": agent.state,
            "current_task": agent.current_task.__dict__ if agent.current_task else None,
            "metrics": agent._get_execution_metrics(),
            "resource_usage": agent.resource_usage
        }

    async def get_system_status(self) -> Dict[str, Any]:
        """Get overall system status."""
        return {
            "state": self.state,
            "agent_count": len(self.agents),
            "active_tasks": len([a for a in self.agents.values() if a.state == AgentState.BUSY]),
            "performance_metrics": self._calculate_system_metrics(),
            "resource_usage": self._calculate_resource_usage()
        }

    def _calculate_system_metrics(self) -> Dict[str, float]:
        """Calculate overall system metrics."""
        metrics = defaultdict(list)
        
        for agent in self.agents.values():
            agent_metrics = agent._get_execution_metrics()
            for key, value in agent_metrics.items():
                metrics[key].append(value)
        
        return {
            key: sum(values) / len(values)
            for key, values in metrics.items()
            if values
        }

    def _calculate_resource_usage(self) -> Dict[str, float]:
        """Calculate overall resource usage."""
        usage = defaultdict(float)
        
        for agent in self.agents.values():
            for resource, amount in agent.resource_usage.items():
                usage[resource] += amount
        
        return dict(usage)