"""
Contractor Agent - The orchestrator that coordinates all agents.
"""

import json
import subprocess
import logging
from pathlib import Path
from typing import Optional, Dict, Any, List
from enum import Enum

from ..base import Orchestrator, Agent, AgentContext, AgentResult
from ..registry import orchestrator
from ..task import Task, Plan, TaskStatus, TaskManager
from ..events import EventEmitter, Event
from ..config import ChainConfig


logger = logging.getLogger(__name__)


class ContractorState(Enum):
    """States of the Contractor."""
    IDLE = "idle"
    PLANNING = "planning"
    EXECUTING = "executing"
    VERIFYING = "verifying"
    COMPLETED = "completed"
    FAILED = "failed"


@orchestrator(name="contractor", role="Project Orchestrator")
class Contractor(Orchestrator):
    """
    The Contractor is the main orchestrator that:
    1. Invokes the Planner to create a task plan
    2. Iterates through tasks, invoking Coder for each
    3. Handles failures with smart retries
    4. Verifies outputs when possible
    
    This is the default orchestrator but can be replaced with a custom one.
    """
    
    name = "contractor"
    role = "Project Orchestrator"
    
    def __init__(
        self,
        agents: Optional[List[Agent]] = None,
        llm=None,
        model: Optional[str] = None,
        events: Optional[EventEmitter] = None,
        config: Optional[ChainConfig] = None,
        **kwargs
    ):
        super().__init__(agents=agents, llm=llm, model=model, **kwargs)
        self.events = events
        self.config = config or ChainConfig()
        self.state = ContractorState.IDLE
        self.task_manager: Optional[TaskManager] = None
        self.current_plan: Optional[Plan] = None
    
    def execute(self, context: AgentContext) -> AgentResult:
        """Execute the full workflow."""
        requirements = context.get("requirements", "")
        output_dir = context.get("output_dir", "./output")
        
        if not requirements:
            return AgentResult.fail("No requirements provided")
        
        # Setup
        self.task_manager = TaskManager(workspace_dir=output_dir)
        
        try:
            # Phase 1: Planning
            self.state = ContractorState.PLANNING
            plan_result = self._run_planning(requirements)
            
            if not plan_result.success:
                return plan_result
            
            # Phase 2: Execution
            self.state = ContractorState.EXECUTING
            exec_result = self._run_execution()
            
            if not exec_result.success:
                self.state = ContractorState.FAILED
                return exec_result
            
            # Phase 3: Final Verification
            self.state = ContractorState.VERIFYING
            verify_result = self._run_final_verification()
            
            self.state = ContractorState.COMPLETED
            
            return AgentResult.ok({
                "plan": self.current_plan.to_dict() if self.current_plan else None,
                "progress": self.task_manager.get_progress(),
                "verification": verify_result,
                "output_dir": output_dir,
            })
            
        except Exception as e:
            self.state = ContractorState.FAILED
            logger.exception("Contractor execution failed")
            return AgentResult.fail(str(e))
    
    def _run_planning(self, requirements: str) -> AgentResult:
        """Run the planning phase."""
        logger.info("Starting planning phase...")
        
        planner = self.get_agent("planner")
        if not planner:
            return AgentResult.fail("Planner agent not found")
        
        context = AgentContext(
            task="Create implementation plan",
            inputs={"requirements": requirements}
        )
        
        result = self.invoke_agent("planner", context)
        
        if not result.success:
            return AgentResult.fail(f"Planning failed: {result.error}")
        
        # Load plan
        plan_data = result.output.get("plan", {})
        self.current_plan = self.task_manager.load_plan(plan_data)
        
        logger.info(f"Plan created with {len(self.current_plan.tasks)} tasks")
        
        if self.events:
            total = len(self.current_plan.tasks)
            self.events.emit(Event.PLANNING_COMPLETE, task_count=total)
        
        return result
    
    def _run_execution(self) -> AgentResult:
        """Execute all tasks in the plan."""
        logger.info("Starting task execution...")
        
        coder = self.get_agent("coder")
        if not coder:
            return AgentResult.fail("Coder agent not found")
        
        completed = 0
        failed = 0
        total = len(self.current_plan.tasks) if self.current_plan else 0
        
        while True:
            # Get next task
            task = self.task_manager.get_next_task()
            
            if task is None:
                # Check if complete or stuck
                if self.task_manager.is_complete():
                    break
                
                # Check for failed tasks
                failed_tasks = self.task_manager.get_failed_tasks()
                if failed_tasks:
                    # Try to reset retryable tasks
                    reset_count = self.task_manager.reset_failed_tasks()
                    if reset_count == 0:
                        logger.error("Tasks have exhausted all retries")
                        return AgentResult.fail(
                            f"{len(failed_tasks)} tasks failed and cannot be retried"
                        )
                    continue
                
                break
            
            # Execute task
            task.status = TaskStatus.IN_PROGRESS
            result = self._execute_task(task)
            
            if result.success:
                task.mark_complete(json.dumps(result.output))
                completed += 1
                
                if self.events:
                    self.events.emit(
                        Event.TASK_COMPLETE,
                        task=task.title,
                        completed=completed,
                        total=total
                    )
                    self.events.emit(
                        Event.PROGRESS,
                        completed=completed,
                        total=total,
                        percent=round(completed / total * 100, 1),
                        task=task.title
                    )
            else:
                # Smart retry
                retry_result = self._smart_retry(task, result)
                
                if retry_result and retry_result.success:
                    task.mark_complete(json.dumps(retry_result.output))
                    completed += 1
                    
                    if self.events:
                        self.events.emit(
                            Event.TASK_COMPLETE,
                            task=task.title,
                            completed=completed,
                            total=total
                        )
                else:
                    task.mark_failed(result.error or "Unknown error")
                    failed += 1
                    
                    if self.events:
                        self.events.emit(
                            Event.TASK_FAILED,
                            task=task.title,
                            error=task.error
                        )
            
            # Save progress
            self.task_manager.save_plan()
        
        progress = self.task_manager.get_progress()
        logger.info(f"Execution complete: {progress['completed']}/{progress['total']} tasks")
        
        if failed > 0:
            return AgentResult.fail(
                f"Completed {completed} tasks but {failed} failed",
                progress=progress
            )
        
        return AgentResult.ok(progress=progress)
    
    def _execute_task(self, task: Task) -> AgentResult:
        """Execute a single task."""
        logger.info(f"Executing: {task.title}")
        
        context = AgentContext(
            task=task.title,
            inputs={
                "task": task.to_dict(),
            },
            workspace_dir=self.config.output_dir if self.config else "./output"
        )
        
        result = self.invoke_agent("coder", context)
        
        # Verify if possible
        if result.success:
            verification = result.output.get("verification", {}) if result.output else {}
            if verification.get("command"):
                verify_result = self._verify_task(verification["command"])
                if not verify_result:
                    return AgentResult.fail(
                        f"Verification failed for task: {task.title}",
                        output=result.output
                    )
        
        return result
    
    def _smart_retry(self, task: Task, failed_result: AgentResult) -> Optional[AgentResult]:
        """Retry a failed task with error context."""
        max_retries = self.config.max_task_retries if self.config else 3
        
        for retry in range(max_retries):
            if not task.can_retry():
                break
            
            logger.info(f"Retry {retry + 1}/{max_retries} for: {task.title}")
            
            if self.events:
                self.events.emit(
                    Event.TASK_RETRY,
                    task=task.title,
                    attempt=retry + 1
                )
            
            # Analyze error
            error_analysis = self._analyze_error(failed_result.error)
            
            context = AgentContext(
                task=task.title,
                inputs={
                    "task": task.to_dict(),
                    "retry_attempt": retry + 1,
                    "previous_error": failed_result.error,
                    "retry_instructions": error_analysis,
                },
                workspace_dir=self.config.output_dir if self.config else "./output"
            )
            
            result = self.invoke_agent("coder", context)
            task.retry_count += 1
            
            if result.success:
                return result
            
            failed_result = result
        
        return None
    
    def _analyze_error(self, error: str) -> str:
        """Analyze error and return fix instructions."""
        if not error:
            return "Please fix the issue and try again"
        
        error_lower = error.lower()
        
        if "json" in error_lower or "parse" in error_lower:
            return "Respond with ONLY valid JSON, no markdown code blocks"
        elif "syntax" in error_lower:
            return "Check all brackets, quotes, and indentation"
        elif "import" in error_lower:
            return "Ensure all imports are at the top and use correct package names"
        elif "attribute" in error_lower:
            return "Check that all methods and attributes are properly defined"
        else:
            return f"Fix this error: {error[:200]}"
    
    def _verify_task(self, command: str) -> bool:
        """Run a verification command."""
        try:
            result = subprocess.run(
                command,
                shell=True,
                cwd=self.config.output_dir if self.config else "./output",
                capture_output=True,
                text=True,
                timeout=30
            )
            return result.returncode == 0
        except subprocess.TimeoutExpired:
            # Timeout might be OK for GUI apps
            return True
        except Exception as e:
            logger.warning(f"Verification failed: {e}")
            return False
    
    def _run_final_verification(self) -> Dict[str, Any]:
        """Run final verification if specified in plan."""
        if not self.current_plan:
            return {"skipped": True}
        
        verification = self.current_plan.metadata.get("final_verification", {})
        command = verification.get("command")
        
        if not command:
            return {"skipped": True, "reason": "No verification command"}
        
        logger.info(f"Running final verification: {command}")
        
        success = self._verify_task(command)
        
        return {
            "command": command,
            "success": success,
            "expected": verification.get("expected", ""),
        }
