#!/usr/bin/env python3
"""
GENESIS MULTI-AGENT COORDINATOR
================================
Orchestrates multiple AI agents (Claude, Gemini, AIVA) working on parallel tasks.

Agent Capabilities:
    - Claude Opus: Complex reasoning, architecture, code review
    - Claude Sonnet: General implementation, balanced
    - Gemini 2.5 Pro: Large context (2M tokens), research
    - Gemini 2.0 Flash: Fast iterations, simple tasks
    - AIVA (Qwen): Zero-cost local execution, monitoring

Coordination Patterns:
    1. Round-Robin: Distribute tasks evenly
    2. Specialist: Route by task type
    3. Auction: Agents bid on tasks
    4. Hierarchical: Lead agent delegates
    5. Consensus: Multiple agents validate

Usage:
    coordinator = MultiAgentCoordinator()
    result = coordinator.execute_task_parallel(tasks)
"""

import asyncio
import json
import os
import sys
import time
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Dict, List, Any, Optional, Callable, Tuple
import threading
import queue

# Learning Loop integration for intelligent agent selection
try:
    from learning_loop import LearningLoop
    LEARNING_AVAILABLE = True
except ImportError:
    LEARNING_AVAILABLE = False
    LearningLoop = None


class AgentType(Enum):
    """Available agent types."""
    CLAUDE_OPUS = "claude-opus"
    CLAUDE_SONNET = "claude-sonnet"
    GEMINI_PRO = "gemini-2.5-pro"
    GEMINI_FLASH = "gemini-2.0-flash"
    AIVA_QWEN = "aiva-qwen"


class TaskComplexity(Enum):
    """Task complexity levels."""
    TRIVIAL = "trivial"      # Simple, one-step
    SIMPLE = "simple"        # Few steps, clear path
    MODERATE = "moderate"    # Multiple steps, some decisions
    COMPLEX = "complex"      # Many steps, architecture decisions
    CRITICAL = "critical"    # High stakes, needs review


class CoordinationPattern(Enum):
    """Patterns for multi-agent coordination."""
    ROUND_ROBIN = "round_robin"
    SPECIALIST = "specialist"
    AUCTION = "auction"
    HIERARCHICAL = "hierarchical"
    CONSENSUS = "consensus"
    PARALLEL = "parallel"


@dataclass
class AgentProfile:
    """Profile defining an agent's capabilities."""
    agent_type: AgentType
    name: str
    strengths: List[str]
    cost_per_token: float
    max_context: int
    speed_rating: float  # 0-1, higher = faster
    quality_rating: float  # 0-1, higher = better
    available: bool = True
    current_load: int = 0
    max_concurrent: int = 1

    def suitability_score(self, task_type: str, complexity: TaskComplexity) -> float:
        """Calculate suitability score for a task."""
        score = 0.5  # Base score

        # Check if task matches strengths
        for strength in self.strengths:
            if strength.lower() in task_type.lower():
                score += 0.2

        # Complexity matching
        if complexity == TaskComplexity.CRITICAL:
            score += self.quality_rating * 0.3
        elif complexity == TaskComplexity.TRIVIAL:
            score += self.speed_rating * 0.3
        else:
            score += (self.quality_rating + self.speed_rating) * 0.15

        # Availability penalty
        if self.current_load >= self.max_concurrent:
            score *= 0.5

        return min(score, 1.0)


@dataclass
class Task:
    """A task to be executed by an agent."""
    id: str
    title: str
    description: str
    task_type: str
    complexity: TaskComplexity = TaskComplexity.MODERATE
    priority: int = 5  # 1-10, higher = more urgent
    dependencies: List[str] = field(default_factory=list)
    assigned_agent: Optional[AgentType] = None
    status: str = "pending"
    result: Optional[Any] = None
    error: Optional[str] = None
    started_at: Optional[str] = None
    completed_at: Optional[str] = None
    cost: float = 0.0


@dataclass
class ExecutionResult:
    """Result of task execution."""
    task_id: str
    success: bool
    output: Any
    agent: AgentType
    duration: float
    cost: float
    error: Optional[str] = None


