"""
GENESIS HYPERDRIVE CONTROLLER
=============================
Multi-agent autonomous loop controller for Genesis.

Manages execution across:
- Claude (Opus 4.5) - Complex reasoning, architecture
- Claude (Sonnet) - Balanced implementation
- Gemini 2.0 Flash - Speed, large context, $300 budget
- Gemini 2.5 Pro - Maximum capability, 2M context
- AIVA (Qwen) - Zero-cost local execution

Budget: $336 total
- $300 Gemini credits (~300-600 iterations)
- $36 Claude credits (~12 iterations)

Usage:
    python hyperdrive_controller.py run     # Run the loop
    python hyperdrive_controller.py status  # Check status
    python hyperdrive_controller.py select "task description"  # Test agent selection
"""

import json
import os
import subprocess
import time
import hashlib
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
from enum import Enum
import re

# Memory Cortex integration for execution history
try:
    from genesis_memory_cortex import GenesisCortex
    MEMORY_CORTEX_AVAILABLE = True
except ImportError:
    MEMORY_CORTEX_AVAILABLE = False
    GenesisCortex = None


class AgentType(Enum):
    """Available agents in the Genesis multi-agent system."""
    CLAUDE_OPUS = "claude-opus"           # $15/M tokens - Complex reasoning
    CLAUDE_SONNET = "claude-sonnet"       # $3/M tokens - Balanced
    GEMINI_FLASH = "gemini-2.0-flash"     # $0.50/M tokens - Speed + cost
    GEMINI_PRO = "gemini-2.5-pro"         # $1.25/M tokens - Max capability
    AIVA_QWEN = "aiva-qwen"               # $0 - Local execution


@dataclass
class AgentBudget:
    """Budget tracking per agent type."""
    agent_type: AgentType
    total_budget: float
    spent: float = 0.0
    iterations: int = 0
    avg_cost_per_iteration: float = 0.0

    @property
    def remaining(self) -> float:
        return self.total_budget - self.spent

    @property
    def estimated_iterations_left(self) -> int:
        # AIVA is unlimited
        if self.agent_type == AgentType.AIVA_QWEN:
            return 999999
        if self.avg_cost_per_iteration > 0:
            return int(self.remaining / self.avg_cost_per_iteration)
        # Default estimates
        defaults = {
            AgentType.CLAUDE_OPUS: 3.0,
            AgentType.CLAUDE_SONNET: 1.0,
            AgentType.GEMINI_FLASH: 0.50,
            AgentType.GEMINI_PRO: 1.0,
            AgentType.AIVA_QWEN: 0.01  # Prevent div by zero
        }
        cost = defaults.get(self.agent_type, 1.0)
        if cost == 0:
            return 999999
        return int(self.remaining / cost)


@dataclass
class TaskClassification:
    """Classification of a task for agent selection."""
    task_id: str
    description: str
    complexity: str  # simple, moderate, complex, architectural
    context_size: str  # small, medium, large, massive
    speed_priority: bool
    cost_sensitive: bool
    requires_reasoning: bool
    recommended_agent: AgentType
    reasoning: str


