"""System One 결정 모델: Qwen3.5 텍스트 모델(LoRA 병합) + 결정 헤드 (Hugging Face `trust_remote_code`).

상태(state)와 타입이 정해진 질문(noul·choice·score)을 받아, 글을 생성하지 않고 한 번의 forward로
질문마다 선택지별 보정된 확률을 낸다.

    from transformers import AutoModel, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
    model = AutoModel.from_pretrained(REPO, trust_remote_code=True).eval()
    model.decide(
        tok,
        state="고객: 결제가 두 번 됐어요. 한 건은 취소해 주세요.",
        questions={
            "refund": {"type": "noul", "instructions": "환불·취소 요청인가요?"},
            "team": {"type": "choice", "instructions": "담당 팀을 고르세요.",
                     "options": ["billing", "shipping", "account"]},
        },
    )

GPU에서 한 건씩 판단할 때는 `model.to("cuda").enable_cuda_graphs()`를 한 번 부르면 decide()가 길이 구간별
CUDA 그래프로 돈다 (A100에서 약 3배 빠르다. 아래 `_GraphRunner`).

입력 형식 v2(config.input_format)부터는 같은 forward에서 언어 모델 글자 경로도 읽는다: 선택지에 글자(A., B., …)를
붙이고 질문마다 답 자리를 두어, 그 자리의 다음 토큰 확률(출력층 = 입력 임베딩)을 결정 헤드의 확률과 섞는다.

이 파일은 저장소에 같이 올라가 받는 쪽에서 실행된다. 학습 코드(sysone)에 기대지 않고, 입력 형식은 학습 때와
같은 규칙을 따른다 (상태 → 질문마다 머리말·지시문·선택지 줄, 선택지 줄의 마지막 토큰 위치로 점수, 상태가 길면 뒤를 자른다).
"""

from __future__ import annotations

import contextlib
import json
import threading
from typing import Any, ClassVar

import torch
from torch import nn
from transformers import AutoConfig, AutoModel, PreTrainedModel

from .configuration_sysone import SysOneConfig

QTYPES = ("noul", "choice", "score")
NOUL_OPTIONS = ("true", "false")
STATE_HEADER = "### 상태 (state)\n"
MIN_STATE_TOKENS = 32
TYPE_HINT = {
    "noul": "예/아니오 (true/false)",
    "choice": "하나 고르기",
    "score": "순서 척도 (낮음→높음)",
}
LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
NOUL_SHOWN = {"true": "true (예 / yes)", "false": "false (아니오 / no)"}
LETTER_HINT = (
    "Answer with one letter."  # 질문마다 따로 묻는 프롬프트(sysone.baseline)의 마지막 줄과 같다
)
ANSWER_CUE = LETTER_HINT + "\n답:"
ALL_HINT = "Answer each question with one letter."
CHAT_SYSTEM = (
    "You are a careful decision model. Read the state and answer each multiple-choice question "
    "with the letter of exactly one option."
)
CHAT_OPEN = f"<|im_start|>system\n{CHAT_SYSTEM}<|im_end|>\n<|im_start|>user\n"
CHAT_CLOSE = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
CHAT_CLOSE_NOTHINK = "<|im_end|>\n<|im_start|>assistant\n"
TRUNC_MARK = "\n(…중략…)\n"
REOPEN = "<|im_end|>\n<|im_start|>user\n"  # v4chat: 빈 어시스턴트 차례를 닫고 다음 질문의 사용자 차례를 연다
TAIL = -2
AFTER = -3


