"""
Base classes for AgentChain framework.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, TYPE_CHECKING
from enum import Enum
import logging

if TYPE_CHECKING:
    from .llm.base import LLMClient


class AgentStatus(Enum):
    """Status of an agent."""
    IDLE = "idle"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class AgentContext:
    """
    Context passed to agents during execution.
    
    Contains all information an agent needs to perform its task.
    """
    task: str
    inputs: Dict[str, Any] = field(default_factory=dict)
    workspace_dir: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    # Chain context - other agents available, shared state
    available_agents: List[str] = field(default_factory=list)
    shared_state: Dict[str, Any] = field(default_factory=dict)
    
    def get(self, key: str, default: Any = None) -> Any:
        """Get a value from inputs."""
        return self.inputs.get(key, default)
    
    def set_shared(self, key: str, value: Any) -> None:
        """Set a shared state value accessible to other agents."""
        self.shared_state[key] = value
    
    def get_shared(self, key: str, default: Any = None) -> Any:
        """Get a shared state value."""
        return self.shared_state.get(key, default)


@dataclass
class AgentResult:
    """Result returned by an agent after execution."""
    success: bool
    output: Optional[Any] = None
    error: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    # For orchestrators - which agents were invoked
    agents_invoked: List[str] = field(default_factory=list)
    
    @classmethod
    def ok(cls, output: Any = None, **metadata) -> "AgentResult":
        """Create a successful result."""
        return cls(success=True, output=output, metadata=metadata)
    
    @classmethod
    def fail(cls, error: str, **metadata) -> "AgentResult":
        """Create a failed result."""
        return cls(success=False, error=error, metadata=metadata)


class Agent(ABC):
    """
    Base class for all agents in AgentChain.
    
    Subclass this to create custom agents. At minimum, implement the `execute` method.
    
    Example:
        class MyAgent(Agent):
            name = "my_agent"
            role = "Does something useful"
            
            def execute(self, context: AgentContext) -> AgentResult:
                # Do work here
                return AgentResult.ok({"result": "done"})
    """
    
    # Override these in subclasses
    name: str = "agent"
    role: str = "Generic Agent"
    description: str = ""
    
    # System prompt for LLM-based agents
    system_prompt: str = ""
    
    def __init__(
        self,
        llm: Optional["LLMClient"] = None,
        model: Optional[str] = None,
        **kwargs
    ):
        """
        Initialize the agent.
        
        Args:
            llm: LLM client to use for this agent
            model: Model name override (if different from default)
            **kwargs: Additional configuration
        """
        self.llm = llm
        self.model = model
        self.config = kwargs
        self.status = AgentStatus.IDLE
        self._logger = logging.getLogger(f"llama2a.{self.name}")
    
    @abstractmethod
    def execute(self, context: AgentContext) -> AgentResult:
        """
        Execute the agent's task.
        
        Args:
            context: The execution context with task details
            
        Returns:
            AgentResult with success/failure and output
        """
        pass
    
    def query_llm(self, prompt: str, system_prompt: Optional[str] = None) -> str:
        """
        Query the LLM with a prompt.
        
        Args:
            prompt: User prompt
            system_prompt: Override system prompt (uses self.system_prompt if not provided)
            
        Returns:
            LLM response text
        """
        if not self.llm:
            raise RuntimeError(f"Agent {self.name} has no LLM client configured")
        
        sys_prompt = system_prompt or self.system_prompt
        response = self.llm.generate(
            prompt=prompt,
            system_prompt=sys_prompt,
            model=self.model
        )
        return response.content
    
    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}(name={self.name}, role={self.role})>"


class Orchestrator(Agent):
    """
    Base class for orchestrating agents.
    
    An orchestrator coordinates other agents, deciding which to invoke
    and in what order. The default Contractor is an orchestrator.
    """
    
    is_orchestrator: bool = True
    
    def __init__(self, agents: Optional[List[Agent]] = None, **kwargs):
        super().__init__(**kwargs)
        self.agents: Dict[str, Agent] = {}
        if agents:
            for agent in agents:
                self.register_agent(agent)
    
    def register_agent(self, agent: Agent) -> None:
        """Register an agent with this orchestrator."""
        self.agents[agent.name] = agent
        self._logger.debug(f"Registered agent: {agent.name}")
    
    def get_agent(self, name: str) -> Optional[Agent]:
        """Get a registered agent by name."""
        return self.agents.get(name)
    
    def invoke_agent(self, name: str, context: AgentContext) -> AgentResult:
        """Invoke a registered agent."""
        agent = self.get_agent(name)
        if not agent:
            return AgentResult.fail(f"Agent '{name}' not found")
        
        self._logger.info(f"Invoking agent: {name}")
        agent.status = AgentStatus.RUNNING
        
        try:
            result = agent.execute(context)
            agent.status = AgentStatus.COMPLETED if result.success else AgentStatus.FAILED
            return result
        except Exception as e:
            agent.status = AgentStatus.FAILED
            self._logger.error(f"Agent {name} failed: {e}")
            return AgentResult.fail(str(e))
