"""Custom Hugging Face Inference Endpoints handler for ProCreations/grug-27b.

Deploy: create an Inference Endpoint from this repo with a GPU holding >=60GB
VRAM for bf16 (1x A100-80GB / H100 / H200). The endpoint auto-detects this
handler.

Request formats accepted:
  {"inputs": "plain user prompt", "parameters": {...}}
  {"inputs": {"messages": [{"role": "user", "content": "..."}],
              "tools": [...optional OpenAI-style tool schemas...]},
   "parameters": {"max_new_tokens": 1024, "temperature": 0.6, ...}}

Response:
  [{"generated_text": "<full raw completion>",
    "think": "<grug reasoning>", "answer": "<final answer>"}]
"""
from typing import Any, Dict, List

import torch
from transformers import AutoModelForImageTextToText, AutoTokenizer

DEFAULTS = {"max_new_tokens": 1024, "temperature": 0.6, "top_p": 0.95, "top_k": 20,
            "do_sample": True}
MAX_NEW_TOKENS_CAP = 8192


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForImageTextToText.from_pretrained(
            path, dtype=torch.bfloat16, device_map="auto",
            attn_implementation="sdpa")
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
        inputs = data.get("inputs", "")
        params = dict(DEFAULTS)
        params.update(data.get("parameters") or {})
        params["max_new_tokens"] = min(int(params.get("max_new_tokens", 1024)),
                                       MAX_NEW_TOKENS_CAP)
        if params.get("temperature", 1.0) <= 0:
            params["do_sample"] = False
            params.pop("temperature", None)
            params.pop("top_p", None)
            params.pop("top_k", None)

        tools = None
        if isinstance(inputs, dict):
            messages = inputs.get("messages") or [
                {"role": "user", "content": str(inputs.get("text", ""))}]
            tools = inputs.get("tools") or None
        elif isinstance(inputs, list):
            messages = inputs
        else:
            messages = [{"role": "user", "content": str(inputs)}]

        prompt = self.tokenizer.apply_chat_template(
            messages, tools=tools, add_generation_prompt=True, tokenize=False)
        enc = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        gen_kwargs = {k: v for k, v in params.items()
                      if k in ("max_new_tokens", "temperature", "top_p", "top_k",
                               "do_sample", "repetition_penalty", "seed")}
        gen_kwargs.pop("seed", None)
        with torch.inference_mode():
            out = self.model.generate(**enc, **gen_kwargs,
                                      pad_token_id=self.tokenizer.pad_token_id)
        text = self.tokenizer.decode(out[0][enc["input_ids"].shape[1]:],
                                     skip_special_tokens=True)
        if "</think>" in text:
            think, answer = text.split("</think>", 1)
            think = think.replace("<think>", "").strip()
            answer = answer.strip()
        else:
            think, answer = "", text.strip()
        return [{"generated_text": text, "think": think, "answer": answer}]