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

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

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

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")),
}

HF_MODEL_ID = os.getenv("HF_MODEL_ID", "ZennyKenny/oss-20b-prereform-to-modern-ru-merged")

def _is_model_dir(path: str) -> bool:
    cfg = os.path.join(path, "config.json")
    if not os.path.exists(cfg):
        return False
    try:
        with open(cfg, "r", encoding="utf-8") as f:
            j = json.load(f)
        return ("model_type" in j) or ("auto_map" in j)
    except Exception:
        return False

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

class EndpointHandler:
    def __init__(self, model_dir: str):
        # Log what paths we got (helps debugging in HF logs)
        print(f"[handler] model_dir passed by toolkit: {model_dir}")
        print(f"[handler] repo path: /repository")

        source = None
        if _is_model_dir(model_dir):
            source = model_dir
            print(f"[handler] using model_dir as source")
        else:
            source = HF_MODEL_ID
            print(f"[handler] model_dir not valid, falling back to HF_MODEL_ID={source}")

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.float16 if torch.cuda.is_available() else torch.float32

        # gpt_oss requires remote code
        self.tokenizer = AutoTokenizer.from_pretrained(
            source, use_fast=True, trust_remote_code=True
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            source,
            torch_dtype=dtype,
            device_map="auto" if torch.cuda.is_available() else None,
            trust_remote_code=True,
        )
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        if not torch.cuda.is_available():
            self.model.config.use_cache = False
        self.model.eval()
        print(f"[handler] model loaded on device: {self.model.device}")

    def _encode(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]]:
        # Accept { "inputs": "…" } or { "inputs": ["…", "…"] }
        if "inputs" not in data:
            return [{"error": "missing 'inputs'"}]

        texts = _as_list(data["inputs"])
        inputs = self._encode(texts)
        outputs = self.model.generate(**inputs, **GEN_KW)

        results: List[Dict[str, str]] = []
        for i, seq in enumerate(outputs):
            in_len = inputs["input_ids"][i].shape[-1]
            gen_only = seq[in_len:]
            text = self.tokenizer.decode(gen_only, skip_special_tokens=True).strip()
            results.append({"generated_text": text})
        return results