class AgentSelector:
    """
    Intelligent agent selection based on task characteristics.

    Heuristics:
    - Complex reasoning → Claude Opus
    - Standard implementation → Claude Sonnet or Gemini Flash
    - Large context needed → Gemini Pro (2M tokens)
    - Fast iterations → Gemini Flash
    - Zero-cost needed → AIVA Qwen
    - Cost-sensitive + complex → Gemini Flash with verification
    """

    # Keywords that suggest specific agent needs
    COMPLEXITY_KEYWORDS = {
        'architectural': ['architect', 'design', 'framework', 'system', 'integration'],
        'complex': ['refactor', 'optimize', 'analyze', 'debug complex', 'migration'],
        'moderate': ['implement', 'add', 'create', 'build', 'fix'],
        'simple': ['update', 'change', 'modify', 'rename', 'typo', 'comment']
    }

    CONTEXT_KEYWORDS = {
        'massive': ['entire codebase', 'all files', 'full analysis', 'comprehensive'],
        'large': ['multiple files', 'cross-module', 'integration'],
        'medium': ['single module', 'one component'],
        'small': ['single file', 'one function', 'quick fix']
    }

    def __init__(self, budget_tracker: 'BudgetTracker' = None):
        self.budget_tracker = budget_tracker

    def classify_task(self, task_description: str, task_id: str = None) -> TaskClassification:
        """Classify a task and recommend the best agent."""
        desc_lower = task_description.lower()
        task_id = task_id or hashlib.md5(task_description.encode()).hexdigest()[:8]

        # Determine complexity
        complexity = 'moderate'  # default
        for level, keywords in self.COMPLEXITY_KEYWORDS.items():
            if any(kw in desc_lower for kw in keywords):
                complexity = level
                break

        # Determine context size
        context_size = 'medium'  # default
        for size, keywords in self.CONTEXT_KEYWORDS.items():
            if any(kw in desc_lower for kw in keywords):
                context_size = size
                break

        # Determine priorities
        speed_priority = any(kw in desc_lower for kw in ['fast', 'quick', 'rapid', 'urgent'])
        cost_sensitive = self.budget_tracker.is_budget_low() if self.budget_tracker else False
        requires_reasoning = any(kw in desc_lower for kw in [
            'why', 'reason', 'explain', 'decide', 'choose', 'evaluate', 'compare'
        ])

        # Select agent based on classification
        recommended_agent, reasoning = self._select_agent(
            complexity, context_size, speed_priority, cost_sensitive, requires_reasoning
        )

        return TaskClassification(
            task_id=task_id,
            description=task_description,
            complexity=complexity,
            context_size=context_size,
            speed_priority=speed_priority,
            cost_sensitive=cost_sensitive,
            requires_reasoning=requires_reasoning,
            recommended_agent=recommended_agent,
            reasoning=reasoning
        )

    def _select_agent(
        self,
        complexity: str,
        context_size: str,
        speed_priority: bool,
        cost_sensitive: bool,
        requires_reasoning: bool
    ) -> Tuple[AgentType, str]:
        """Select the best agent based on task characteristics."""

        # Massive context → Must use Gemini Pro (2M tokens)
        if context_size == 'massive':
            return AgentType.GEMINI_PRO, "Task requires massive context (2M tokens) - only Gemini Pro can handle"

        # Architectural complexity → Claude Opus (unless budget is low)
        if complexity == 'architectural' and not cost_sensitive:
            return AgentType.CLAUDE_OPUS, "Architectural task requires deep reasoning - Claude Opus recommended"

        # Complex reasoning without speed requirement → Claude Opus
        if requires_reasoning and complexity in ['complex', 'architectural'] and not speed_priority:
            if cost_sensitive:
                return AgentType.GEMINI_FLASH, "Complex reasoning but cost-sensitive - Gemini Flash with verification"
            return AgentType.CLAUDE_OPUS, "Complex reasoning task - Claude Opus for best results"

        # Speed priority → Gemini Flash
        if speed_priority:
            return AgentType.GEMINI_FLASH, "Speed priority - Gemini Flash for fast iteration"

        # Cost sensitive → Gemini Flash or AIVA
        if cost_sensitive:
            if complexity == 'simple':
                return AgentType.AIVA_QWEN, "Simple task + cost sensitive - AIVA (zero cost)"
            return AgentType.GEMINI_FLASH, "Cost sensitive - Gemini Flash ($0.50/M tokens)"

        # Large context → Gemini Pro
        if context_size == 'large':
            return AgentType.GEMINI_PRO, "Large context required - Gemini Pro for extended context window"

        # Moderate complexity → Gemini Flash (best cost/performance)
        if complexity == 'moderate':
            return AgentType.GEMINI_FLASH, "Moderate complexity - Gemini Flash for optimal cost/performance"

        # Simple tasks → AIVA or Gemini Flash
        if complexity == 'simple':
            return AgentType.GEMINI_FLASH, "Simple task - Gemini Flash for quick execution"

        # Default: Gemini Flash (best balance of cost and capability)
        return AgentType.GEMINI_FLASH, "Default selection - Gemini Flash for balanced performance"


