"""
AgentChain - Main fluent API for the framework.
"""

import logging
from pathlib import Path
from typing import Optional, Type, Callable, Dict, Any, List, Union

from .base import Agent, Orchestrator, AgentContext, AgentResult
from .config import ChainConfig, AgentConfig
from .events import EventEmitter, Event, EventHandler, console_progress_handler
from .registry import AgentRegistry, get_registry
from .llm import OllamaClient, LLMClient
from .task import TaskManager, Task, Plan, TaskStatus


logger = logging.getLogger(__name__)


class AgentChain:
    """
    Main entry point for AgentChain framework.
    
    Provides a fluent API for configuring and running agent chains.
    
    Example:
        result = AgentChain() \\
            .requirements("Build a snake game with Python") \\
            .model("deepseek-coder-v2:16b") \\
            .configure("planner", model="llama3:70b") \\
            .configure("coder", model="codellama:34b") \\
            .output("./my_game") \\
            .on_progress(lambda p: print(f"{p}%")) \\
            .run()
    """
    
    def __init__(self):
        self._config = ChainConfig()
        self._events = EventEmitter()
        self._registry = get_registry()
        self._custom_agents: Dict[str, Type[Agent]] = {}
        self._custom_orchestrator: Optional[Type[Orchestrator]] = None
        self._llm_client: Optional[LLMClient] = None
        
        # Set up default console progress
        self._use_console_progress = True
    
    # ==================== Fluent Configuration Methods ====================
    
    def requirements(self, requirements: str) -> "AgentChain":
        """Set requirements as a string."""
        self._config.requirements = requirements
        return self
    
    def from_file(self, filepath: str) -> "AgentChain":
        """Load requirements from a file."""
        self._config.requirements_file = filepath
        return self
    
    @classmethod
    def from_requirements(cls, filepath: str) -> "AgentChain":
        """Create AgentChain from a requirements file."""
        return cls().from_file(filepath)
    
    def model(self, model_name: str) -> "AgentChain":
        """Set the default model for all agents."""
        self._config.default_model = model_name
        return self
    
    def ollama_url(self, url: str) -> "AgentChain":
        """Set the Ollama server URL."""
        self._config.ollama_url = url
        return self
    
    def output(self, directory: str) -> "AgentChain":
        """Set the output directory for generated files."""
        self._config.output_dir = directory
        return self
    
    def output_dir(self, directory: str) -> "AgentChain":
        """Alias for output()."""
        return self.output(directory)
    
    def configure(
        self,
        agent_name: str,
        *,
        model: Optional[str] = None,
        enabled: bool = True,
        max_retries: int = 3,
        **extra
    ) -> "AgentChain":
        """
        Configure a specific agent.
        
        Args:
            agent_name: Name of the agent to configure
            model: Model to use for this agent (overrides default)
            enabled: Whether this agent is enabled
            max_retries: Max retries for this agent
            **extra: Additional agent-specific configuration
        """
        config = AgentConfig(
            enabled=enabled,
            model=model,
            max_retries=max_retries,
            extra=extra
        )
        self._config.agents[agent_name] = config
        return self
    
    def disable(self, agent_name: str) -> "AgentChain":
        """Disable a specific agent."""
        return self.configure(agent_name, enabled=False)
    
    def temperature(self, temp: float) -> "AgentChain":
        """Set LLM temperature."""
        self._config.temperature = temp
        return self
    
    def timeout(self, seconds: int) -> "AgentChain":
        """Set request timeout in seconds."""
        self._config.timeout = seconds
        return self
    
    def verbose(self, enabled: bool = True) -> "AgentChain":
        """Enable verbose logging."""
        self._config.verbose = enabled
        self._config.log_level = "DEBUG" if enabled else "INFO"
        return self
    
    def infinity_mode(self, enabled: bool = True) -> "AgentChain":
        """Enable infinity mode (continuous improvement)."""
        self._config.infinity_mode = enabled
        return self
    
    # ==================== Agent Registration ====================
    
    def register(self, agent_class: Type[Agent], name: Optional[str] = None) -> "AgentChain":
        """
        Register a custom agent.
        
        Args:
            agent_class: The agent class to register
            name: Optional name override
        """
        agent_name = name or agent_class.name
        self._custom_agents[agent_name] = agent_class
        self._registry.register(agent_class, agent_name)
        return self
    
    def orchestrator(self, orchestrator_class: Type[Orchestrator]) -> "AgentChain":
        """
        Set a custom orchestrator (replaces default Contractor).
        
        Args:
            orchestrator_class: The orchestrator class to use
        """
        self._custom_orchestrator = orchestrator_class
        self._config.orchestrator_name = orchestrator_class.name
        return self
    
    # ==================== Event Handlers ====================
    
    def on(self, event: Union[Event, str], handler: EventHandler) -> "AgentChain":
        """
        Register an event handler.
        
        Args:
            event: Event type or string event name
            handler: Callback function
        """
        if isinstance(event, str):
            event = Event(event)
        self._events.on(event, handler)
        return self
    
    def on_progress(self, handler: Callable[[int, int, str], None]) -> "AgentChain":
        """
        Register a progress handler.
        
        Handler receives: (completed_count, total_count, current_task_name)
        """
        def wrapper(event_data):
            handler(
                event_data.get("completed", 0),
                event_data.get("total", 0),
                event_data.get("task", "")
            )
        self._events.on(Event.PROGRESS, wrapper)
        return self
    
    def on_task_complete(self, handler: Callable[[str], None]) -> "AgentChain":
        """Register handler for task completion."""
        def wrapper(event_data):
            handler(event_data.get("task", ""))
        self._events.on(Event.TASK_COMPLETE, wrapper)
        return self
    
    def on_error(self, handler: Callable[[str, str], None]) -> "AgentChain":
        """Register error handler. Receives (task_name, error_message)."""
        def wrapper(event_data):
            handler(event_data.get("task", ""), event_data.get("error", ""))
        self._events.on(Event.TASK_FAILED, wrapper)
        return self
    
    def quiet(self) -> "AgentChain":
        """Disable console progress output."""
        self._use_console_progress = False
        return self
    
    # ==================== Execution ====================
    
    def run(self) -> AgentResult:
        """
        Execute the agent chain.
        
        Returns:
            AgentResult with the final output
        """
        # Setup logging
        log_level = getattr(logging, self._config.log_level.upper(), logging.INFO)
        logging.basicConfig(
            level=log_level,
            format='%(asctime)s | %(levelname)-8s | %(name)s | %(message)s'
        )
        
        # Validate config
        errors = self._config.validate()
        if errors:
            return AgentResult.fail("; ".join(errors))
        
        # Setup console progress if enabled
        if self._use_console_progress:
            self._events.on(Event.TASK_COMPLETE, console_progress_handler)
            self._events.on(Event.TASK_FAILED, console_progress_handler)
        
        # Emit chain start
        self._events.emit(Event.CHAIN_START, config=self._config.to_dict())
        
        try:
            result = self._execute()
            self._events.emit(Event.CHAIN_COMPLETE, result=result.output)
            return result
        except Exception as e:
            logger.exception("Chain execution failed")
            self._events.emit(Event.CHAIN_ERROR, error=str(e))
            return AgentResult.fail(str(e))
    
    def _execute(self) -> AgentResult:
        """Internal execution logic."""
        # Load requirements
        requirements = self._load_requirements()
        
        # Create output directory
        output_dir = Path(self._config.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # Initialize LLM client
        self._llm_client = OllamaClient(
            model=self._config.default_model,
            base_url=self._config.ollama_url,
            timeout=self._config.timeout,
            temperature=self._config.temperature
        )
        
        # Check LLM availability
        if not self._llm_client.is_available():
            return AgentResult.fail(
                f"Cannot connect to Ollama at {self._config.ollama_url}"
            )
        
        # Get orchestrator
        orchestrator = self._get_orchestrator()
        
        # Create context
        context = AgentContext(
            task="Build software from requirements",
            inputs={
                "requirements": requirements,
                "output_dir": str(output_dir),
            },
            workspace_dir=str(output_dir),
            available_agents=list(self._get_available_agents().keys())
        )
        
        # Execute
        logger.info("Starting AgentChain execution...")
        result = orchestrator.execute(context)
        
        return result
    
    def _load_requirements(self) -> str:
        """Load requirements from config."""
        if self._config.requirements:
            return self._config.requirements
        
        if self._config.requirements_file:
            path = Path(self._config.requirements_file)
            return path.read_text()
        
        raise ValueError("No requirements provided")
    
    def _get_orchestrator(self) -> Orchestrator:
        """Get and initialize the orchestrator."""
        if self._custom_orchestrator:
            orchestrator_class = self._custom_orchestrator
        else:
            # Import default Contractor
            from .agents.contractor import Contractor
            orchestrator_class = Contractor
        
        # Get available agents
        agents = self._get_available_agents()
        
        # Create orchestrator with agents
        orchestrator = orchestrator_class(
            agents=list(agents.values()),
            llm=self._llm_client,
            model=self._config.get_agent_model(orchestrator_class.name),
            events=self._events,
            config=self._config
        )
        
        return orchestrator
    
    def _get_available_agents(self) -> Dict[str, Agent]:
        """Get all available and enabled agents."""
        agents = {}
        
        # Import default agents
        from .agents.planner import Planner
        from .agents.coder import Coder
        
        default_agents = {
            "planner": Planner,
            "coder": Coder,
        }
        
        # Add default agents if enabled
        for name, agent_class in default_agents.items():
            config = self._config.agents.get(name, AgentConfig())
            if config.enabled:
                agent = agent_class(
                    llm=self._llm_client,
                    model=config.model or self._config.default_model,
                    events=self._events,
                    workspace_dir=self._config.output_dir
                )
                agents[name] = agent
        
        # Add custom agents
        for name, agent_class in self._custom_agents.items():
            config = self._config.agents.get(name, AgentConfig())
            if config.enabled:
                agent = agent_class(
                    llm=self._llm_client,
                    model=config.model or self._config.default_model,
                    events=self._events,
                    workspace_dir=self._config.output_dir
                )
                agents[name] = agent
        
        return agents
    
    # ==================== Utility Methods ====================
    
    def get_config(self) -> ChainConfig:
        """Get the current configuration."""
        return self._config
    
    def list_agents(self) -> List[str]:
        """List all available agents."""
        from .agents.planner import Planner
        from .agents.coder import Coder
        
        agents = ["planner", "coder"] + list(self._custom_agents.keys())
        return agents
    
    def __repr__(self) -> str:
        return f"<AgentChain(model={self._config.default_model}, agents={self.list_agents()})>"