class DecisionHead(nn.Module):
    """선택지 위치 hidden을 쿼리로: 선택지끼리 self-attention + 전체 시퀀스 cross-attention + 선형 prior."""

    def __init__(
        self,
        d_model: int,
        d_head: int = 512,
        n_layers: int = 2,
        n_heads: int = 8,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.norm_in = nn.LayerNorm(d_model)
        self.mem_proj = nn.Linear(d_model, d_head)
        self.q_proj = nn.Linear(d_model, d_head)
        self.type_emb = nn.Embedding(len(QTYPES), d_head)
        layer = nn.TransformerDecoderLayer(
            d_head, n_heads, 4 * d_head, dropout, batch_first=True, norm_first=True
        )
        self.layers = nn.TransformerDecoder(layer, n_layers)
        self.norm_out = nn.LayerNorm(d_head)
        self.out = nn.Linear(d_head, 1)
        self.prior = nn.Linear(d_model, 1)

    def forward(self, hidden, attention_mask, opt_pos, opt_mask, opt_type):
        h = self.norm_in(hidden.float())
        opt_h = h.gather(1, opt_pos[..., None].expand(-1, -1, h.size(-1)))
        q = self.q_proj(opt_h) + self.type_emb(opt_type)
        x = self.layers(
            q,
            self.mem_proj(h),
            tgt_key_padding_mask=~opt_mask,
            memory_key_padding_mask=~attention_mask,
        )
        return (self.out(self.norm_out(x)) + self.prior(opt_h)).squeeze(-1)


def parse_questions(questions: dict[str, Any]) -> list[dict[str, Any]]:
    out = []
    for name, q in questions.items():
        qtype = q.get("type")
        if qtype not in QTYPES:
            raise ValueError(f"질문 {name}: type은 {QTYPES} 중 하나 (받은 값: {qtype})")
        if qtype == "noul":
            options = list(NOUL_OPTIONS)
        else:
            options = [str(o) for o in q.get("options") or []]
            if len(options) < 2 or len(set(options)) != len(options):
                raise ValueError(f"질문 {name}: 서로 다른 선택지가 2개 이상 필요하다")
        out.append({"name": name, "type": qtype, "instructions": q.get("instructions", ""),
                    "criteria": q.get("criteria", ""), "options": options})  # fmt: skip
    if not out:
        raise ValueError("질문이 없다")
    return out


def _ids(tok, text: str) -> list[int]:
    return tok(text, add_special_tokens=False)["input_ids"]


def _cut_state(tok, s_all: list[int], budget: int, fmt: str) -> list[int]:
    """v1은 상태 뒤를 자르고, v2는 앞 1/4과 뒤 3/4를 남겨 가운데에 줄임 표시를 넣는다 (최신 발화·글의 끝을 지킨다)."""
    if len(s_all) <= budget:
        return s_all
    if fmt == "v1":
        return s_all[:budget]
    mark = _ids(tok, TRUNC_MARK)
    keep = max(0, budget - len(mark))
    head = keep // 4
    return s_all[:head] + mark + s_all[len(s_all) - (keep - head) :]


def encode(
    tok, state: Any, questions: dict[str, Any], max_len: int, fmt: str = "v1"
) -> dict[str, Any]:
    """[상태 머리말 + 상태] [질문 1 머리말·지시문·기준·선택지 줄 (+ v2 답 자리)] [질문 2 ...] → 토큰 id와 위치.

    v2: 선택지 줄 `A. 선택지`, 선택지 뒤 `답:` (선택지 26개 이하일 때). v2chat: 채팅 형식으로 감싸고 마지막 질문은
    어시스턴트 차례 시작에서 답한다. 학습 코드(sysone.encode)와 같은 규칙이다.
    """
    qs = parse_questions(questions)
    v2, after = fmt != "v1", fmt == "v3chat"
    chat = fmt in ("v2chat", "v3chat", "v4chat", "v4chat-nothink")
    turns, in_answer = fmt in ("v4chat", "v4chat-nothink"), False
    close = CHAT_CLOSE_NOTHINK if fmt == "v4chat-nothink" else CHAT_CLOSE
    prefix = [t for t in (tok.cls_token_id, tok.bos_token_id) if t is not None][:1]
    q_ids: list[int] = []
    opt_rel, opt_q, opt_type, slot_rel, n_letters = [], [], [], [], []
    for qi, q in enumerate(qs):
        head = f"\n### 질문 {qi + 1}: {q['name']} [{TYPE_HINT[q['type']]}]\n{q['instructions']}\n"
        if q["criteria"]:
            head += f"기준: {q['criteria']}\n"
        if in_answer:
            q_ids += _ids(tok, REOPEN)
            in_answer = False
        q_ids += _ids(tok, head)
        lm = v2 and len(q["options"]) <= len(LETTERS)
        for j, opt in enumerate(q["options"]):
            shown = NOUL_SHOWN[opt] if v2 and q["type"] == "noul" else opt
            q_ids += _ids(tok, f"{LETTERS[j]}. {shown}\n" if lm else f"- {shown}\n")
            opt_rel.append(len(q_ids) - 1)
            opt_q.append(qi)
            opt_type.append(QTYPES.index(q["type"]))
        if lm and turns:
            q_ids += _ids(tok, LETTER_HINT + close)
            slot_rel.append(len(q_ids) - 1)
            in_answer = True
        elif lm and after:
            slot_rel.append(AFTER)
        elif lm and chat and qi == len(qs) - 1:
            q_ids += _ids(tok, LETTER_HINT)
            slot_rel.append(TAIL)
        elif lm:
            q_ids += _ids(tok, ANSWER_CUE)
            slot_rel.append(len(q_ids) - 1)
        else:
            slot_rel.append(-1)
        n_letters.append(len(q["options"]) if lm else 0)
    a_ids: list[int] = []
    lm_qs = [qi for qi, p in enumerate(slot_rel) if p == AFTER]
    if lm_qs:
        q_ids += _ids(tok, LETTER_HINT if len(lm_qs) == 1 else ALL_HINT)
        if len(lm_qs) == 1:
            slot_rel[lm_qs[0]] = TAIL
        else:
            for i, qi in enumerate(lm_qs):
                a_ids += _ids(tok, ("\n" if i else "") + f"질문 {qi + 1}:")
                slot_rel[qi] = AFTER - len(a_ids)
    open_ids = _ids(tok, CHAT_OPEN) if chat else []
    close_ids = _ids(tok, CHAT_CLOSE) if chat and not turns else []
    header = _ids(tok, STATE_HEADER)
    budget = (max_len - len(prefix) - len(open_ids) - len(header) - len(q_ids) - len(close_ids)
              - len(a_ids))  # fmt: skip
    if budget < MIN_STATE_TOKENS:
        raise ValueError(
            f"질문·선택지만 {len(q_ids)} 토큰이라 max_len={max_len}에 상태가 들어갈 자리가 없다"
        )
    text = state if isinstance(state, str) else json.dumps(state, ensure_ascii=False, indent=1)
    s_all = _ids(tok, text)
    s_ids = _cut_state(tok, s_all, budget, fmt)
    base = len(prefix) + len(open_ids) + len(header) + len(s_ids)
    input_ids = prefix + open_ids + header + s_ids + q_ids + close_ids + a_ids
    turn_end = base + len(q_ids) + len(close_ids)

    def slot(p: int) -> int:
        if p >= 0:
            return base + p
        if p == TAIL:
            return turn_end - 1
        return turn_end + (AFTER - p) - 1 if p < AFTER else -1

    slot_pos = [slot(p) for p in slot_rel]
    return {"input_ids": input_ids, "opt_pos": [base + p for p in opt_rel], "opt_q": opt_q,
            "opt_type": opt_type, "slot_pos": slot_pos, "n_letters": n_letters, "questions": qs,
            "truncated": len(s_all) > budget}  # fmt: skip


def letter_logits(hidden, slot_pos, emb, letter_ids):
    """답 자리 hidden (B, Q, d) × 출력층(= 입력 임베딩)의 글자 행 → 글자 logit (B, Q, 26). 변형끼리 logsumexp."""
    d = hidden.size(-1)
    h = hidden.gather(1, slot_pos.clamp(min=0)[..., None].expand(-1, -1, d)).float()
    w = emb[letter_ids.reshape(-1)].float().view(*letter_ids.shape, d)
    return torch.einsum("bqd,kld->bqkl", h, w).logsumexp(2)


def temperature_bucket(qtype: str, n_opts: int, head_only: bool = False) -> str:
    """온도 구간. 형식 v2에서 글자 경로가 없는 질문(선택지 27개 이상)은 따로 쓴다."""
    if head_only:
        return f"{qtype}:27+"
    return f"{qtype}:{n_opts if n_opts <= 2 else '3-5' if n_opts <= 5 else '6+'}"


def graph_buckets(max_len: int) -> list[int]:
    """CUDA 그래프 길이 구간: 512까지 32씩, 1024까지 64씩, 그 뒤는 128씩 max_len을 덮을 때까지.

    입력을 구간 길이까지 채워 계산하므로 구간이 촘촘할수록 덜 낭비한다 (형식 v4chat에서 221토큰이 256이 아니라 224로).
    """
    out, n = [], 0
    while n < max_len:
        n += 32 if n < 512 else 64 if n < 1024 else 128
        out.append(n)
    return out


class _GraphRunner:
    """배치 1 백본 forward를 길이 구간마다 CUDA 그래프로 한 번 녹화해 두고 다시 재생한다 (`enable_cuda_graphs`).

    - eager 실행은 층마다 작은 커널을 하나씩 부른다 (4B는 forward 한 번에 2천여 개). 짧은 입력에서는 GPU 계산보다
      이 호출 비용이 커서 길이와 상관없이 지연이 거의 같았다. 재생은 호출 비용을 없앤다
    - 입력은 구간 길이까지 오른쪽을 채운다. 인과 모델이라 채운 토큰은 앞쪽 hidden을 바꾸지 않는다 (수치 오차만 다르다)
    - 그래프들은 메모리 풀 하나를 같이 쓴다. 출력 버퍼도 공유라 재생부터 헤드 계산까지 잠금 안에서 한다
    """

    def __init__(self, backbone: nn.Module, max_len: int):
        self.backbone = backbone
        self.device = next(backbone.parameters()).device
        self.buckets = graph_buckets(max_len)
        self.pool = torch.cuda.graph_pool_handle()
        self.graphs: dict[int, tuple[Any, torch.Tensor, torch.Tensor]] = {}
        self.lock = threading.Lock()

    def capture(self, length: int) -> None:
        static = torch.zeros(1, length, dtype=torch.long, device=self.device)
        stream = torch.cuda.Stream(self.device)
        stream.wait_stream(torch.cuda.current_stream(self.device))
        with torch.no_grad(), torch.cuda.stream(stream):
            for _ in range(2):  # Triton 커널 자동 조정과 cuBLAS 준비를 녹화 전에 끝낸다
                self.backbone(input_ids=static, use_cache=False)
        torch.cuda.current_stream(self.device).wait_stream(stream)
        graph = torch.cuda.CUDAGraph()
        with torch.no_grad(), torch.cuda.graph(graph, pool=self.pool):
            out = self.backbone(input_ids=static, use_cache=False).last_hidden_state
        self.graphs[length] = (graph, static, out)

    def __call__(self, input_ids: torch.Tensor) -> torch.Tensor | None:
        """(1, n) → (1, n, d) hidden. 가장 긴 구간보다 길면 None이다 (호출한 쪽이 eager로 계산한다)."""
        n = input_ids.size(1)
        length = next((b for b in self.buckets if b >= n), None)
        if input_ids.size(0) != 1 or length is None:
            return None
        if length not in self.graphs:
            self.capture(length)
        graph, static, out = self.graphs[length]
        static[:, :n].copy_(input_ids)
        static[:, n:].zero_()
        graph.replay()
        return out[:, :n]


class SysOneForDecision(PreTrainedModel):
    config_class = SysOneConfig
    base_model_prefix = "sysone"
    # 헤드는 float32로 학습했다. _keep_in_fp32_modules는 fp16에만 적용되어 bf16까지 막으려면 _strict가 필요하다
    _keep_in_fp32_modules: ClassVar[list[str]] = ["head"]
    _keep_in_fp32_modules_strict: ClassVar[list[str]] = ["head"]
    _supports_sdpa = True

    def __init__(self, config: SysOneConfig):
        super().__init__(config)
        text = dict(config.text_config)
        self.backbone = AutoModel.from_config(AutoConfig.for_model(text.pop("model_type"), **text))
        self.head = DecisionHead(self.backbone.config.hidden_size, **config.head)
        self._graphs: _GraphRunner | None = None
        self._letter_ids: torch.Tensor | None = None
        self.post_init()

    def forward(self, input_ids, attention_mask, opt_pos, opt_mask, opt_type):
        hidden = self.backbone(
            input_ids=input_ids, attention_mask=attention_mask.long(), use_cache=False
        ).last_hidden_state
        return self.head(hidden, attention_mask, opt_pos, opt_mask, opt_type)

    def enable_cuda_graphs(self, capture_all: bool = True) -> SysOneForDecision:
        """decide()의 백본을 길이 구간별 CUDA 그래프로 돌린다 (배치 1).

        GPU로 옮긴 뒤 부른다. capture_all이면 모든 구간을 지금 녹화하고 (A100에서 약 5초, 추가 메모리 1GiB 안쪽),
        아니면 처음 쓰는 구간을 그때 녹화한다. 장치·dtype·가중치를 바꾸면 다시 불러야 한다.
        """
        device = next(self.backbone.parameters()).device
        if device.type != "cuda":
            raise RuntimeError(
                "CUDA 그래프는 CUDA 장치에서만 쓸 수 있다. model.to('cuda') 뒤에 부를 것"
            )
        runner = _GraphRunner(self.backbone, self.config.max_len)
        if capture_all:
            for length in runner.buckets:
                runner.capture(length)
        self._graphs = runner
        return self

    def disable_cuda_graphs(self) -> None:
        self._graphs = None

    @torch.no_grad()
    def decide(
        self, tokenizer, state: Any, questions: dict[str, Any], max_len: int | None = None
    ) -> dict[str, Any]:
        """질문마다 선택지별 확률(보정 온도 적용).

        choice·score: answer(가장 높은 선택지)·confidence, score는 expected(선택지 순번 1부터의 기댓값)도 준다.
        noul: answer(bool)·p_true. truncated는 상태가 max_len을 넘어 잘렸는지다.
        입력 형식 v2부터는 글자 경로가 있는 질문에 paths(결정 헤드·언어 모델 글자의 경로별 확률, 보정 전)도 준다.
        """
        if (
            self.head.out.weight.dtype != torch.float32
        ):  # 다른 dtype으로 옮겨졌어도 헤드는 float32로 계산한다
            self.head.float()
        fmt = getattr(self.config, "input_format", "v1")
        enc = encode(tokenizer, state, questions, max_len or self.config.max_len, fmt)
        dev = self.head.out.weight.device
        ids = torch.tensor([enc["input_ids"]], device=dev)
        n = len(enc["opt_pos"])
        mask = torch.ones_like(ids, dtype=torch.bool)
        pos = torch.tensor([enc["opt_pos"]], device=dev)
        opt_mask = torch.ones(1, n, dtype=torch.bool, device=dev)
        opt_type = torch.tensor([enc["opt_type"]], device=dev)
        use_lm = fmt != "v1" and any(enc["n_letters"])
        runner = self._graphs if self._graphs is not None and dev == self._graphs.device else None
        letters = None
        with runner.lock if runner is not None else contextlib.nullcontext():
            hidden = runner(ids) if runner is not None else None
            if hidden is None:
                hidden = self.backbone(
                    input_ids=ids, attention_mask=mask.long(), use_cache=False
                ).last_hidden_state
            logits = self.head(hidden, mask, pos, opt_mask, opt_type)[0].float().cpu()
            if use_lm:
                if self._letter_ids is None or self._letter_ids.device != dev:
                    self._letter_ids = torch.tensor(self.config.letter_ids, device=dev)
                slots = torch.tensor([enc["slot_pos"]], device=dev)
                emb = self.backbone.get_input_embeddings().weight
                letters = letter_logits(hidden, slots, emb, self._letter_ids)[0].float().cpu()
        temps = self.config.temperatures or {}
        w = float(getattr(self.config, "lm_weight", 0.5))
        answers = {}
        for qi, q in enumerate(enc["questions"]):
            cols = [j for j, x in enumerate(enc["opt_q"]) if x == qi]
            n_let = enc["n_letters"][qi] if fmt != "v1" else None
            final, paths = logits[cols], None
            if n_let and letters is not None:
                lm = letters[qi, :n_let]
                paths = {"head": final.softmax(-1).tolist(), "lm": lm.softmax(-1).tolist()}
                final = (1 - w) * final.log_softmax(-1) + w * lm.log_softmax(-1)
            t = temps.get(temperature_bucket(q["type"], len(cols), n_let == 0), temps.get("*", 1.0))
            probs = (final / t).softmax(-1).tolist()
            out: dict[str, Any] = {"probs": dict(zip(q["options"], (round(p, 4) for p in probs)))}
            if paths:
                out["paths"] = {k: dict(zip(q["options"], (round(p, 4) for p in v)))
                                for k, v in paths.items()}  # fmt: skip
            if q["type"] == "noul":
                out |= {"answer": probs[0] >= 0.5, "p_true": round(probs[0], 4)}
            else:
                best = max(range(len(probs)), key=probs.__getitem__)
                out |= {"answer": q["options"][best], "confidence": round(probs[best], 4)}
                if q["type"] == "score":
                    out["expected"] = round(sum((k + 1) * p for k, p in enumerate(probs)), 3)
            answers[q["name"]] = out
        return {"answers": answers, "truncated": enc["truncated"]}

    def decide_record(
        self, tokenizer, record: dict[str, Any], max_len: int | None = None
    ) -> dict[str, Any]:
        """{"id", "state", "questions"} 레코드 하나 (학습 데이터와 같은 형식)."""
        return {
            "id": record.get("id"),
            **self.decide(tokenizer, record["state"], record["questions"], max_len),
        }