"""HF Inference Endpoint custom handler for state-spaces/mamba-2.8b-hf.

Deploys pure-state-space Mamba via the standard HF Endpoints custom-handler
interface. The handler loads the base model at init time and serves a single
endpoint that accepts Bench 1.6-A's completion-format prompts.

Input schema (Bench 1.6-A concatenated completion format):
    {
        "inputs": "<flat text prompt with system + user turns concatenated>",
        "parameters": {
            "max_new_tokens": 512,
            "temperature": 0.7,
            "top_p": 0.95,
            "do_sample": true,
        }
    }

Output schema:
    {
        "generated_text": "<model completion>",
        "input_tokens": <int>,
        "output_tokens": <int>,
        "model": "state-spaces/mamba-2.8b-hf"
    }

Preregistered per docs/BENCH-1.6A-PREREG.md §5.5 (base-model asymmetry):
Base Mamba receives completion-format prompts, NOT chat-template formatted
messages. The caller (scripts/nsi_bench_hf.py) is responsible for
concatenating the Bench 1 [system, user_1, assistant_1, user_2, ...]
structure into the flat text the base model expects.
"""
from __future__ import annotations

from typing import Any

import torch
from transformers import AutoTokenizer, MambaForCausalLM

MODEL_ID = "state-spaces/mamba-2.8b-hf"


class EndpointHandler:
    """HF Endpoints custom handler entry point.

    HF Endpoints constructs this class once at boot and calls __call__ per
    request. The class name MUST be EndpointHandler.
    """

    def __init__(self, path: str = "") -> None:
        # HF Endpoints passes the model path in `path`. For our handler the
        # model ID is canonical and the weights are either shipped in `path`
        # or pulled from the hub. We prefer the hub ID so the handler is
        # independent of the deployment repo layout.
        self.model_id = MODEL_ID
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.bfloat16 if self.device == "cuda" else torch.float32

        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
        self.model = MambaForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype=dtype,
        ).to(self.device)
        self.model.eval()

    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
        prompt: str = data.get("inputs", "")
        params: dict[str, Any] = data.get("parameters", {}) or {}

        max_new_tokens: int = int(params.get("max_new_tokens", 512))
        temperature: float = float(params.get("temperature", 0.7))
        top_p: float = float(params.get("top_p", 0.95))
        do_sample: bool = bool(params.get("do_sample", True))

        if not prompt:
            return {
                "generated_text": "",
                "input_tokens": 0,
                "output_tokens": 0,
                "model": self.model_id,
                "error": "empty_input",
            }

        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        input_tokens = int(inputs["input_ids"].shape[-1])

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature if do_sample else 1.0,
                top_p=top_p,
                do_sample=do_sample,
                pad_token_id=self.tokenizer.eos_token_id
                if self.tokenizer.pad_token_id is None
                else self.tokenizer.pad_token_id,
            )

        full_text = self.tokenizer.decode(
            outputs[0],
            skip_special_tokens=True,
        )
        generated_only = full_text[len(prompt):] if full_text.startswith(prompt) else full_text
        output_tokens = int(outputs.shape[-1]) - input_tokens

        return {
            "generated_text": generated_only,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "model": self.model_id,
        }