import json
import logging
from abc import ABC, abstractmethod
from contextlib import AsyncExitStack
from functools import cached_property
from typing import Any, Optional, Type, cast

from pydantic import BaseModel, Field
from tenacity import before_sleep_log, retry, stop_after_attempt, wait_exponential

from proxy_lite.client import BaseClient, ClientConfigTypes, OpenAIClientConfig
from proxy_lite.history import (
    AssistantMessage,
    MessageHistory,
    MessageLabel,
    SystemMessage,
    Text,
    ToolCall,
    ToolMessage,
    UserMessage,
)
from proxy_lite.logger import logger
from proxy_lite.tools import Tool

# if TYPE_CHECKING:
#     from proxy_lite.tools import Tool


class BaseAgentConfig(BaseModel):
    client: ClientConfigTypes = Field(default_factory=OpenAIClientConfig)
    history_messages_limit: dict[MessageLabel, int] = Field(default_factory=lambda: dict())
    history_messages_include: Optional[dict[MessageLabel, int]] = Field(
        default=None,
        description="If set, overrides history_messages_limit by setting all message types to 0 except those specified",
    )

    def model_post_init(self, __context: Any) -> None:
        if self.history_messages_include is not None:
            self.history_messages_limit = {label: 0 for label in MessageLabel}
            self.history_messages_limit.update(self.history_messages_include)


class BaseAgent(BaseModel, ABC):
    config: BaseAgentConfig
    temperature: float = Field(default=0.7, ge=0, le=2)
    history: MessageHistory = Field(default_factory=MessageHistory)
    client: Optional[BaseClient] = None
    env_tools: list[Tool] = Field(default_factory=list)
    task: Optional[str] = Field(default=None)
    seed: Optional[int] = Field(default=None)

    class Config:
        arbitrary_types_allowed = True

    def __init__(self, **data) -> None:
        super().__init__(**data)
        self._exit_stack = AsyncExitStack()
        self._tools_init_task = None

    def model_post_init(self, __context: Any) -> None:
        super().model_post_init(__context)
        self.client = BaseClient.create(self.config.client)

    @property
    @abstractmethod
    def system_prompt(self) -> str: ...

    @cached_property
    @abstractmethod
    def tools(self) -> list[Tool]: ...

    @cached_property
    def tool_descriptions(self) -> str:
        tool_descriptions = []
        for tool in self.tools:
            func_descriptions = "\n".join("- {name}: {description}".format(**schema) for schema in tool.schema)
            tool_title = f"{tool.__class__.__name__}:\n" if len(self.tools) > 1 else ""
            tool_descriptions.append(f"{tool_title}{func_descriptions}")
        return "\n\n".join(tool_descriptions)

    async def get_history_view(self) -> MessageHistory:
        return MessageHistory(
            messages=[SystemMessage(content=[Text(text=self.system_prompt)])],
        ) + self.history.history_view(
            limits=self.config.history_messages_limit,
        )

    @retry(
        wait=wait_exponential(multiplier=1, min=4, max=10),
        stop=stop_after_attempt(3),
        reraise=True,
        before_sleep=before_sleep_log(logger, logging.ERROR),
    )
    async def generate_output(
        self,
        use_tool: bool = False,
        response_format: Optional[type[BaseModel]] = None,
        append_assistant_message: bool = True,
    ) -> AssistantMessage:
        messages: MessageHistory = await self.get_history_view()
        response_content = (
            await self.client.create_completion(
                messages=messages,
                temperature=self.temperature,
                seed=self.seed,
                response_format=response_format,
                tools=self.tools if use_tool else None,
            )
        ).model_dump()
        response_content = response_content["choices"][0]["message"]
        assistant_message = AssistantMessage(
            role=response_content["role"],
            content=[Text(text=response_content["content"])] if response_content["content"] else [],
            tool_calls=response_content["tool_calls"],
        )
        if append_assistant_message:
            self.history.append(message=assistant_message, label=self.message_label)
        return assistant_message

    def receive_user_message(
        self,
        text: Optional[str] = None,
        image: list[bytes] = None,
        label: MessageLabel = None,
        is_base64: bool = False,
    ) -> None:
        message = UserMessage.from_media(
            text=text,
            image=image,
            is_base64=is_base64,
        )
        self.history.append(message=message, label=label)

    def receive_system_message(
        self,
        text: Optional[str] = None,
        label: MessageLabel = None,
    ) -> None:
        message = SystemMessage.from_media(text=text)
        self.history.append(message=message, label=label)

    def receive_assistant_message(
        self,
        content: Optional[str] = None,
        tool_calls: Optional[list[ToolCall]] = None,
        label: MessageLabel = None,
    ) -> None:
        message = AssistantMessage(
            content=[Text(text=content)] if content else [],
            tool_calls=tool_calls,
        )
        self.history.append(message=message, label=label)

    async def use_tool(self, tool_call: ToolCall):
        function = tool_call.function
        for tool in self.tools:
            if hasattr(tool, function["name"]):
                return await getattr(tool, function["name"])(
                    **json.loads(function["arguments"]),
                )
        msg = f'No tool function with name "{function["name"]}"'
        raise ValueError(msg)

    async def receive_tool_message(
        self,
        text: str,
        tool_id: str,
        label: MessageLabel = None,
    ) -> None:
        self.history.append(
            message=ToolMessage(content=[Text(text=text)], tool_call_id=tool_id),
            label=label,
        )


class Agents:
    _agent_registry: dict[str, type[BaseAgent]] = {}
    _agent_config_registry: dict[str, type[BaseAgentConfig]] = {}

    @classmethod
    def register_agent(cls, name: str):
        """
        Decorator to register an Agent class under a given name.

        Example:
            @Agents.register_agent("browser")
            class BrowserAgent(BaseAgent):
                ...
        """

        def decorator(agent_cls: type[BaseAgent]) -> type[BaseAgent]:
            cls._agent_registry[name] = agent_cls
            return agent_cls

        return decorator

    @classmethod
    def register_agent_config(cls, name: str):
        """
        Decorator to register a configuration class under a given name.

        Example:
            @Agents.register_agent_config("browser")
            class BrowserAgentConfig(BaseAgentConfig):
                ...
        """

        def decorator(config_cls: type[BaseAgentConfig]) -> type[BaseAgentConfig]:
            cls._agent_config_registry[name] = config_cls
            return config_cls

        return decorator

    @classmethod
    def get(cls, name: str) -> type[BaseAgent]:
        """
        Retrieve a registered Agent class by its name.

        Raises:
            ValueError: If no such agent is found.
        """
        try:
            return cast(Type[BaseAgent], cls._agent_registry[name])
        except KeyError:
            raise ValueError(f"Agent '{name}' not found.")

    @classmethod
    def get_config(cls, name: str) -> type[BaseAgentConfig]:
        """
        Retrieve a registered Agent configuration class by its name.

        Raises:
            ValueError: If no such config is found.
        """
        try:
            return cast(type[BaseAgentConfig], cls._agent_config_registry[name])
        except KeyError:
            raise ValueError(f"Agent config for '{name}' not found.")