# -*- coding: utf-8 -*-
"""VERITAS 부착 핸들러 — JGOS-398B-fp8 (9-팀장 2026-09-17)

무엇을 하나
  한 번의 호출로 **답과 속 신호를 같이** 돌려준다.
  속 신호 = 마지막 층·마지막 토큰의 은닉벡터(가공 안 된 원본값).
  이것이 VERITAS 가 읽는 바로 그 값이며, **토큰을 한 개도 더 쓰지 않는다**
  (이미 계산된 것을 꺼내 올 뿐이다).

두 가지 모드
  mode="generate" : 답 생성 + 그 답까지 포함한 은닉 (문항 풀이용)
  mode="embed"    : 주어진 글자열의 은닉만 (이미 뽑아 둔 답을 채점할 때)

🔴 규율
  - 잘림(length)을 결과에 반드시 표시한다. 자원 소진을 오답으로 세지 않기 위해서다.
  - 은닉은 정규화하지 않는다. 정규화하면 VERITAS 가 쓰는 크기 정보가 날아간다.
  - 생성분이 reasoning 쪽으로 빠지는 템플릿이면 그것까지 합쳐 돌려준다(버리지 않는다).
"""
from __future__ import annotations

from typing import Any, Dict, List

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.tok = AutoTokenizer.from_pretrained(path)
        if self.tok.pad_token is None:
            self.tok.pad_token = self.tok.eos_token
        self.tok.padding_side = "left"          # 🔴 디코더 모델 배치의 필수 조건
        self.model = AutoModelForCausalLM.from_pretrained(
            path, dtype="auto", device_map="auto"
        ).eval()
        cfg = self.model.config
        cfg = cfg.get_text_config() if hasattr(cfg, "get_text_config") else cfg
        self.n_layers = cfg.num_hidden_layers

    # ---- 내부: 마지막 층·마지막 토큰 은닉 ------------------------------
    @torch.no_grad()
    def _hidden(self, ids) -> List[List[float]]:
        out = self.model(**ids, output_hidden_states=True, use_cache=False)
        h = out.hidden_states[-1]                       # 마지막 층
        mask = ids["attention_mask"]
        last = mask.sum(dim=1) - 1                      # 좌패딩이라 실제 마지막 위치
        picked = h[torch.arange(h.size(0), device=h.device), last, :]
        return picked.float().cpu().tolist()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inputs = data.get("inputs")
        if isinstance(inputs, str):
            inputs = [inputs]
        p = data.get("parameters") or {}
        mode = p.get("mode", "generate")

        # ── 은닉만 ──────────────────────────────────────────────────
        if mode == "embed":
            ids = self.tok(inputs, return_tensors="pt", padding=True,
                           truncation=True, max_length=p.get("max_length", 32768))
            ids = {k: v.to(self.model.device) for k, v in ids.items()}
            return {"hidden": self._hidden(ids), "layer": self.n_layers,
                    "normalized": False}

        # ── 생성 + 은닉 ─────────────────────────────────────────────
        max_new = int(p.get("max_new_tokens", 16384))
        n = int(p.get("n", 1))
        temp = float(p.get("temperature", 0.0))
        texts = []
        for s in inputs:
            if p.get("chat", True):
                msgs = [{"role": "user", "content": s}]
                kw = {}
                if p.get("enable_thinking") is not None:
                    kw["enable_thinking"] = p["enable_thinking"]
                try:
                    texts.append(self.tok.apply_chat_template(
                        msgs, tokenize=False, add_generation_prompt=True, **kw))
                except TypeError:
                    texts.append(self.tok.apply_chat_template(
                        msgs, tokenize=False, add_generation_prompt=True))
            else:
                texts.append(s)
        enc = self.tok(texts, return_tensors="pt", padding=True)
        enc = {k: v.to(self.model.device) for k, v in enc.items()}
        gen_kw = dict(max_new_tokens=max_new, pad_token_id=self.tok.pad_token_id)
        if temp > 0:
            gen_kw.update(do_sample=True, temperature=temp, num_return_sequences=n)
        else:
            gen_kw.update(do_sample=False)
        with torch.no_grad():
            g = self.model.generate(**enc, **gen_kw)
        plen = enc["input_ids"].shape[1]
        new = g[:, plen:]
        outs = self.tok.batch_decode(new, skip_special_tokens=True)
        ntok = (new != self.tok.pad_token_id).sum(dim=1).tolist()
        # 🔴 잘림 표시 — 자원 소진을 오답으로 세지 않기 위해 결과에 남긴다
        hit_cap = [int(t >= max_new - 2) for t in ntok]

        att = torch.ones_like(g)
        att[g == self.tok.pad_token_id] = 0
        hid = self._hidden({"input_ids": g, "attention_mask": att})
        return {"generated": outs, "hidden": hid, "n_tokens": ntok,
                "hit_cap": hit_cap, "layer": self.n_layers, "normalized": False}