"""Base Agent - Unified foundation for all AI agent implementations.

This module provides the single base class for all agent operations,
following DRY principles and ensuring consistent behavior across all agents.
"""

from __future__ import annotations

import os
import asyncio
import logging
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Generic, TypeVar, Optional, Protocol, runtime_checkable
from pathlib import Path
from datetime import datetime
from dataclasses import field, dataclass

from .model_registry import ModelConfig, registry

logger = logging.getLogger(__name__)


# Type variables for generic context
TContext = TypeVar("TContext")
TResult = TypeVar("TResult")


@runtime_checkable
class AgentContext(Protocol):
    """Protocol for agent execution context."""

    async def log(self, message: str, level: str = "info") -> None:
        """Log a message."""
        pass

    async def progress(self, message: str, percentage: Optional[float] = None) -> None:
        """Report progress."""
        pass


@dataclass
class AgentConfig:
    """Configuration for agent execution."""

    model: str = "claude-3-5-sonnet-20241022"
    timeout: int = 300
    max_retries: int = 3
    working_dir: Optional[Path] = None
    environment: Dict[str, str] = field(default_factory=dict)
    stream_output: bool = False
    use_worktree: bool = False

    def __post_init__(self) -> None:
        """Resolve model name and validate configuration."""
        self.model = registry.resolve(self.model)
        if self.working_dir and not isinstance(self.working_dir, Path):
            self.working_dir = Path(self.working_dir)


@dataclass
class AgentResult:
    """Result from agent execution."""

    success: bool
    output: Optional[str] = None
    error: Optional[str] = None
    duration: Optional[float] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def content(self) -> str:
        """Get the primary content (output or error)."""
        return self.output if self.success else (self.error or "Unknown error")


class BaseAgent(ABC, Generic[TContext]):
    """Base class for all AI agents.

    This is the single foundation for all agent implementations,
    ensuring consistent behavior and eliminating code duplication.
    """

    def __init__(self, config: Optional[AgentConfig] = None) -> None:
        """Initialize agent with configuration.

        Args:
            config: Agent configuration
        """
        self.config = config or AgentConfig()
        self._start_time: Optional[datetime] = None
        self._end_time: Optional[datetime] = None

    @property
    @abstractmethod
    def name(self) -> str:
        """Agent name."""
        pass

    @property
    @abstractmethod
    def description(self) -> str:
        """Agent description."""
        pass

    async def execute(
        self,
        prompt: str,
        context: Optional[TContext] = None,
        **kwargs: Any,
    ) -> AgentResult:
        """Execute agent with prompt.

        Args:
            prompt: The prompt or task
            context: Execution context
            **kwargs: Additional parameters

        Returns:
            Agent execution result
        """
        self._start_time = datetime.now()

        try:
            # Setup environment
            env = self._prepare_environment()

            # Log start
            if context and isinstance(context, AgentContext):
                await context.log(f"Starting {self.name} with model {self.config.model}")

            # Execute with retries
            result = await self._execute_with_retries(prompt, context, env, **kwargs)

            # Calculate duration
            self._end_time = datetime.now()
            duration = (self._end_time - self._start_time).total_seconds()

            return AgentResult(
                success=True,
                output=result,
                duration=duration,
                metadata={"model": self.config.model, "agent": self.name},
            )

        except Exception as e:
            self._end_time = datetime.now()
            duration = (self._end_time - self._start_time).total_seconds() if self._start_time else None

            logger.error(f"Agent {self.name} failed: {e}")

            return AgentResult(
                success=False,
                error=str(e),
                duration=duration,
                metadata={"model": self.config.model, "agent": self.name},
            )

    def _prepare_environment(self) -> Dict[str, str]:
        """Prepare environment variables for execution.

        Returns:
            Environment variables dictionary
        """
        env = os.environ.copy()

        # Add model-specific API key
        model_config = registry.get(self.config.model)
        if model_config and model_config.api_key_env:
            key_var = model_config.api_key_env
            if key_var in os.environ:
                env[key_var] = os.environ[key_var]

        # Add Hanzo unified auth
        if "HANZO_API_KEY" in os.environ:
            env["HANZO_API_KEY"] = os.environ["HANZO_API_KEY"]

        # Add custom environment
        env.update(self.config.environment)

        return env

    async def _execute_with_retries(
        self,
        prompt: str,
        context: Optional[TContext],
        env: Dict[str, str],
        **kwargs: Any,
    ) -> str:
        """Execute with retry logic.

        Args:
            prompt: The prompt
            context: Execution context
            env: Environment variables
            **kwargs: Additional parameters

        Returns:
            Execution output

        Raises:
            Exception: If all retries fail
        """
        last_error = None

        for attempt in range(self.config.max_retries):
            try:
                # Call the implementation
                result = await self._execute_impl(prompt, context, env, **kwargs)
                return result

            except asyncio.TimeoutError:
                last_error = f"Timeout after {self.config.timeout} seconds"
                if context and isinstance(context, AgentContext):
                    await context.log(f"Attempt {attempt + 1} timed out", "warning")

            except Exception as e:
                last_error = str(e)
                if context and isinstance(context, AgentContext):
                    await context.log(f"Attempt {attempt + 1} failed: {e}", "warning")

                # Don't retry on certain errors
                if "unauthorized" in str(e).lower() or "forbidden" in str(e).lower():
                    raise

            # Wait before retry (exponential backoff)
            if attempt < self.config.max_retries - 1:
                await asyncio.sleep(2**attempt)

        raise Exception(f"All {self.config.max_retries} attempts failed. Last error: {last_error}")

    @abstractmethod
    async def _execute_impl(
        self,
        prompt: str,
        context: Optional[TContext],
        env: Dict[str, str],
        **kwargs: Any,
    ) -> str:
        """Implementation-specific execution.

        Args:
            prompt: The prompt
            context: Execution context
            env: Environment variables
            **kwargs: Additional parameters

        Returns:
            Execution output
        """
        pass


