# ui/agent/minicpm/model.py
from __future__ import annotations

import logging
import os
import threading
from typing import TYPE_CHECKING, Any

from langchain_core.messages import AIMessage

from ..config import ENABLE_THINKING, MODEL_ID
from .messages import append_tool_instructions, normalize_messages

if TYPE_CHECKING:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

logger = logging.getLogger(__name__)

_GENERATE_LOCK = threading.Lock()
_MODEL: AutoModelForCausalLM | None = None
_TOKENIZER: AutoTokenizer | None = None
_DEVICE: torch.device | None = None


def _resolve_device() -> torch.device:
    import torch

    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _hub_login() -> None:
    from huggingface_hub import login

    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        login(token=hf_token)
        logger.info("Logged in to Hugging Face Hub for MiniCPM weights")
    else:
        logger.warning("HF_TOKEN not set — gated MiniCPM weights may be inaccessible")


def _load_model() -> tuple[AutoTokenizer, AutoModelForCausalLM]:
    global _MODEL, _TOKENIZER, _DEVICE
    from transformers import AutoModelForCausalLM, AutoTokenizer

    import torch

    device = _resolve_device()

    if _MODEL is not None and _TOKENIZER is not None and _DEVICE is not None:
        if device.type != _DEVICE.type:
            logger.info("Moving MiniCPM model from %s to %s", _DEVICE, device)
            _MODEL = _MODEL.to(device)
            _DEVICE = device
        return _TOKENIZER, _MODEL

    _hub_login()
    logger.info("Loading MiniCPM model %s on %s", MODEL_ID, device)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        dtype=torch.bfloat16,
        trust_remote_code=True,
    ).to(device)
    _TOKENIZER = tokenizer
    _MODEL = model
    _DEVICE = device
    return tokenizer, model


def _apply_chat_template(
    tokenizer: AutoTokenizer,
    messages: list[dict[str, str]],
    *,
    enable_thinking: bool,
) -> str:
    kwargs: dict[str, Any] = {
        "tokenize": False,
        "add_generation_prompt": True,
    }
    try:
        return tokenizer.apply_chat_template(
            messages,
            enable_thinking=enable_thinking,
            **kwargs,
        )
    except TypeError:
        return tokenizer.apply_chat_template(messages, **kwargs)


def _split_think_output(text: str) -> tuple[str, str]:
    open_tag = "<" + "think" + ">"
    close_tag = "</" + "think" + ">"
    start = text.find(open_tag)
    end = text.find(close_tag)
    if start != -1 and end != -1 and end > start:
        reasoning = text[start + len(open_tag) : end].strip()
        content = (text[:start] + text[end + len(close_tag) :]).strip()
        return content, reasoning
    return text.strip(), ""


def chat_complete(
    messages: list[Any],
    *,
    tools: list[dict[str, Any]] | None = None,
    max_tokens: int = 1800,
    temperature: float = 0.35,
    top_p: float = 0.9,
    enable_thinking: bool | None = None,
) -> AIMessage:
    """Run one MiniCPM chat turn and return a LangChain AIMessage."""
    tokenizer, model = _load_model()
    assert _DEVICE is not None
    normalized = normalize_messages(messages)
    if tools:
        normalized = append_tool_instructions(normalized, tools)

    thinking = ENABLE_THINKING if enable_thinking is None else enable_thinking
    prompt_text = _apply_chat_template(tokenizer, normalized, enable_thinking=thinking)
    model_inputs = tokenizer([prompt_text], return_tensors="pt").to(_DEVICE)

    gen_kwargs: dict[str, Any] = {
        **model_inputs,
        "max_new_tokens": max_tokens,
    }
    if temperature > 0:
        gen_kwargs.update(
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
        )
    else:
        gen_kwargs["do_sample"] = False

    with _GENERATE_LOCK:
        output_ids = model.generate(**gen_kwargs)

    generated = output_ids[0][model_inputs["input_ids"].shape[1] :]
    raw_text = tokenizer.decode(generated, skip_special_tokens=False)
    content, reasoning = _split_think_output(raw_text)

    additional_kwargs: dict[str, Any] = {}
    if reasoning:
        additional_kwargs["reasoning_content"] = reasoning

    return AIMessage(content=content or raw_text, additional_kwargs=additional_kwargs)