"""Direct decision head on a small decoder (plan §6: 'same family with a direct decision head').

A causal LM (1-4B, LoRA-adapted) reads the serialized query + candidate list with symbols A, B, C, ...
and NONE. The decision is read from the next-token logits at the final position, restricted to the
symbol tokens. Single-pass; no autoregressive generation. Tokenization of every symbol is verified to be
exactly one token (plan §6: a decimal integer is not necessarily one token).

prompt_format "chat" (Nov 8, 2026; post hoc arm) wraps the same prompt in the backbone's chat template. For
hybrid-thinking models (Qwen3) the thinking-off template ends in an empty "<think>\n\n</think>\n\n" block, so
training and thinking-off scoring read the symbol right after it. Thinking-on scoring (decoder_trainer.
predict_decoder_thinking) lets the model write its own reasoning inside that block first, then reads the symbol
logits at the same position after "</think>\n\n".
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F

from ..data.serialization import NONE_SYMBOL, SYMBOLS, serialize_record


@dataclass
class DecoderHeadConfig:
    backbone: str = "Qwen/Qwen2.5-1.5B-Instruct"
    max_len: int = 2048
    serialization: str = "compact_v1"
    lora_r: int = 16
    lora_alpha: int = 32
    lora_dropout: float = 0.05
    lora_targets: tuple[str, ...] = ("q_proj", "k_proj", "v_proj", "o_proj")
    dtype: str = "bfloat16"
    symbol_prefix: str = " "  # token form tried first: " A" (leading space), then "A"
    gradient_checkpointing: bool = True
    prompt_format: str = "raw"  # raw: completion prompt ending in "Answer:" | chat: chat template (see module doc)
    max_think_tokens: int = 1024  # thinking-on scoring: reasoning budget before "</think>" is forced

    def to_dict(self) -> dict[str, Any]:
        d = self.__dict__.copy()
        d["lora_targets"] = list(self.lora_targets)
        return d


PROMPT_HEADER = (
    "You are matching a query record to at most one catalog record.\n"
    "Identity policy: {policy}\n"
    "Query: {query}\n"
    "Candidates:\n{candidates}\n"
    "Answer with the single symbol of the matching candidate, or {none} if no candidate is the same entity.\n"
    "Answer:"
)


def render_prompt(query_fields: dict[str, Any], candidates: list[dict[str, Any]], policy: str, serialization: str, field_order: list[str] | None = None, compact_fields: list[str] | None = None) -> str:
    q = serialize_record(query_fields, field_order, serialization, compact_fields)
    lines = [f"{SYMBOLS[j]}. {serialize_record(c, field_order, serialization, compact_fields)}" for j, c in enumerate(candidates)]
    return PROMPT_HEADER.format(policy=policy, query=q, candidates="\n".join(lines), none=NONE_SYMBOL)


def verify_symbol_tokens(tokenizer, k_max: int, prefix_candidates: tuple[str, ...] = (" ", "")) -> tuple[str, list[int], int]:
    """Return (prefix, candidate_token_ids[k_max], none_token_id) such that every symbol is ONE token.

    Raises if no prefix variant makes all symbols single tokens.
    """
    problems = {}
    for prefix in prefix_candidates:
        ids = []
        ok = True
        for s in SYMBOLS[:k_max] + [NONE_SYMBOL]:
            toks = tokenizer.encode(prefix + s, add_special_tokens=False)
            if len(toks) != 1:
                ok = False
                problems[prefix + s] = toks
                break
            ids.append(toks[0])
        if ok and len(set(ids)) == len(ids):
            return prefix, ids[:-1], ids[-1]
    raise ValueError(f"no single-token symbol encoding for K={k_max}; offending: {problems}. Reduce K or change symbols.")


class DecoderDecisionHead(nn.Module):
    def __init__(self, cfg: DecoderHeadConfig, k_max: int, use_lora: bool = True):
        super().__init__()
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.cfg = cfg
        self.tokenizer = AutoTokenizer.from_pretrained(cfg.backbone)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = "left"  # last position = decision position
        dtype = getattr(torch, cfg.dtype)
        self.lm = AutoModelForCausalLM.from_pretrained(cfg.backbone, torch_dtype=dtype, attn_implementation="sdpa")
        if cfg.gradient_checkpointing:
            self.lm.gradient_checkpointing_enable()
            self.lm.config.use_cache = False
        if use_lora:
            from peft import LoraConfig, get_peft_model

            self.lm = get_peft_model(self.lm, LoraConfig(r=cfg.lora_r, lora_alpha=cfg.lora_alpha, lora_dropout=cfg.lora_dropout, target_modules=list(cfg.lora_targets), task_type="CAUSAL_LM"))
        self.prefix, cand_ids, none_id = verify_symbol_tokens(self.tokenizer, k_max, (cfg.symbol_prefix, "" if cfg.symbol_prefix else " "))
        self.register_buffer("symbol_ids", torch.tensor(cand_ids + [none_id], dtype=torch.long), persistent=False)
        self.k_max = k_max

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, slot_mask: torch.Tensor) -> torch.Tensor:
        """Return log-probs [B, k_max+1] over symbols (padded slots -inf). Left-padded inputs."""
        # logits_to_keep=1: only the decision position's vocabulary logits (same values; avoids [B, L, V] at long L)
        out = self.lm(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=1)
        last = out.logits[:, -1, :].float()  # [B, V]
        sym = last[:, self.symbol_ids]  # [B, k_max+1]
        full_mask = torch.cat([slot_mask, torch.ones(slot_mask.shape[0], 1, dtype=torch.bool, device=slot_mask.device)], 1)
        sym = sym.masked_fill(~full_mask, float("-inf"))
        return F.log_softmax(sym, dim=-1)

    def format_prompts(self, prompts: list[str], thinking: bool = False) -> list[str]:
        """Model input text for rendered prompts: unchanged for "raw"; chat template for "chat"."""
        if self.cfg.prompt_format == "raw":
            if thinking:
                raise ValueError("thinking needs prompt_format: chat")
            return prompts
        if self.cfg.prompt_format != "chat":
            raise ValueError(self.cfg.prompt_format)
        suffix = "\nAnswer:"
        return [self.tokenizer.apply_chat_template([{"role": "user", "content": p.removesuffix(suffix)}], tokenize=False,
                                                   add_generation_prompt=True, enable_thinking=thinking) for p in prompts]

    def tokenize_prompts(self, prompts: list[str], formatted: bool = False, max_len: int | None = None) -> dict[str, torch.Tensor]:
        texts = prompts if formatted else self.format_prompts(prompts)
        enc = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=False, add_special_tokens=self.cfg.prompt_format == "raw")
        limit = max_len or self.cfg.max_len
        if enc["input_ids"].shape[1] > limit:
            # right truncation would cut the decision position (the last token); never truncate silently
            raise ValueError(f"prompt of {enc['input_ids'].shape[1]} tokens exceeds max_len {limit}")
        return {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]}

    def save_adapter(self, path: str) -> None:
        self.lm.save_pretrained(path)
        self.tokenizer.save_pretrained(path)