class CLIAgent(BaseAgent[TContext]):
    """Base class for CLI-based agents."""

    @property
    @abstractmethod
    def cli_command(self) -> str:
        """CLI command to execute."""
        pass

    def build_command(self, prompt: str, **kwargs: Any) -> List[str]:
        """Build the CLI command.

        Args:
            prompt: The prompt
            **kwargs: Additional parameters

        Returns:
            Command arguments list
        """
        command = [self.cli_command]

        # Add model if specified
        model_config = registry.get(self.config.model)
        if model_config:
            command.extend(["--model", model_config.full_name])

        # Add prompt
        command.append(prompt)

        return command

    async def _execute_impl(
        self,
        prompt: str,
        context: Optional[TContext],
        env: Dict[str, str],
        **kwargs: Any,
    ) -> str:
        """Execute CLI command.

        Args:
            prompt: The prompt
            context: Execution context
            env: Environment variables
            **kwargs: Additional parameters

        Returns:
            Command output
        """
        command = self.build_command(prompt, **kwargs)

        # Determine if we need stdin
        needs_stdin = self.cli_command in ["claude", "cline"]

        # Execute command
        process = await asyncio.create_subprocess_exec(
            *command,
            stdin=asyncio.subprocess.PIPE if needs_stdin else None,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(self.config.working_dir) if self.config.working_dir else None,
            env=env,
        )

        # Handle timeout
        try:
            stdout, stderr = await asyncio.wait_for(
                process.communicate(prompt.encode() if needs_stdin else None),
                timeout=self.config.timeout,
            )
        except asyncio.TimeoutError:
            process.kill()
            raise asyncio.TimeoutError(f"Command timed out after {self.config.timeout} seconds")

        # Check for errors
        if process.returncode != 0:
            error_msg = stderr.decode() if stderr else "Command failed"
            raise Exception(error_msg)

        return stdout.decode()


