from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator
from dataclasses import dataclass
from enum import StrEnum, auto

from donkit.llm import (
    GenerateRequest,
    LLMModelAbstract,
    Message,
    ModelCapability,
    Tool,
)
from loguru import logger

from donkit.llm_agent.history import HistoryCompressor
from donkit.llm_agent.protocol import MCPClientProtocol
from donkit.llm_agent.tool import AgentTool


class EventType(StrEnum):
    CONTENT = auto()
    TOOL_CALL_START = auto()
    TOOL_CALL_END = auto()
    TOOL_CALL_ERROR = auto()
    HISTORY_COMPRESSED = auto()


@dataclass
class StreamEvent:
    """Event yielded during streaming response."""

    type: EventType
    content: str | None = None
    tool_name: str | None = None
    tool_args: dict | None = None
    tool_result: str | None = None
    error: str | None = None


class LLMAgent:
    def __init__(
        self,
        provider: LLMModelAbstract,
        tools: list[AgentTool] | None = None,
        mcp_clients: list[MCPClientProtocol] | None = None,
        max_iterations: int = 500,
        history_compressor: HistoryCompressor | None = None,
    ) -> None:
        self.provider = provider
        self.local_tools = tools or []
        self.mcp_clients = mcp_clients or []
        self.mcp_tools: dict[str, tuple[dict, MCPClientProtocol]] = {}
        self.max_iterations = max_iterations
        self.history_compressor = history_compressor or HistoryCompressor()

    async def ainit_mcp_tools(self) -> None:
        """Initialize MCP tools asynchronously. Call this after creating the agent."""
        for client in self.mcp_clients:
            try:
                await client.connect()
                discovered = await client.alist_tools()
                for t in discovered:
                    self.mcp_tools[t["name"]] = (t, client)
                logger.debug(
                    "[MCP] Registered {} tools from {}",
                    len(discovered),
                    client.identifier,
                )
            except Exception as ex:
                logger.error(
                    "Failed to load MCP tools from {}: {}", client.identifier, ex
                )
                try:
                    await client.disconnect()
                except Exception:
                    pass

    def _tool_specs(self) -> list[Tool]:
        local_specs = [t.to_tool_spec() for t in self.local_tools]
        mcp_specs = [
            Tool(
                **{
                    "function": {
                        "name": info["name"],
                        "description": info["description"],
                        "parameters": info["parameters"],
                    }
                }
            )
            for info, _ in self.mcp_tools.values()
        ]
        return local_specs + mcp_specs

    def _find_tool(
        self, name: str
    ) -> tuple[AgentTool | None, tuple[dict, MCPClientProtocol] | None]:
        for t in self.local_tools:
            if t.name == name:
                return t, None
        if name in self.mcp_tools:
            return None, self.mcp_tools[name]
        return None, None

    def _should_execute_tools(self, resp) -> bool:
        return bool(resp.tool_calls)

    def _append_synthetic_assistant_turn(
        self, messages: list[Message], tool_calls, content: str | None = None
    ) -> None:
        messages.append(
            Message(
                role="assistant",
                content=content or None,
                tool_calls=tool_calls,
            )
        )

    def _parse_tool_args(self, tc) -> dict:
        try:
            raw = tc.function.arguments
            if isinstance(raw, dict):
                return raw
            return json.loads(raw or "{}")
        except Exception as e:
            logger.error("Failed to parse tool arguments: {}", e)
            return {}

    async def _aexecute_tool_call(self, tc, args: dict) -> str:
        try:
            local_tool, mcp_tool_info = self._find_tool(tc.function.name)
            if not local_tool and not mcp_tool_info:
                logger.warning("Tool not found: {}", tc.function.name)
                return f"Error: Tool '{tc.function.name}' not found."

            if local_tool:
                if local_tool.is_async:
                    result = await local_tool.handler(args)
                else:
                    result = local_tool.handler(args)
            else:
                tool_meta, client = mcp_tool_info
                result = await client.acall_tool(tool_meta["name"], args)

        except KeyboardInterrupt:
            return "Tool execution cancelled by user (Ctrl+C)"
        except asyncio.CancelledError:
            return "Tool execution cancelled by user (Ctrl+C)"
        except Exception as e:
            logger.error("Tool {} error: {}", tc.function.name, e, exc_info=True)
            return f"Error: {e!s}"

        return self._serialize_tool_result(result)

    def _serialize_tool_result(self, result) -> str:
        if isinstance(result, str):
            return result
        try:
            return json.dumps(result)
        except Exception:
            return str(result)

    async def _aexecute_and_record(
        self, tc, args: dict, messages: list[Message]
    ) -> str:
        """Execute a tool call, append the result message, and compress history.

        Returns the result string. Raises on tool execution errors
        so callers can handle them (e.g. emit error events in streaming).
        """
        result_str = await self._aexecute_tool_call(tc, args)
        messages.append(
            Message(
                role="tool",
                name=tc.function.name,
                tool_call_id=tc.id,
                content=result_str,
            )
        )
        messages[:] = await self.history_compressor.compress_if_needed(
            messages, self.provider
        )
        return result_str

    async def _ahandle_tool_calls(
        self, messages: list[Message], tool_calls, content: str | None = None
    ) -> None:
        self._append_synthetic_assistant_turn(messages, tool_calls, content=content)
        for tc in tool_calls:
            args = self._parse_tool_args(tc)
            await self._aexecute_and_record(tc, args, messages)

    async def arespond(self, messages: list[Message]) -> str:
        """Agentic loop: generate, execute tool calls, repeat until text response."""
        tools = (
            self._tool_specs()
            if self.provider.supports_capability(ModelCapability.TOOL_CALLING)
            else None
        )

        for _ in range(self.max_iterations):
            messages[:] = await self.history_compressor.compress_if_needed(
                messages, self.provider
            )
            resp = await self.provider.generate(
                GenerateRequest(messages=messages, tools=tools)
            )

            if self._should_execute_tools(resp):
                await self._ahandle_tool_calls(
                    messages, resp.tool_calls, content=resp.content
                )
                continue

            if not resp.content:
                retry_resp = await self.provider.generate(
                    GenerateRequest(messages=messages)
                )
                return retry_resp.content or ""
            return resp.content

        logger.warning(
            "Agent exhausted max_iterations ({}) without producing a final response",
            self.max_iterations,
        )
        return ""

    async def arespond_stream(
        self, messages: list[Message]
    ) -> AsyncIterator[StreamEvent]:
        """Agentic loop with streaming: yields StreamEvent for content and tool calls."""
        tools = (
            self._tool_specs()
            if self.provider.supports_capability(ModelCapability.TOOL_CALLING)
            else None
        )

        for _ in range(self.max_iterations):
            prev_len = len(messages)
            messages[:] = await self.history_compressor.compress_if_needed(
                messages, self.provider
            )
            if len(messages) < prev_len:
                yield StreamEvent(type=EventType.HISTORY_COMPRESSED)

            streamed_content = ""
            saw_tool_calls = False

            async for chunk in self.provider.generate_stream(
                GenerateRequest(messages=messages, tools=tools)
            ):
                if chunk.content:
                    streamed_content += chunk.content
                    yield StreamEvent(type=EventType.CONTENT, content=chunk.content)

                if chunk.tool_calls:
                    saw_tool_calls = True
                    self._append_synthetic_assistant_turn(
                        messages, chunk.tool_calls, content=streamed_content or None
                    )
                    for tc in chunk.tool_calls:
                        args = self._parse_tool_args(tc)
                        yield StreamEvent(
                            type=EventType.TOOL_CALL_START,
                            tool_name=tc.function.name,
                            tool_args=args,
                        )
                        try:
                            prev_len_inter = len(messages)
                            result_str = await self._aexecute_and_record(tc, args, messages)
                            yield StreamEvent(
                                type=EventType.TOOL_CALL_END,
                                tool_name=tc.function.name,
                                tool_result=result_str,
                            )
                            if len(messages) < prev_len_inter + 1:
                                yield StreamEvent(type=EventType.HISTORY_COMPRESSED)
                        except Exception as e:
                            error_msg = str(e)
                            messages.append(
                                Message(
                                    role="tool",
                                    name=tc.function.name,
                                    tool_call_id=tc.id,
                                    content=f"Error: {error_msg}",
                                )
                            )
                            yield StreamEvent(
                                type=EventType.TOOL_CALL_ERROR,
                                tool_name=tc.function.name,
                                error=error_msg,
                            )

            if not saw_tool_calls:
                return

        logger.warning(
            "Agent exhausted max_iterations ({}) without producing a final response (stream)",
            self.max_iterations,
        )
