import asyncio
import logging
import lmstudio as lms
from typing import Any, Callable
from core.interfaces import Agent
from .prompt import CAVEMAN_PROMPT

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


class _ActResponseCapture:
    """Captures the AI response from LMStudio act() callback."""

    def __init__(self) -> None:
        self.responses: list[str] = []

    def __call__(self, message: Any) -> None:
        content: str = ""
        if hasattr(message, 'content'):
            content = message.content
        elif hasattr(message, 'text'):
            content = message.text
        elif hasattr(message, 'response'):
            content = message.response
        elif hasattr(message, 'message'):
            content = message.message
        else:
            return

        if isinstance(content, list):
            parts: list[str] = []
            for item in content:
                if isinstance(item, dict):
                    text: str = item.get('text', '')
                    if isinstance(text, list):
                        parts.extend([str(t) for t in text])
                    else:
                        parts.append(str(text))
                elif isinstance(item, str):
                    parts.append(item)
                elif hasattr(item, 'text'):
                    parts.append(str(item.text))
                elif hasattr(item, 'content'):
                    parts.append(str(item.content))
            content = ''.join(parts)
        elif not isinstance(content, str):
            content = str(content)

        if content.strip():
            self.responses.append(content.strip())

    @property
    def full_response(self) -> str:
        return '\n'.join(self.responses) if self.responses else "No response captured."


class BaseAgent(Agent):
    """Base AI agent implementing common LMStudio interaction patterns."""

    def __init__(self, model_name: str) -> None:
        self.model_name: str = model_name
        self.model: Any | None = None
        self.system_prompt: str = ""

    async def initialize(self) -> None:
        """Initialize the LM Studio model."""
        logger.info(f"Initializing agent with model: {self.model_name}")
        self.model = lms.llm(self.model_name)

    async def run(self, user_input: str) -> str:
        """Run a single interaction with the agent."""
        if self.model is None:
            await self.initialize()
        assert self.model is not None

        messages: list[dict[str, str]] = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_input},
        ]

        try:
            logger.info(f"Running agent interactively (input length: {len(user_input)})")
            response = await self.model.respond(user_input, messages=messages)
            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[Any]) -> str:
        """Run the agent with tool calling capability."""
        if self.model is None:
            await self.initialize()
        assert self.model is not None

        try:
            capture = _ActResponseCapture()
            logger.info(f"Calling LMStudio act() on agent with {len(tools)} tools...")
            result: lms.ActResult = self.model.act(user_input, tools=tools, on_message=capture)
            logger.info(f"act() on agent returned: {result}")
            response: str = capture.full_response
            if not response or response == "No response captured.":
                logger.warning(f"Act completed with {result.rounds} rounds but no response was captured.")
                return f"Act completed with {result.rounds} rounds but no response captured."
            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 = CAVEMAN_PROMPT