class BudgetTracker:
    """Track budget across all agents."""

    def __init__(self, state_path: str = None):
        self.state_path = Path(state_path or "E:/genesis-system/data/hyperdrive_budget.json")
        self.budgets: Dict[AgentType, AgentBudget] = {}
        self._initialize_budgets()
        self._load_state()

    def _initialize_budgets(self):
        """Initialize budgets for all agent types."""
        self.budgets = {
            AgentType.CLAUDE_OPUS: AgentBudget(AgentType.CLAUDE_OPUS, 18.0),  # ~6 iterations at $3
            AgentType.CLAUDE_SONNET: AgentBudget(AgentType.CLAUDE_SONNET, 18.0),  # ~18 iterations at $1
            AgentType.GEMINI_FLASH: AgentBudget(AgentType.GEMINI_FLASH, 200.0),  # ~400 iterations
            AgentType.GEMINI_PRO: AgentBudget(AgentType.GEMINI_PRO, 100.0),  # ~100 iterations
            AgentType.AIVA_QWEN: AgentBudget(AgentType.AIVA_QWEN, float('inf')),  # Unlimited
        }

    def _load_state(self):
        """Load budget state from disk."""
        if self.state_path.exists():
            try:
                with open(self.state_path) as f:
                    data = json.load(f)
                    for agent_name, budget_data in data.get('budgets', {}).items():
                        agent_type = AgentType(agent_name)
                        if agent_type in self.budgets:
                            self.budgets[agent_type].spent = budget_data.get('spent', 0)
                            self.budgets[agent_type].iterations = budget_data.get('iterations', 0)
                            self.budgets[agent_type].avg_cost_per_iteration = budget_data.get('avg_cost', 0)
            except Exception:
                pass

    def _save_state(self):
        """Save budget state to disk."""
        self.state_path.parent.mkdir(parents=True, exist_ok=True)
        data = {
            'updated_at': datetime.now().isoformat(),
            'budgets': {
                agent.value: {
                    'total': budget.total_budget,
                    'spent': budget.spent,
                    'remaining': budget.remaining,
                    'iterations': budget.iterations,
                    'avg_cost': budget.avg_cost_per_iteration
                }
                for agent, budget in self.budgets.items()
            },
            'total_remaining': self.get_total_remaining()
        }
        with open(self.state_path, 'w') as f:
            json.dump(data, f, indent=2)

    def record_usage(self, agent_type: AgentType, cost: float):
        """Record cost for an agent iteration."""
        if agent_type in self.budgets:
            budget = self.budgets[agent_type]
            budget.spent += cost
            budget.iterations += 1
            budget.avg_cost_per_iteration = budget.spent / budget.iterations
            self._save_state()

    def get_remaining(self, agent_type: AgentType) -> float:
        """Get remaining budget for an agent type."""
        return self.budgets.get(agent_type, AgentBudget(agent_type, 0)).remaining

    def get_total_remaining(self) -> float:
        """Get total remaining budget across all paid agents."""
        total = 0.0
        for b in self.budgets.values():
            if b.agent_type != AgentType.AIVA_QWEN:
                if b.remaining != float('inf'):
                    total += b.remaining
        return total

    def is_budget_low(self, threshold: float = 50.0) -> bool:
        """Check if total budget is running low."""
        return self.get_total_remaining() < threshold

    def check_budget_alert(self, threshold: float = 50.0) -> bool:
        """Check budget and send alert if low."""
        if not self.is_budget_low(threshold):
            return False

        remaining = self.get_total_remaining()
        print(f"[BUDGET ALERT] Budget low: ${remaining:.2f} remaining (threshold: ${threshold:.2f})")

        # Try to send Slack alert
        try:
            from monitoring.slack_bridge import SlackBridge
            bridge = SlackBridge()
            bridge.send_budget_alert(remaining, threshold)
        except ImportError:
            pass

        return True

    def budget_alert(self, threshold: float = 50.0) -> Dict:
        """Get budget alert status and details."""
        remaining = self.get_total_remaining()
        is_low = remaining < threshold
        is_critical = remaining < 10.0

        return {
            "alert": is_low,
            "critical": is_critical,
            "remaining": remaining,
            "threshold": threshold,
            "message": f"Budget {'CRITICAL' if is_critical else 'LOW' if is_low else 'OK'}: ${remaining:.2f}"
        }

    def get_summary(self) -> Dict:
        """Get budget summary."""
        return {
            'total_remaining': self.get_total_remaining(),
            'budgets': {
                agent.value: {
                    'remaining': budget.remaining,
                    'iterations': budget.iterations,
                    'estimated_remaining': budget.estimated_iterations_left
                }
                for agent, budget in self.budgets.items()
            }
        }


