"""Base agent class for Hikaflow AI debugging agents.

This module provides the abstract base class and common utilities
for all Hikaflow debugging agents.
"""

from __future__ import annotations

import abc
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, Generic, List, Optional, Protocol, TypeVar

logger = logging.getLogger(__name__)

# Optional OpenTelemetry integration for agent observability.
# When the opentelemetry-api package is installed, tool calls and agent
# executions are emitted as spans. If OTel is not installed, a lightweight
# no-op wrapper is used so that all tracing code is a no-op at runtime.
try:
    from opentelemetry import trace as _otel_trace

    _tracer = _otel_trace.get_tracer("hikaflow.agent")
except ImportError:  # pragma: no cover
    _tracer = None  # type: ignore[assignment]


class AgentStatus(str, Enum):
    """Status of an agent execution."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class AgentCapability(str, Enum):
    """Capabilities an agent can have."""

    READ_CODE = "read_code"
    WRITE_CODE = "write_code"
    EXECUTE_TESTS = "execute_tests"
    SEARCH_CODE = "search_code"
    READ_TRANSCRIPTS = "read_transcripts"
    GENERATE_FIX = "generate_fix"
    CREATE_PR = "create_pr"
    ANALYZE_ERROR = "analyze_error"


@dataclass
class AgentContext:
    """Context provided to agents during execution.

    Contains all information an agent needs to perform its task.
    """

    # Project context
    project_path: str
    project_id: Optional[str] = None

    # Authentication
    api_base: str = "https://get.hikaflow.com"
    token: Optional[str] = None
    github_token: Optional[str] = None

    # Task context
    run_id: Optional[str] = None
    test_name: Optional[str] = None
    error_type: Optional[str] = None
    error_message: Optional[str] = None
    stack_trace: Optional[str] = None

    # Configuration
    max_iterations: int = 10
    timeout_seconds: int = 300
    temperature: float = 0.0

    # Execution state
    start_time: Optional[datetime] = None
    current_iteration: int = 0

    # Memory/history
    messages: List[Dict[str, Any]] = field(default_factory=list)
    tool_calls: List[Dict[str, Any]] = field(default_factory=list)

    def elapsed_seconds(self) -> float:
        """Get elapsed time since start."""
        if self.start_time is None:
            return 0.0
        return (datetime.now(timezone.utc) - self.start_time).total_seconds()


@dataclass
class AgentResult:
    """Result of an agent execution."""

    status: AgentStatus
    output: Optional[str] = None
    error: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    iterations: int = 0
    elapsed_seconds: float = 0.0
    tool_calls: int = 0

    @property
    def success(self) -> bool:
        """Whether the agent completed successfully."""
        return self.status == AgentStatus.COMPLETED


class MaxIterationsExceeded(Exception):
    """Raised when an agent exceeds its maximum iteration count."""

    def __init__(self, max_iterations: int):
        self.max_iterations = max_iterations
        super().__init__(f"Agent exceeded max iterations ({max_iterations})")


class Tool(Protocol):
    """Protocol for agent tools.

    Tools should define a parameters_schema property returning a JSON Schema
    dict describing the expected input parameters. This allows the LLM to
    know what arguments a tool accepts and enables pre-execution validation.
    """

    @property
    def name(self) -> str:
        """Tool name."""
        ...

    @property
    def description(self) -> str:
        """Tool description for LLM."""
        ...

    @property
    def parameters_schema(self) -> Dict[str, Any]:
        """JSON Schema describing the tool's input parameters.

        Returns a dict conforming to JSON Schema (draft-07 or later).
        Example::

            {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "File path"},
                    "content": {"type": "string", "description": "File content"},
                },
                "required": ["path", "content"],
            }

        Returns an empty schema by default for backwards compatibility.
        """
        return {"type": "object", "properties": {}}

    async def execute(self, **kwargs) -> Any:
        """Execute the tool with given arguments."""
        ...


T = TypeVar("T")


class BaseAgent(abc.ABC, Generic[T]):
    """Abstract base class for all Hikaflow debugging agents.

    Provides common functionality for:
    - Tool registration and management
    - Execution loop with iteration limits
    - Logging and observability
    - Error handling

    Subclasses must implement:
    - `run()`: Main execution logic
    - `capabilities`: List of capabilities

    Example:
        class MyAgent(BaseAgent[MyResult]):
            capabilities = [AgentCapability.READ_CODE, AgentCapability.ANALYZE_ERROR]

            async def run(self, context: AgentContext) -> MyResult:
                # Implementation
                pass
    """

    # Class attributes to be overridden by subclasses
    name: str = "base_agent"
    description: str = "Base agent"
    capabilities: List[AgentCapability] = []

    def __init__(
        self,
        model: Optional[str] = None,
        temperature: float = 0.0,
        tools: Optional[List[Tool]] = None,
    ):
        """Initialize the agent.

        Args:
            model: LLM model to use (e.g., 'claude-sonnet-4-5-20250929')
            temperature: LLM temperature (0.0 = deterministic)
            tools: List of tools available to the agent
        """
        self.model = model or "claude-sonnet-4-5-20250929"
        self.temperature = temperature
        self._tools: Dict[str, Tool] = {}

        if tools:
            for tool in tools:
                self.register_tool(tool)

    def register_tool(self, tool: Tool) -> None:
        """Register a tool for the agent to use.

        Args:
            tool: Tool instance to register
        """
        self._tools[tool.name] = tool
        logger.debug(f"Registered tool: {tool.name}")

    def get_tool(self, name: str) -> Optional[Tool]:
        """Get a registered tool by name.

        Args:
            name: Tool name

        Returns:
            Tool if found, None otherwise
        """
        return self._tools.get(name)

    @property
    def tools(self) -> List[Tool]:
        """Get all registered tools."""
        return list(self._tools.values())

    @property
    def tool_names(self) -> List[str]:
        """Get names of all registered tools."""
        return list(self._tools.keys())

    def has_capability(self, capability: AgentCapability) -> bool:
        """Check if agent has a specific capability.

        Args:
            capability: Capability to check

        Returns:
            True if agent has the capability
        """
        return capability in self.capabilities

    @abc.abstractmethod
    async def run(self, context: AgentContext) -> T:
        """Execute the agent's main logic.

        Args:
            context: Execution context

        Returns:
            Agent-specific result type
        """
        ...

    async def execute(self, context: AgentContext) -> AgentResult:
        """Execute the agent with full lifecycle management.

        This is the main entry point for running an agent.
        Handles timing, error handling, and result packaging.

        Args:
            context: Execution context

        Returns:
            AgentResult with status and output
        """
        context.start_time = datetime.now(timezone.utc)

        try:
            logger.info(f"Starting agent: {self.name}")

            import asyncio
            result = await asyncio.wait_for(
                self.run(context),
                timeout=context.timeout_seconds,
            )

            elapsed = context.elapsed_seconds()
            logger.info(f"Agent {self.name} completed in {elapsed:.2f}s")

            return AgentResult(
                status=AgentStatus.COMPLETED,
                output=str(result) if result else None,
                metadata={"result": result} if result else {},
                iterations=context.current_iteration,
                elapsed_seconds=elapsed,
                tool_calls=len(context.tool_calls),
            )

        except MaxIterationsExceeded as e:
            elapsed = context.elapsed_seconds()
            logger.warning(f"Agent {self.name} exceeded max iterations: {e}")
            return AgentResult(
                status=AgentStatus.FAILED,
                error=str(e),
                iterations=context.current_iteration,
                elapsed_seconds=elapsed,
                tool_calls=len(context.tool_calls),
            )

        except TimeoutError as e:
            elapsed = context.elapsed_seconds()
            logger.warning(f"Agent {self.name} timed out after {elapsed:.2f}s")
            return AgentResult(
                status=AgentStatus.FAILED,
                error=f"Timeout: {e}",
                iterations=context.current_iteration,
                elapsed_seconds=elapsed,
                tool_calls=len(context.tool_calls),
            )

        except Exception as e:
            elapsed = context.elapsed_seconds()
            logger.exception(f"Agent {self.name} failed: {e}")
            return AgentResult(
                status=AgentStatus.FAILED,
                error=str(e),
                iterations=context.current_iteration,
                elapsed_seconds=elapsed,
                tool_calls=len(context.tool_calls),
            )

    async def call_tool(
        self,
        context: AgentContext,
        tool_name: str,
        **kwargs,
    ) -> Any:
        """Call a registered tool.

        Args:
            context: Execution context
            tool_name: Name of tool to call
            **kwargs: Arguments to pass to tool

        Returns:
            Tool result

        Raises:
            ValueError: If tool not found
        """
        tool = self.get_tool(tool_name)
        if tool is None:
            raise ValueError(f"Tool not found: {tool_name}")

        # Enforce iteration limit (defense-in-depth alongside timeout)
        context.current_iteration += 1
        if context.current_iteration > context.max_iterations:
            raise MaxIterationsExceeded(context.max_iterations)

        logger.debug(f"Calling tool: {tool_name} with {kwargs}")

        # Wrap in OTel span when tracing is available
        if _tracer is not None:
            with _tracer.start_as_current_span(
                f"tool.{tool_name}",
                attributes={
                    "tool.name": tool_name,
                    "tool.arguments": str(kwargs)[:200],
                    "agent.name": self.name,
                    "agent.iteration": context.current_iteration,
                },
            ) as span:
                result = await tool.execute(**kwargs)
                result_str = str(result)[:500] if result else ""
                span.set_attribute("tool.result_length", len(result_str))
        else:
            result = await tool.execute(**kwargs)

        # Record tool call
        context.tool_calls.append({
            "tool": tool_name,
            "arguments": kwargs,
            "result": str(result)[:500] if result else None,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })

        return result

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__} name={self.name} tools={len(self._tools)}>"


class CompositeAgent(BaseAgent[T]):
    """Agent that orchestrates multiple sub-agents.

    Use this when a task requires coordination between
    multiple specialized agents.
    """

    def __init__(
        self,
        agents: List[BaseAgent],
        **kwargs,
    ):
        """Initialize composite agent.

        Args:
            agents: List of sub-agents to orchestrate
            **kwargs: Passed to BaseAgent
        """
        super().__init__(**kwargs)
        self._agents = {agent.name: agent for agent in agents}

    @property
    def agents(self) -> List[BaseAgent]:
        """Get all sub-agents."""
        return list(self._agents.values())

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

    async def delegate(
        self,
        agent_name: str,
        context: AgentContext,
    ) -> AgentResult:
        """Delegate execution to a sub-agent.

        Args:
            agent_name: Name of sub-agent
            context: Execution context

        Returns:
            Sub-agent result

        Raises:
            ValueError: If agent not found
        """
        agent = self.get_agent(agent_name)
        if agent is None:
            raise ValueError(f"Agent not found: {agent_name}")

        logger.info(f"Delegating to agent: {agent_name}")
        return await agent.execute(context)


def adapt_tool_for_pydantic_ai(tool: Tool) -> dict:
    """Create a Pydantic AI-compatible tool definition from a BaseAgent Tool.

    This adapter bridges the two agent architectures (BaseAgent's Tool protocol
    and Pydantic AI's @agent.tool decorator system) so that tools written for
    one system can be reused in the other.

    Args:
        tool: A Tool instance conforming to the BaseAgent Tool protocol.

    Returns:
        A dict with 'name', 'description', 'parameters_schema', and 'function'
        keys suitable for registration with a Pydantic AI agent.

    Example::

        for tool in my_base_agent.tools:
            spec = adapt_tool_for_pydantic_ai(tool)
            # Register spec['function'] with Pydantic AI agent
    """
    return {
        "name": tool.name,
        "description": tool.description,
        "parameters_schema": tool.parameters_schema,
        "function": tool.execute,
    }
