# ui/agent/graph/llm.py
from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Any

from huggingface_hub import InferenceClient
from langchain_core.messages import AIMessage
from langchain_core.runnables import RunnableConfig
from langchain_openai import ChatOpenAI

from ..completion import complete_turn
from ..config import INFERENCE_MODE, MODEL_ID, local_inference_enabled
from ..messages import api_turn_to_ai_message, langchain_messages_to_api

HF_ROUTER_BASE_URL = "https://router.huggingface.co/v1"


def _resolve_hf_token(configurable: dict[str, Any]) -> str:
    oauth_token = str(configurable.get("hf_token") or "").strip()
    if oauth_token:
        return oauth_token
    return os.environ.get("HF_TOKEN", "").strip()


@dataclass
class HubInferenceChatModel:
    """LangGraph-compatible wrapper around huggingface_hub InferenceClient."""

    hf_token: str
    max_tokens: int = 1800
    temperature: float = 0.35
    top_p: float = 0.9
    tools: list[dict[str, Any]] = field(default_factory=list)

    def bind_tools(self, tools: list[dict[str, Any]]) -> HubInferenceChatModel:
        return HubInferenceChatModel(
            hf_token=self.hf_token,
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            top_p=self.top_p,
            tools=list(tools),
        )

    def invoke(self, messages: list[Any]) -> AIMessage:
        client = InferenceClient(api_key=self.hf_token, model=MODEL_ID)
        api_messages = langchain_messages_to_api(messages)
        content, reasoning, tool_calls = complete_turn(
            client,
            api_messages,
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            top_p=self.top_p,
            tools=self.tools or None,
        )
        return api_turn_to_ai_message(content, reasoning, tool_calls)


@dataclass
class MiniCPMChatModel:
    max_tokens: int = 1800
    temperature: float = 0.35
    top_p: float = 0.9
    tools: list[dict[str, Any]] = field(default_factory=list)

    def bind_tools(self, tools: list[dict[str, Any]]) -> MiniCPMChatModel:
        return MiniCPMChatModel(
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            top_p=self.top_p,
            tools=list(tools),
        )

    def invoke(self, messages: list[Any]) -> AIMessage:
        from ..minicpm.model import chat_complete

        return chat_complete(
            messages,
            tools=self.tools or None,
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            top_p=self.top_p,
        )


def build_llm(
    config: RunnableConfig,
    **overrides: Any,
) -> HubInferenceChatModel | MiniCPMChatModel | ChatOpenAI:
    """Build a chat model from the per-request configurable values."""
    configurable = config.get("configurable", {})
    max_tokens = int(overrides.pop("max_tokens", configurable.get("max_tokens", 1800)))
    temperature = float(
        overrides.pop("temperature", configurable.get("temperature", 0.35))
    )
    top_p = float(overrides.pop("top_p", configurable.get("top_p", 0.9)))
    hf_token = _resolve_hf_token(configurable)

    if local_inference_enabled():
        return MiniCPMChatModel(
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )

    if INFERENCE_MODE == "router":
        params: dict[str, Any] = {
            "model": MODEL_ID,
            "api_key": hf_token,
            "base_url": HF_ROUTER_BASE_URL,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
        }
        params.update(overrides)
        return ChatOpenAI(**params)

    return HubInferenceChatModel(
        hf_token=hf_token,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )