"""LangChain integration: RotatorChatModel(BaseChatModel)."""

from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
from typing import Any

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_core.runnables import Runnable
from langchain_core.tools import BaseTool

from llm_rotator._types import LLMResponse, RoutingContext, StreamChunk, ToolCall
from llm_rotator.rotator import LLMRotator


def _to_openai_messages(messages: list[BaseMessage]) -> list[dict[str, Any]]:
    """Convert LangChain messages to OpenAI-format dicts."""
    result: list[dict[str, Any]] = []
    for msg in messages:
        if isinstance(msg, ToolMessage):
            result.append(
                {
                    "role": "tool",
                    "content": msg.content,
                    "tool_call_id": msg.tool_call_id,
                }
            )
        elif isinstance(msg, SystemMessage):
            result.append({"role": "system", "content": msg.content})
        elif isinstance(msg, HumanMessage):
            result.append({"role": "user", "content": msg.content})
        else:
            # AIMessage or any other — treat as assistant
            entry: dict[str, Any] = {"role": "assistant", "content": msg.content}
            if hasattr(msg, "tool_calls") and msg.tool_calls:
                entry["tool_calls"] = [
                    {
                        "id": tc["id"],
                        "type": "function",
                        "function": {
                            "name": tc["name"],
                            "arguments": json.dumps(tc["args"]),
                        },
                    }
                    for tc in msg.tool_calls
                ]
            result.append(entry)
    return result


def _convert_tool_calls(tool_calls: list[ToolCall]) -> list[dict[str, Any]]:
    """Convert llm_rotator ToolCall list to LangChain tool_calls format."""
    return [
        {
            "name": tc.name,
            "args": json.loads(tc.arguments),
            "id": tc.id,
        }
        for tc in tool_calls
    ]


def _to_chat_result(response: LLMResponse) -> ChatResult:
    """Convert LLMResponse to LangChain ChatResult."""
    kwargs: dict[str, Any] = {}

    if response.tool_calls:
        kwargs["tool_calls"] = _convert_tool_calls(response.tool_calls)

    kwargs["usage_metadata"] = {
        "input_tokens": response.usage.prompt_tokens,
        "output_tokens": response.usage.completion_tokens,
        "total_tokens": response.usage.total_tokens,
    }

    message = AIMessage(
        content=response.content or "",
        **kwargs,
    )

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


def _to_chat_generation_chunk(chunk: StreamChunk) -> ChatGenerationChunk:
    """Convert StreamChunk to LangChain ChatGenerationChunk."""
    kwargs: dict[str, Any] = {}

    if chunk.tool_calls:
        kwargs["tool_calls"] = _convert_tool_calls(chunk.tool_calls)

    if chunk.usage:
        kwargs["usage_metadata"] = {
            "input_tokens": chunk.usage.prompt_tokens,
            "output_tokens": chunk.usage.completion_tokens,
            "total_tokens": chunk.usage.total_tokens,
        }

    message = AIMessageChunk(content=chunk.delta, **kwargs)
    return ChatGenerationChunk(message=message)


class RotatorChatModel(BaseChatModel):
    """LangChain ChatModel backed by LLMRotator.

    Usage::

        from llm_rotator import LLMRotator, RotatorChatModel

        rotator = LLMRotator(config=...)
        model = RotatorChatModel(rotator=rotator)
        response = model.invoke([HumanMessage(content="Hello")])
    """

    rotator: LLMRotator
    default_routing: RoutingContext | None = None

    @property
    def _llm_type(self) -> str:
        return "llm-rotator"

    def bind_tools(
        self,
        tools: Sequence[dict[str, Any] | type | Callable | BaseTool],
        *,
        tool_choice: str | None = None,
        **kwargs: Any,
    ) -> Runnable:
        """Bind tools to the model, forwarding them to the rotator."""
        from langchain_core.utils.function_calling import convert_to_openai_tool

        formatted = [convert_to_openai_tool(t) for t in tools]
        bind_kwargs: dict[str, Any] = {"tools": formatted, **kwargs}
        if tool_choice is not None:
            bind_kwargs["tool_choice"] = tool_choice
        return self.bind(**bind_kwargs)

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        loop = asyncio.new_event_loop()
        try:
            return loop.run_until_complete(self._agenerate(messages, stop, run_manager, **kwargs))
        finally:
            loop.close()

    async def _agenerate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        routing = kwargs.pop("routing", self.default_routing)
        openai_msgs = _to_openai_messages(messages)

        extra: dict[str, Any] = {}
        if stop:
            extra["stop"] = stop
        # Forward tools/tool_choice/response_format if present
        for key in ("tools", "tool_choice", "response_format"):
            if key in kwargs:
                extra[key] = kwargs.pop(key)
        # Pass remaining kwargs
        extra.update(kwargs)

        response = await self.rotator.complete(messages=openai_msgs, routing=routing, **extra)
        return _to_chat_result(response)

    def _stream(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        loop = asyncio.new_event_loop()
        try:
            ait = self._astream(messages, stop, run_manager, **kwargs)
            while True:
                try:
                    yield loop.run_until_complete(ait.__anext__())
                except StopAsyncIteration:
                    break
        finally:
            loop.close()

    async def _astream(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> AsyncIterator[ChatGenerationChunk]:
        routing = kwargs.pop("routing", self.default_routing)
        openai_msgs = _to_openai_messages(messages)

        extra: dict[str, Any] = {}
        if stop:
            extra["stop"] = stop
        for key in ("tools", "tool_choice", "response_format"):
            if key in kwargs:
                extra[key] = kwargs.pop(key)
        extra.update(kwargs)

        async for chunk in self.rotator.stream(messages=openai_msgs, routing=routing, **extra):
            yield _to_chat_generation_chunk(chunk)