class APIAgent(BaseAgent[TContext]):
    """Base class for API-based agents."""

    async def _execute_impl(
        self,
        prompt: str,
        context: Optional[TContext],
        env: Dict[str, str],
        **kwargs: Any,
    ) -> str:
        """Execute via API.

        Args:
            prompt: The prompt
            context: Execution context
            env: Environment variables
            **kwargs: Additional parameters

        Returns:
            API response
        """
        # This would be implemented by specific API agents
        # using the appropriate client library
        raise NotImplementedError("API agents must implement _execute_impl")


class AgentOrchestrator:
    """Orchestrator for managing multiple agents."""

    def __init__(self, default_config: Optional[AgentConfig] = None) -> None:
        """Initialize orchestrator.

        Args:
            default_config: Default configuration for agents
        """
        self.default_config = default_config or AgentConfig()
        self._agents: Dict[str, BaseAgent] = {}
        self._semaphore: Optional[asyncio.Semaphore] = None

    def register(self, agent: BaseAgent) -> None:
        """Register an agent.

        Args:
            agent: Agent to register
        """
        self._agents[agent.name] = agent

    def get_agent(self, name: str) -> Optional[BaseAgent]:
        """Get agent by name.

        Args:
            name: Agent name

        Returns:
            Agent instance or None
        """
        return self._agents.get(name)

    async def execute_single(
        self,
        agent_name: str,
        prompt: str,
        context: Optional[Any] = None,
        **kwargs: Any,
    ) -> AgentResult:
        """Execute single agent.

        Args:
            agent_name: Name of agent to use
            prompt: The prompt
            context: Execution context
            **kwargs: Additional parameters

        Returns:
            Execution result
        """
        agent = self.get_agent(agent_name)
        if not agent:
            return AgentResult(
                success=False,
                error=f"Agent '{agent_name}' not found",
            )

        return await agent.execute(prompt, context, **kwargs)

    async def execute_parallel(
        self,
        tasks: List[Dict[str, Any]],
        max_concurrent: int = 5,
    ) -> List[AgentResult]:
        """Execute multiple agents in parallel.

        Args:
            tasks: List of task definitions
            max_concurrent: Maximum concurrent executions

        Returns:
            List of results
        """
        self._semaphore = asyncio.Semaphore(max_concurrent)

        async def run_with_semaphore(task: Dict[str, Any]) -> AgentResult:
            async with self._semaphore:
                return await self.execute_single(
                    task["agent"],
                    task["prompt"],
                    task.get("context"),
                    **task.get("kwargs", {}),
                )

        return await asyncio.gather(
            *[run_with_semaphore(task) for task in tasks],
            return_exceptions=False,
        )

    async def execute_consensus(
        self,
        prompt: str,
        agents: List[str],
        threshold: float = 0.66,
    ) -> Dict[str, Any]:
        """Execute consensus operation with multiple agents.

        Args:
            prompt: The prompt
            agents: List of agent names
            threshold: Agreement threshold

        Returns:
            Consensus results
        """
        # Execute all agents in parallel
        tasks = [{"agent": agent, "prompt": prompt} for agent in agents]
        results = await self.execute_parallel(tasks)

        # Analyze consensus
        successful = [r for r in results if r.success]
        agreement = len(successful) / len(results) if results else 0

        return {
            "consensus_reached": agreement >= threshold,
            "agreement_score": agreement,
            "individual_results": results,
            "agents_used": agents,
        }

    async def execute_chain(
        self,
        initial_prompt: str,
        agents: List[str],
    ) -> List[AgentResult]:
        """Execute agents in a chain, passing output forward.

        Args:
            initial_prompt: Initial prompt
            agents: List of agent names

        Returns:
            List of results from each step
        """
        results = []
        current_prompt = initial_prompt

        for agent_name in agents:
            result = await self.execute_single(agent_name, current_prompt)
            results.append(result)

            if result.success and result.output:
                # Use output as input for next agent
                current_prompt = f"Review and improve:\n{result.output}"
            else:
                # Chain broken
                break

        return results


__all__ = [
    "AgentContext",
    "AgentConfig",
    "AgentResult",
    "BaseAgent",
    "CLIAgent",
    "APIAgent",
    "AgentOrchestrator",
]
