from __future__ import annotations

from typing import Any, Dict, List, Optional

from aiel_sdk.errors import RuntimeDependencyError

try:
    from langchain_core.language_models.chat_models import BaseChatModel
    from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
    from langchain_core.outputs import ChatGeneration, ChatResult
except Exception as e:  # pragma: no cover
    raise RuntimeDependencyError("langchain-core", "pip install 'aiel-sdk[langchain]'") from e

try:
    from pydantic import PrivateAttr
except Exception as e:  # pragma: no cover
    raise RuntimeDependencyError("pydantic", "pip install pydantic") from e

from ...core.protocols import ChatModel
from ...core.types import Message


def _from_langchain_message(msg: BaseMessage) -> Message:
    if isinstance(msg, HumanMessage):
        return Message(role="user", content=msg.content)
    if isinstance(msg, SystemMessage):
        return Message(role="system", content=msg.content)
    if isinstance(msg, ToolMessage):
        return Message(role="tool", content=msg.content, tool_call_id=msg.tool_call_id)
    return Message(role="assistant", content=msg.content)


class AIELLangChainChatModel(BaseChatModel):
    model_config = {"arbitrary_types_allowed": True}

    _inner: ChatModel = PrivateAttr()

    def __init__(self, inner: ChatModel, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._inner = inner

    @property
    def _llm_type(self) -> str:
        return "aiel-chat-model"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        return {}

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        extra = dict(kwargs)
        if stop is not None:
            extra["stop"] = stop

        response = self._inner.invoke(
            [_from_langchain_message(m) for m in messages],
            **extra,
        )

        ai_message = AIMessage(
            content=response.message.content,
            additional_kwargs={"raw": response.raw},
        )

        return ChatResult(generations=[ChatGeneration(message=ai_message)])