class HyperdriveController:
    """
    Main controller for the Genesis Hyperdrive multi-agent loop.

    Coordinates task execution across multiple AI agents,
    selecting the optimal agent for each task based on:
    - Task complexity
    - Context requirements
    - Budget constraints
    - Speed requirements
    """

    def __init__(self, tasks_path: str = None):
        self.tasks_path = Path(tasks_path or "E:/genesis-system/loop/tasks.json")
        self.budget_tracker = BudgetTracker()
        self.agent_selector = AgentSelector(self.budget_tracker)
        self.log_path = Path("E:/genesis-system/data/hyperdrive_log.jsonl")
        self.max_iterations = 100

        # Memory Cortex for storing execution history
        self.memory_cortex = None
        if MEMORY_CORTEX_AVAILABLE and GenesisCortex:
            try:
                self.memory_cortex = GenesisCortex()
                print("[Hyperdrive] Memory Cortex connected")
            except Exception as e:
                print(f"[Hyperdrive] Memory Cortex unavailable: {e}")

    def store_execution_memory(self, task_id: str, task_title: str, agent: str,
                               success: bool, response: str, cost: float):
        """Store task execution in Memory Cortex for learning."""
        if not self.memory_cortex:
            return

        try:
            memory_entry = {
                "type": "hyperdrive_execution",
                "task_id": task_id,
                "task_title": task_title,
                "agent": agent,
                "success": success,
                "response_preview": response[:500] if response else "",
                "cost": cost,
                "timestamp": datetime.now().isoformat()
            }

            # Store in working memory for immediate access
            self.memory_cortex.store_working(
                key=f"hyperdrive:{task_id}",
                value=memory_entry,
                ttl=3600  # 1 hour TTL
            )

            # Store in episodic memory for long-term learning
            self.memory_cortex.store_episodic(
                event_type="task_execution",
                content=f"Task '{task_title}' executed by {agent}: {'SUCCESS' if success else 'FAILED'}",
                metadata=memory_entry
            )

            print(f"[Memory] Stored execution history for {task_id}")
        except Exception as e:
            print(f"[Memory] Failed to store execution: {e}")

    def load_tasks(self) -> List[Dict]:
        """Load tasks from tasks.json."""
        if not self.tasks_path.exists():
            return []
        try:
            with open(self.tasks_path) as f:
                data = json.load(f)
                return [s for s in data.get('stories', []) if not s.get('passes', False)]
        except Exception as e:
            print(f"Error loading tasks: {e}")
            return []

    def select_agent_for_task(self, task: Dict) -> TaskClassification:
        """Select the best agent for a task."""
        description = task.get('description', task.get('title', ''))
        task_id = task.get('id', 'unknown')
        return self.agent_selector.classify_task(description, task_id)

    def execute_with_agent(
        self,
        task: Dict,
        classification: TaskClassification
    ) -> Tuple[bool, str]:
        """
        Execute a task with the selected agent.

        Returns (success, output).
        """
        agent = classification.recommended_agent
        task_title = task.get('title', 'Unknown task')

        print(f"\n{'='*60}")
        print(f"HYPERDRIVE: Executing with {agent.value}")
        print(f"Task: {task_title}")
        print(f"Reasoning: {classification.reasoning}")
        print(f"{'='*60}\n")

        # Build the prompt
        prompt = self._build_execution_prompt(task, classification)

        # Execute based on agent type
        if agent == AgentType.AIVA_QWEN:
            return self._execute_aiva(prompt)
        elif agent in [AgentType.GEMINI_FLASH, AgentType.GEMINI_PRO]:
            return self._execute_gemini(prompt, agent)
        else:
            return self._execute_claude(prompt, agent)

    def _build_execution_prompt(self, task: Dict, classification: TaskClassification) -> str:
        """Build the execution prompt for the agent."""
        criteria = task.get('acceptance_criteria', [])
        criteria_text = "\n".join([f"- {c.get('description', c)}" for c in criteria])

        return f"""# GENESIS HYPERDRIVE TASK EXECUTION

## Task: {task.get('title', 'Unknown')}

## Description
{task.get('description', '')}

## Acceptance Criteria
{criteria_text}

## Context
{task.get('context', 'No additional context')}

## Agent Assignment
You are: {classification.recommended_agent.value}
Reason: {classification.reasoning}

## Instructions
1. Implement this task completely
2. Test against ALL acceptance criteria
3. Report: PASS or FAIL for each criterion
4. If all PASS, respond with: TASK_COMPLETE
5. If any FAIL, explain what failed and why

## Working Directory
E:\\genesis-system

## Verification Protocol
Every claim requires proof. Show command output for any verification.

Execute now. Self-verify. Report results."""

    def _execute_gemini(self, prompt: str, agent: AgentType) -> Tuple[bool, str]:
        """Execute via Gemini API DIRECTLY ($300 credits)."""
        try:
            # Import Gemini executor for direct API access
            from gemini_executor import GeminiExecutor

            # Set API key
            os.environ['GEMINI_API_KEY'] = 'AIzaSyALfbAdHfJ6aRnqNyiTRmKmGVoena1JsdU'

            executor = GeminiExecutor()

            # Map agent type to model
            model_map = {
                AgentType.GEMINI_FLASH: "flash",
                AgentType.GEMINI_PRO: "pro"
            }
            model = model_map.get(agent, "flash")

            # Execute directly with Gemini API
            result = executor.execute(prompt, model=model)

            output = result.response
            success = result.task_complete or "TASK_COMPLETE" in output

            # Record actual cost
            self.budget_tracker.record_usage(agent, result.cost_estimate)

            return success, output

        except Exception as e:
            return False, f"Gemini execution error: {e}"

    def _execute_claude(self, prompt: str, agent: AgentType) -> Tuple[bool, str]:
        """Execute via Claude API directly."""
        try:
            result = subprocess.run(
                ['claude', '-p', prompt, '--max-turns', '5'],
                capture_output=True,
                text=True,
                timeout=300,
                cwd="E:/genesis-system"
            )

            output = result.stdout
            success = "TASK_COMPLETE" in output

            # Estimate cost
            token_estimate = len(prompt.split()) * 2 + len(output.split()) * 2
            cost_per_token = 0.000003 if agent == AgentType.CLAUDE_OPUS else 0.000001
            cost = token_estimate * cost_per_token

            self.budget_tracker.record_usage(agent, cost)

            return success, output

        except Exception as e:
            return False, f"Execution error: {e}"

    def _execute_aiva(self, prompt: str) -> Tuple[bool, str]:
        """Execute via AIVA (local Qwen) - zero cost."""
        try:
            # Call AIVA API on Elestio
            import urllib.request
            import urllib.parse

            data = json.dumps({"prompt": prompt}).encode()
            req = urllib.request.Request(
                "http://152.53.201.152:23405/api/generate",
                data=data,
                headers={'Content-Type': 'application/json'}
            )

            with urllib.request.urlopen(req, timeout=120) as response:
                result = json.loads(response.read().decode())
                output = result.get('response', '')
                success = "TASK_COMPLETE" in output
                return success, output

        except Exception as e:
            return False, f"AIVA execution error: {e}"

    def log_execution(self, task: Dict, classification: TaskClassification, success: bool, output: str):
        """Log execution to JSONL file."""
        self.log_path.parent.mkdir(parents=True, exist_ok=True)

        log_entry = {
            'timestamp': datetime.now().isoformat(),
            'task_id': task.get('id'),
            'task_title': task.get('title'),
            'agent': classification.recommended_agent.value,
            'reasoning': classification.reasoning,
            'success': success,
            'output_preview': output[:500] if output else None
        }

        with open(self.log_path, 'a') as f:
            f.write(json.dumps(log_entry) + '\n')

    def run_loop(self, max_iterations: int = None):
        """Run the main hyperdrive loop."""
        max_iterations = max_iterations or self.max_iterations
        iteration = 0

        print(f"\n{'='*60}")
        print("GENESIS HYPERDRIVE CONTROLLER")
        print(f"{'='*60}")
        print(f"Budget: ${self.budget_tracker.get_total_remaining():.2f} remaining")
        print(f"Max iterations: {max_iterations}")
        print(f"{'='*60}\n")

        while iteration < max_iterations:
            iteration += 1

            # Load pending tasks
            tasks = self.load_tasks()
            if not tasks:
                print("All tasks complete!")
                break

            # Get next task
            task = tasks[0]
            print(f"\n--- Iteration {iteration} ---")
            print(f"Task: {task.get('title')}")

            # Select agent
            classification = self.select_agent_for_task(task)
            print(f"Agent: {classification.recommended_agent.value}")
            print(f"Reason: {classification.reasoning}")

            # Execute
            success, output = self.execute_with_agent(task, classification)

            # Log
            self.log_execution(task, classification, success, output)

            if success:
                print(f"SUCCESS: Task completed")
                # Update tasks.json to mark as passed
                self._mark_task_complete(task.get('id'))
            else:
                print(f"FAILED: Task not completed")
                print(f"Output: {output[:200]}...")

            # Budget check
            if self.budget_tracker.get_total_remaining() < 5:
                print("\nBudget critically low! Stopping loop.")
                break

            # Small delay between iterations
            time.sleep(2)

        print(f"\n{'='*60}")
        print("HYPERDRIVE LOOP COMPLETE")
        print(f"Iterations: {iteration}")
        print(f"Budget remaining: ${self.budget_tracker.get_total_remaining():.2f}")
        print(f"{'='*60}")

    def _mark_task_complete(self, task_id: str):
        """Mark a task as complete in tasks.json."""
        try:
            with open(self.tasks_path) as f:
                data = json.load(f)

            for story in data.get('stories', []):
                if story.get('id') == task_id:
                    story['passes'] = True
                    story['completed_at'] = datetime.now().isoformat()
                    break

            with open(self.tasks_path, 'w') as f:
                json.dump(data, f, indent=2)

        except Exception as e:
            print(f"Error updating task: {e}")

    def get_status(self) -> Dict:
        """Get current hyperdrive status."""
        tasks = self.load_tasks()
        return {
            'pending_tasks': len(tasks),
            'budget': self.budget_tracker.get_summary(),
            'next_task': tasks[0] if tasks else None
        }


def main():
    import sys

    controller = HyperdriveController()

    if len(sys.argv) < 2:
        print("Usage: python hyperdrive_controller.py [run|status|select]")
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "run":
        max_iter = int(sys.argv[2]) if len(sys.argv) > 2 else 10
        controller.run_loop(max_iterations=max_iter)

    elif cmd == "status":
        status = controller.get_status()
        print(json.dumps(status, indent=2, default=str))

    elif cmd == "select":
        if len(sys.argv) < 3:
            print("Usage: python hyperdrive_controller.py select 'task description'")
            sys.exit(1)
        task_desc = sys.argv[2]
        classification = controller.agent_selector.classify_task(task_desc)
        print(f"\nTask: {task_desc}")
        print(f"Complexity: {classification.complexity}")
        print(f"Context Size: {classification.context_size}")
        print(f"Recommended Agent: {classification.recommended_agent.value}")
        print(f"Reasoning: {classification.reasoning}")

    else:
        print(f"Unknown command: {cmd}")
        sys.exit(1)


if __name__ == "__main__":
    main()