"""
Custom Inference Endpoint handler for Gemma4 text (drive-thru merged model).
Requires a recent transformers build that knows `gemma4_text` / Gemma4ForCausalLM.
"""

from __future__ import annotations

from typing import Any, Dict, List

import torch
from transformers import AutoTokenizer, Gemma4ForCausalLM


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.model = Gemma4ForCausalLM.from_pretrained(
            path,
            torch_dtype=dtype,
            device_map="auto" if self.device == "cuda" else None,
        )
        if self.device == "cpu":
            self.model = self.model.to(self.device)
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
        # OpenAI-style or classic HF inputs
        if "inputs" in data:
            prompt = data["inputs"]
            params = data.get("parameters") or {}
        elif "messages" in data:
            messages = data["messages"]
            system = ""
            user = ""
            for m in messages:
                if m.get("role") == "system":
                    system = m.get("content") or ""
                elif m.get("role") == "user":
                    user = m.get("content") or ""
            prompt = (
                f"<bos><start_of_turn>system\n{system}<end_of_turn>\n"
                f"<start_of_turn>user\n{user}<end_of_turn>\n"
                f"<start_of_turn>model\n"
            )
            params = {
                "max_new_tokens": data.get("max_tokens") or data.get("max_new_tokens") or 120,
                "temperature": data.get("temperature", 0.2),
            }
        else:
            prompt = str(data)
            params = {}

        max_new = int(params.get("max_new_tokens") or params.get("max_tokens") or 120)
        temperature = float(params.get("temperature", 0.2))
        do_sample = temperature > 0

        inputs = self.tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
        with torch.no_grad():
            out = self.model.generate(
                **inputs,
                max_new_tokens=max_new,
                do_sample=do_sample,
                temperature=max(temperature, 1e-5) if do_sample else None,
                pad_token_id=self.tokenizer.pad_token_id,
            )
        text = self.tokenizer.decode(
            out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
        )
        text = text.split("<end_of_turn>")[0].strip()
        return [{"generated_text": text}]