"""
Hugging Face Inference Endpoint custom handler.

Loads a base model with multiple LoRA adapters (one per mode).
Adapters stay resident in memory; set_adapter() switches cheaply per request.
The model is trained to emit  <ui>{...json...}</ui>  then prose.
"""
import json
import os
import re
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE_MODEL = os.getenv("BASE_MODEL", "Qwen/Qwen2.5-7B-Instruct")

# Override via env vars to point at your own adapter repos on the Hub
ADAPTERS: dict[str, str] = {
    "support":   os.getenv("ADAPTER_SUPPORT",   "your-org/adapter-support"),
    "analytics": os.getenv("ADAPTER_ANALYTICS", "your-org/adapter-analytics"),
    "form":      os.getenv("ADAPTER_FORM",       "your-org/adapter-form"),
}

_UI_RE = re.compile(r"<ui>(.*?)</ui>", re.DOTALL)


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.tok = AutoTokenizer.from_pretrained(BASE_MODEL)

        base = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )

        names = list(ADAPTERS.keys())
        # First adapter via from_pretrained so PEFT wraps the model
        self.model = PeftModel.from_pretrained(
            base, ADAPTERS[names[0]], adapter_name=names[0]
        )
        # Remaining adapters attached in-place
        for name in names[1:]:
            self.model.load_adapter(ADAPTERS[name], adapter_name=name)

        self.model.eval()

    def __call__(self, data: dict) -> dict:
        inp      = data.get("inputs", {})
        messages = inp.get("messages", [])
        mode     = inp.get("mode", "support")

        if mode not in ADAPTERS:
            mode = "support"

        self.model.set_adapter(mode)

        prompt = self.tok.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False
        )
        ids = self.tok(prompt, return_tensors="pt").to(self.model.device)

        with torch.no_grad():
            out = self.model.generate(
                **ids,
                max_new_tokens=512,
                do_sample=False,
            )

        new_tokens = out[0][ids.input_ids.shape[1]:]
        text = self.tok.decode(new_tokens, skip_special_tokens=True)

        ui_spec, clean_text = _extract_ui_spec(text)

        return {"text": clean_text, "ui_spec": ui_spec, "adapter": mode}


def _extract_ui_spec(text: str) -> tuple[dict | None, str]:
    m = _UI_RE.search(text)
    if not m:
        return None, text

    try:
        spec = json.loads(m.group(1).strip())
    except json.JSONDecodeError:
        return None, text

    return spec, _UI_RE.sub("", text).strip()