import logging
from typing import Any, Callable
from pydantic_ai import Agent
from pydantic_ai.exceptions import ModelRetry
from core.prompt import CAVEMAN_PROMPT

logger: logging.Logger = logging.getLogger("agent-base")


class BaseAgent:
    """Base AI agent implementing Pydantic AI interaction patterns."""

    def __init__(self, model_name: str) -> None:
        self.model_name: str = model_name
        self.system_prompt: str = ""
        self.pydantic_agent: Agent[Any, str] | None = None

    async def initialize(self) -> None:
        """Initialize the Pydantic AI agent instance."""
        logger.info(f"Initializing Pydantic AI agent with model: {self.model_name}")
        model_str: str = self.model_name if ":" in self.model_name else f"openai:{self.model_name}"
        self.pydantic_agent = Agent(
            model_str,
            system_prompt=self.system_prompt,
        )

    async def run(self, user_input: str) -> str:
        """Run a single interaction with the agent."""
        if self.pydantic_agent is None:
            await self.initialize()
        if self.pydantic_agent is None:
            raise RuntimeError("Model initialization failed: pydantic_agent is None")

        try:
            logger.info(f"Running agent interactively (input length: {len(user_input)})")
            result = await self.pydantic_agent.run(user_input)
            response: str = str(result.data)
            logger.info(f"Agent responded successfully (response length: {len(response)})")
            return response
        except Exception as e:
            logger.error(f"Agent execution error: {e}")
            return f"Error in agent execution: {str(e)}"

    async def run_with_tools(self, user_input: str, tools: list[Callable[..., Any]]) -> str:
        """Run the agent with tool calling capability using Pydantic AI."""
        if self.pydantic_agent is None:
            await self.initialize()

        try:
            logger.info(f"Calling Pydantic AI agent with {len(tools)} tools...")
            model_str: str = self.model_name if ":" in self.model_name else f"openai:{self.model_name}"
            agent: Agent[Any, str] = Agent(
                model_str,
                system_prompt=self.system_prompt,
                tools=tools,
            )
            result = await agent.run(user_input)
            response: str = str(result.data)
            return response
        except Exception as e:
            logger.error(f"Agent tool execution error: {e}")
            return f"Error in agent tool execution: {str(e)}"


class CavemanAgent(BaseAgent):
    """Caveman AI agent - minimal token usage variant."""

    def __init__(self, model_name: str) -> None:
        super().__init__(model_name)
        self.system_prompt: str = CAVEMAN_PROMPT