"""HF Inference Endpoint custom handler for ai21labs/Jamba-v0.1.

Deploys hybrid SSM-Transformer MoE base model via the HF Endpoints custom-handler
interface. Jamba requires trust_remote_code=True and the mamba-ssm + causal-conv1d
dependencies. 52B total / 12B active parameters; needs A100 80GB minimum.

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": "ai21labs/Jamba-v0.1"
    }

Preregistered per docs/BENCH-1.6A-PREREG.md §5.5 (base-model asymmetry):
Base Jamba 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 AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "ai21labs/Jamba-v0.1"


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:
        self.model_id = MODEL_ID
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_id,
            trust_remote_code=True,
        )
        # BF16 with device_map=auto for A100 80GB. Per ai21labs model card,
        # this fits the model in a single 80GB GPU at BF16.
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_id,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
        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.model.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,
        }