class BaseAgent(ABC):
    """Abstract base class for all agents."""

    def __init__(self, profile: AgentProfile):
        self.profile = profile
        self.execution_history: List[ExecutionResult] = []

    @abstractmethod
    def execute(self, task: Task) -> ExecutionResult:
        """Execute a task and return result."""
        pass

    @abstractmethod
    def health_check(self) -> bool:
        """Check if agent is available."""
        pass

    def get_stats(self) -> Dict:
        """Get execution statistics."""
        if not self.execution_history:
            return {"executions": 0, "success_rate": 0, "avg_duration": 0, "total_cost": 0}

        successes = sum(1 for r in self.execution_history if r.success)
        return {
            "executions": len(self.execution_history),
            "success_rate": successes / len(self.execution_history),
            "avg_duration": sum(r.duration for r in self.execution_history) / len(self.execution_history),
            "total_cost": sum(r.cost for r in self.execution_history)
        }


class GeminiAgent(BaseAgent):
    """Gemini-based agent using direct API."""

    def __init__(self, model: str = "gemini-2.0-flash-exp"):
        profile = AgentProfile(
            agent_type=AgentType.GEMINI_FLASH if "flash" in model else AgentType.GEMINI_PRO,
            name=f"Gemini-{model}",
            strengths=["research", "large_context", "fast_iteration", "code_generation"],
            cost_per_token=0.0001 if "flash" in model else 0.0005,
            max_context=2000000 if "pro" in model else 1000000,
            speed_rating=0.9 if "flash" in model else 0.7,
            quality_rating=0.7 if "flash" in model else 0.9,
            max_concurrent=5
        )
        super().__init__(profile)
        self.model = model
        self.api_key = os.environ.get('GEMINI_API_KEY', 'AIzaSyALfbAdHfJ6aRnqNyiTRmKmGVoena1JsdU')

    def execute(self, task: Task) -> ExecutionResult:
        """Execute task using Gemini API."""
        start_time = time.time()
        self.profile.current_load += 1

        try:
            import google.generativeai as genai
            genai.configure(api_key=self.api_key)
            model = genai.GenerativeModel(self.model)

            prompt = f"""Task: {task.title}

Description: {task.description}

Please complete this task and provide a clear response."""

            response = model.generate_content(prompt)
            output = response.text

            duration = time.time() - start_time
            # Estimate cost based on tokens
            tokens = len(prompt.split()) + len(output.split())
            cost = tokens * self.profile.cost_per_token

            result = ExecutionResult(
                task_id=task.id,
                success=True,
                output=output,
                agent=self.profile.agent_type,
                duration=duration,
                cost=cost
            )

        except Exception as e:
            result = ExecutionResult(
                task_id=task.id,
                success=False,
                output=None,
                agent=self.profile.agent_type,
                duration=time.time() - start_time,
                cost=0,
                error=str(e)
            )

        finally:
            self.profile.current_load -= 1

        self.execution_history.append(result)
        return result

    def health_check(self) -> bool:
        """Check Gemini API connectivity."""
        try:
            import google.generativeai as genai
            genai.configure(api_key=self.api_key)
            model = genai.GenerativeModel(self.model)
            response = model.generate_content("ping")
            return bool(response.text)
        except:
            return False


class ClaudeAgent(BaseAgent):
    """Claude-based agent (stub for API integration)."""

    def __init__(self, model: str = "claude-sonnet-4"):
        is_opus = "opus" in model.lower()
        profile = AgentProfile(
            agent_type=AgentType.CLAUDE_OPUS if is_opus else AgentType.CLAUDE_SONNET,
            name=f"Claude-{model}",
            strengths=["reasoning", "architecture", "code_review", "complex_analysis"] if is_opus else ["implementation", "balanced", "general"],
            cost_per_token=0.015 if is_opus else 0.003,
            max_context=200000,
            speed_rating=0.5 if is_opus else 0.7,
            quality_rating=0.95 if is_opus else 0.85,
            max_concurrent=2
        )
        super().__init__(profile)
        self.model = model

    def execute(self, task: Task) -> ExecutionResult:
        """Execute task (stub - would use Anthropic API)."""
        start_time = time.time()
        self.profile.current_load += 1

        try:
            # Stub implementation - would use anthropic library
            output = f"[Claude {self.model}] Task '{task.title}' would be executed here"
            duration = time.time() - start_time

            result = ExecutionResult(
                task_id=task.id,
                success=True,
                output=output,
                agent=self.profile.agent_type,
                duration=duration,
                cost=0.01  # Estimated
            )

        except Exception as e:
            result = ExecutionResult(
                task_id=task.id,
                success=False,
                output=None,
                agent=self.profile.agent_type,
                duration=time.time() - start_time,
                cost=0,
                error=str(e)
            )

        finally:
            self.profile.current_load -= 1

        self.execution_history.append(result)
        return result

    def health_check(self) -> bool:
        """Check Claude availability."""
        return True  # Stub


