# handler.py
from __future__ import annotations
import os
from typing import Any, Dict, List, Union

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

GEN_KW = {
    "temperature": float(os.getenv("GEN_TEMPERATURE", "0.2")),
    "do_sample": os.getenv("GEN_DO_SAMPLE", "false").lower() == "true",
    "max_new_tokens": int(os.getenv("GEN_MAX_NEW_TOKENS", "512")),
    "repetition_penalty": float(os.getenv("GEN_REP_PENALTY", "1.0")),
}

PROMPT_PREFIX = (
    "Преобразуй дореформенный русский текст в современную орфографию, "
    "сохранив смысл и пунктуацию. Верни только преобразованный текст.\n\nТекст:\n"
)
PROMPT_SUFFIX = "\n\nСовременный вариант:"

def _to_list(x: Union[str, List[str]]) -> List[str]:
    if isinstance(x, list):
        return [str(t) for t in x]
    return [str(x)]

class EndpointHandler:
    def __init__(self, model_dir: str):
        # model_dir is the local path downloaded by the endpoint
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # IMPORTANT: trust_remote_code to load custom arch "gpt_oss"
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_dir, use_fast=True, trust_remote_code=True
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            model_dir,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
            trust_remote_code=True,
        )
        if not torch.cuda.is_available():
            # For safety on CPU, disable KV cache to reduce RAM spikes
            self.model.config.use_cache = False
        self.model.eval()

    def _prepare_inputs(self, texts: List[str]) -> Dict[str, Any]:
        prompts = [f"{PROMPT_PREFIX}{t}{PROMPT_SUFFIX}" for t in texts]
        toks = self.tokenizer(
            prompts, return_tensors="pt", padding=True, truncation=True
        )
        return {k: v.to(self.model.device) for k, v in toks.items()}

    @torch.inference_mode()
    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
        """
        Accepts { "inputs": "text" } or { "inputs": ["t1","t2",...] }
        Returns [{ "generated_text": "..." }, ...]
        """
        if "inputs" not in data:
            return [{"error": "missing 'inputs'"}]

        texts = _to_list(data["inputs"])
        inputs = self._prepare_inputs(texts)
        out = self.model.generate(**inputs, **GEN_KW)

        results: List[Dict[str, str]] = []
        # decode only the newly generated part
        for i, seq in enumerate(out):
            inp_len = inputs["input_ids"][i].shape[-1]
            gen_part = seq[inp_len:]
            text = self.tokenizer.decode(gen_part, skip_special_tokens=True).strip()
            results.append({"generated_text": text})
        return results