"""Pinned text-only inference. Model weights are loaded only on a cloud GPU."""
from __future__ import annotations

import json
import os
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
CONFIG = json.loads(Path(__file__).with_name("config.json").read_text(encoding="utf-8"))
os.environ.setdefault("HF_HOME", str(ROOT / ".cache/huggingface"))
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false")


def conversation(prompt: str, history: list[dict] | None = None) -> list[dict]:
    if not prompt.strip():
        raise ValueError("Enter a message first.")
    if len(prompt) > 24000:
        raise ValueError("Please split this message into smaller parts.")
    messages = [{"role": "system", "content": CONFIG["system_prompt"]}]
    previous = history or []
    if len(previous) % 2:
        raise ValueError("Conversation history is incomplete. Start a new chat.")
    for index, item in enumerate(previous):
        if item.get("role") != ("user" if index % 2 == 0 else "assistant"):
            raise ValueError("Conversation history is invalid. Start a new chat.")
        if not isinstance(item.get("content"), str):
            raise ValueError("This chat accepts text messages only.")
        messages.append({"role": item["role"], "content": item["content"]})
    messages.append({"role": "user", "content": prompt.strip()})
    return messages


def gpu_available() -> bool:
    import torch
    return bool(torch.cuda.is_available())


def load_engine(adapter: str | None = None):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if not gpu_available():
        raise RuntimeError("The cloud model is not connected. This computer shows the interface preview only.")
    tokenizer = AutoTokenizer.from_pretrained(
        CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False)
    model = AutoModelForCausalLM.from_pretrained(
        CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False,
        dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
        device_map={"": "cuda:0"}, attn_implementation="sdpa")
    if adapter:
        from peft import PeftConfig, PeftModel
        adapter_config = PeftConfig.from_pretrained(adapter)
        if (adapter_config.base_model_name_or_path != CONFIG["base_model"]
                or adapter_config.revision != CONFIG["base_revision"]):
            raise ValueError("Adapter does not match the pinned chat-and-code model.")
        model = PeftModel.from_pretrained(model, adapter)
    return tokenizer, model.eval()


def answer(tokenizer, model, prompt: str, history: list[dict] | None = None,
           max_new_tokens: int = CONFIG["max_new_tokens"]) -> str:
    if not 1 <= max_new_tokens <= 2048:
        raise ValueError("Response length must be between 1 and 2,048 tokens.")
    messages = conversation(prompt, history)
    while True:
        inputs = tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt")
        if inputs["input_ids"].shape[1] + max_new_tokens <= 4096:
            break
        if len(messages) <= 2:
            raise ValueError("This message is too long. Please send a smaller code example.")
        del messages[1:3]
    return generate_messages(tokenizer, model, messages, max_new_tokens=max_new_tokens)


def generate_messages(tokenizer, model, messages, max_new_tokens=768,
                      context_limit=4096, do_sample=True):
    import torch
    if not 1 <= max_new_tokens <= 2048 or not 256 <= context_limit <= 8192:
        raise ValueError("Invalid response or context length.")
    inputs = tokenizer.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt")
    if inputs["input_ids"].shape[1] + max_new_tokens > context_limit:
        raise ValueError("Source packet is too long. Shorten the excerpts; no sources were silently removed.")
    inputs = inputs.to(model.device)
    sampling = {"temperature": 0.7, "top_p": 0.8, "top_k": 20} if do_sample else {}
    with torch.inference_mode():
        output = model.generate(
            **inputs, max_new_tokens=max_new_tokens, do_sample=do_sample, **sampling,
            pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)
    reply = tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
    if not reply:
        raise RuntimeError("The model returned an empty response. Please try again.")
    return reply