class AIVAAgent(BaseAgent):
    """AIVA/Qwen local agent (zero cost)."""

    def __init__(self):
        profile = AgentProfile(
            agent_type=AgentType.AIVA_QWEN,
            name="AIVA-Qwen",
            strengths=["monitoring", "simple_tasks", "local_execution", "zero_cost"],
            cost_per_token=0,
            max_context=32000,
            speed_rating=0.8,
            quality_rating=0.6,
            max_concurrent=10
        )
        super().__init__(profile)

    def execute(self, task: Task) -> ExecutionResult:
        """Execute task locally (stub for AIVA integration)."""
        start_time = time.time()
        self.profile.current_load += 1

        try:
            # Would connect to local AIVA/Qwen instance
            output = f"[AIVA] Task '{task.title}' executed locally"
            duration = time.time() - start_time

            result = ExecutionResult(
                task_id=task.id,
                success=True,
                output=output,
                agent=self.profile.agent_type,
                duration=duration,
                cost=0  # Always free
            )

        except Exception as e:
            result = ExecutionResult(
                task_id=task.id,
                success=False,
                output=None,
                agent=self.profile.agent_type,
                duration=time.time() - start_time,
                cost=0,
                error=str(e)
            )

        finally:
            self.profile.current_load -= 1

        self.execution_history.append(result)
        return result

    def health_check(self) -> bool:
        """Check AIVA availability."""
        # Would check Elestio connection
        return True  # Stub


class MultiAgentCoordinator:
    """
    Coordinates multiple AI agents for parallel task execution.

    Features:
    - Automatic agent selection based on task requirements
    - Parallel execution with thread pool
    - Load balancing across agents
    - Consensus validation for critical tasks
    - Budget tracking and optimization
    """

    def __init__(self, max_workers: int = 5, enable_learning: bool = True):
        self.max_workers = max_workers
        self.agents: Dict[AgentType, BaseAgent] = {}
        self.task_queue: queue.Queue = queue.Queue()
        self.results: Dict[str, ExecutionResult] = {}
        self.budget_spent: float = 0.0
        self.budget_limit: float = 336.0  # Default limit

        # Initialize learning loop for intelligent agent selection
        self.learning_loop = None
        if enable_learning and LEARNING_AVAILABLE:
            try:
                self.learning_loop = LearningLoop()
            except Exception:
                pass

        # Initialize default agents
        self._init_agents()

    def _init_agents(self):
        """Initialize available agents."""
        self.agents[AgentType.GEMINI_FLASH] = GeminiAgent("gemini-2.0-flash-exp")
        self.agents[AgentType.GEMINI_PRO] = GeminiAgent("gemini-1.5-pro")
        self.agents[AgentType.CLAUDE_SONNET] = ClaudeAgent("claude-sonnet-4")
        self.agents[AgentType.CLAUDE_OPUS] = ClaudeAgent("claude-opus-4")
        self.agents[AgentType.AIVA_QWEN] = AIVAAgent()

    def select_agent(self, task: Task) -> AgentType:
        """Select best agent for a task using learning-enhanced selection."""
        best_score = 0
        best_agent = AgentType.GEMINI_FLASH  # Default fallback

        # Check learning loop for recommendations first
        learning_recommendation = None
        if self.learning_loop:
            try:
                recs = self.learning_loop.get_recommendations(
                    task_type=task.task_type,
                    complexity=task.complexity.value if hasattr(task.complexity, 'value') else str(task.complexity)
                )
                if recs.get("agent") and recs.get("agent_confidence", 0) > 0.7:
                    # Map string agent name to AgentType
                    agent_name = recs["agent"]
                    for at in AgentType:
                        if at.value == agent_name or agent_name in at.value:
                            learning_recommendation = at
                            break
            except Exception:
                pass

        for agent_type, agent in self.agents.items():
            if not agent.profile.available:
                continue

            score = agent.profile.suitability_score(task.task_type, task.complexity)

            # Learning boost: If learning loop recommends this agent, boost score
            if learning_recommendation and agent_type == learning_recommendation:
                score += 0.3  # Significant boost for learned preference

            # Budget consideration
            if self.budget_spent + agent.profile.cost_per_token * 1000 > self.budget_limit:
                if agent.profile.cost_per_token > 0:
                    score *= 0.3  # Heavily penalize if near budget

            if score > best_score:
                best_score = score
                best_agent = agent_type

        return best_agent

    def execute_task(self, task: Task) -> ExecutionResult:
        """Execute a single task with the best agent."""
        if not task.assigned_agent:
            task.assigned_agent = self.select_agent(task)

        agent = self.agents[task.assigned_agent]
        task.status = "running"
        task.started_at = datetime.now().isoformat()

        result = agent.execute(task)

        task.status = "completed" if result.success else "failed"
        task.completed_at = datetime.now().isoformat()
        task.result = result.output
        task.error = result.error
        task.cost = result.cost

        self.budget_spent += result.cost
        self.results[task.id] = result

        # Record outcome to learning loop for future agent selection improvement
        if self.learning_loop:
            try:
                error_type = None
                if result.error:
                    # Extract error type from error message
                    if "timeout" in result.error.lower():
                        error_type = "timeout"
                    elif "rate" in result.error.lower():
                        error_type = "rate_limit"
                    elif "context" in result.error.lower():
                        error_type = "context_overflow"
                    else:
                        error_type = "execution_error"

                self.learning_loop.record_outcome(
                    task_id=task.id,
                    task_type=task.task_type,
                    task_title=task.title,
                    complexity=task.complexity.value if hasattr(task.complexity, 'value') else str(task.complexity),
                    agent_used=task.assigned_agent.value,
                    success=result.success,
                    duration=result.duration,
                    cost=result.cost,
                    error_type=error_type,
                    error_message=result.error
                )
            except Exception:
                pass  # Don't let learning failures affect execution

        return result

    def execute_parallel(
        self,
        tasks: List[Task],
        pattern: CoordinationPattern = CoordinationPattern.SPECIALIST
    ) -> Dict[str, ExecutionResult]:
        """Execute multiple tasks in parallel."""

        # Apply coordination pattern
        if pattern == CoordinationPattern.SPECIALIST:
            for task in tasks:
                task.assigned_agent = self.select_agent(task)
        elif pattern == CoordinationPattern.ROUND_ROBIN:
            agent_types = list(self.agents.keys())
            for i, task in enumerate(tasks):
                task.assigned_agent = agent_types[i % len(agent_types)]
        elif pattern == CoordinationPattern.HIERARCHICAL:
            # Lead agent (Opus) assigns others
            for task in tasks:
                if task.complexity in [TaskComplexity.CRITICAL, TaskComplexity.COMPLEX]:
                    task.assigned_agent = AgentType.CLAUDE_OPUS
                else:
                    task.assigned_agent = AgentType.GEMINI_FLASH

        # Execute with thread pool
        results = {}
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_task = {
                executor.submit(self.execute_task, task): task
                for task in tasks
            }

            for future in as_completed(future_to_task):
                task = future_to_task[future]
                try:
                    result = future.result()
                    results[task.id] = result
                except Exception as e:
                    results[task.id] = ExecutionResult(
                        task_id=task.id,
                        success=False,
                        output=None,
                        agent=task.assigned_agent or AgentType.GEMINI_FLASH,
                        duration=0,
                        cost=0,
                        error=str(e)
                    )

        return results

    def execute_consensus(
        self,
        task: Task,
        agent_types: List[AgentType] = None
    ) -> Tuple[Any, float]:
        """
        Execute task with multiple agents and reach consensus.
        Returns (consensus_output, confidence).
        """
        if agent_types is None:
            agent_types = [AgentType.GEMINI_FLASH, AgentType.CLAUDE_SONNET]

        outputs = []
        with ThreadPoolExecutor(max_workers=len(agent_types)) as executor:
            futures = []
            for agent_type in agent_types:
                task_copy = Task(
                    id=f"{task.id}_{agent_type.value}",
                    title=task.title,
                    description=task.description,
                    task_type=task.task_type,
                    complexity=task.complexity,
                    assigned_agent=agent_type
                )
                futures.append(executor.submit(self.execute_task, task_copy))

            for future in as_completed(futures):
                result = future.result()
                if result.success:
                    outputs.append(result.output)

        if not outputs:
            return None, 0.0

        # Simple consensus: return most common output or first
        # In practice, would use semantic similarity
        consensus = outputs[0]
        confidence = 1.0 if len(set(outputs)) == 1 else 0.5 + (0.5 / len(set(outputs)))

        return consensus, confidence

    def get_agent_stats(self) -> Dict:
        """Get statistics for all agents."""
        return {
            agent_type.value: agent.get_stats()
            for agent_type, agent in self.agents.items()
        }

    def get_budget_status(self) -> Dict:
        """Get current budget status."""
        return {
            "spent": self.budget_spent,
            "remaining": self.budget_limit - self.budget_spent,
            "limit": self.budget_limit,
            "utilization": self.budget_spent / self.budget_limit if self.budget_limit > 0 else 0
        }

    def health_check_all(self) -> Dict[str, bool]:
        """Check health of all agents."""
        return {
            agent_type.value: agent.health_check()
            for agent_type, agent in self.agents.items()
        }


def main():
    """Test the multi-agent coordinator."""
    import argparse
    parser = argparse.ArgumentParser(description="Genesis Multi-Agent Coordinator")
    parser.add_argument("--health", action="store_true", help="Check agent health")
    parser.add_argument("--stats", action="store_true", help="Show agent stats")
    parser.add_argument("--demo", action="store_true", help="Run demo execution")
    args = parser.parse_args()

    coordinator = MultiAgentCoordinator()

    if args.health:
        print("Agent Health Check:")
        for agent, healthy in coordinator.health_check_all().items():
            status = "✅" if healthy else "❌"
            print(f"  {agent}: {status}")
        return

    if args.stats:
        print("Agent Statistics:")
        print(json.dumps(coordinator.get_agent_stats(), indent=2))
        print("\nBudget Status:")
        print(json.dumps(coordinator.get_budget_status(), indent=2))
        return

    if args.demo:
        # Demo tasks
        tasks = [
            Task(
                id="demo-001",
                title="Simple Code Generation",
                description="Generate a hello world function",
                task_type="code_generation",
                complexity=TaskComplexity.SIMPLE
            ),
            Task(
                id="demo-002",
                title="Architecture Review",
                description="Review the Genesis memory architecture",
                task_type="architecture",
                complexity=TaskComplexity.COMPLEX
            ),
            Task(
                id="demo-003",
                title="Quick Status Check",
                description="Check system status",
                task_type="monitoring",
                complexity=TaskComplexity.TRIVIAL
            )
        ]

        print("Executing demo tasks in parallel...")
        results = coordinator.execute_parallel(tasks, CoordinationPattern.SPECIALIST)

        for task_id, result in results.items():
            status = "✅" if result.success else "❌"
            print(f"\n{task_id}: {status}")
            print(f"  Agent: {result.agent.value}")
            print(f"  Duration: {result.duration:.2f}s")
            print(f"  Cost: ${result.cost:.4f}")
            if result.error:
                print(f"  Error: {result.error}")

        print(f"\nTotal cost: ${coordinator.budget_spent:.4f}")
        return

    # Default: show available agents
    print("Genesis Multi-Agent Coordinator")
    print("=" * 40)
    print("\nAvailable Agents:")
    for agent_type, agent in coordinator.agents.items():
        print(f"\n  {agent.profile.name}:")
        print(f"    Strengths: {', '.join(agent.profile.strengths)}")
        print(f"    Speed: {agent.profile.speed_rating:.0%}")
        print(f"    Quality: {agent.profile.quality_rating:.0%}")
        print(f"    Cost: ${agent.profile.cost_per_token}/token")


if __name__ == "__main__":
    main()