"""量化变体加载器（laya-ad 发布格式）。

支持 laya-ad Release 的五种 safetensors 变体并统一加载为可推理模型：
- fp16 / bf16：transformers 原生加载
- fp8（float8_e4m3fn 存储）：cast 回 bf16
- int8（per-channel 对称 weight-only，`__scale__` 伴生 fp16）
- int4（group=64 对称、双 nibble 打包 uint8、`__scale__`/`__pad__` 伴生）

由 scripts/quantize_variants.py 产出；格式与 laya-ad 仓库 scripts/load_quantized.py 一致。
"""
import json
from pathlib import Path
from types import SimpleNamespace

import torch
from safetensors import safe_open

INT4_GROUP = 64


class _DualHeadBinary(torch.nn.Module):
    """把多任务模型包装成服务端可用的模型：binary/kind/scope/severity 四组 logits。"""

    def __init__(self, model, kind_labels: list[str] | None = None,
                 scope_labels: list[str] | None = None,
                 severity_labels: list[str] | None = None):
        super().__init__()
        self.model = model
        self.config = model.config
        self.kind_labels = list(kind_labels or [])
        self.scope_labels = list(scope_labels or [])
        self.severity_labels = list(severity_labels or [])
        self.has_scope = bool(self.scope_labels) and getattr(model, "scope_head", None) is not None
        self.has_severity = (bool(self.severity_labels)
                             and getattr(model, "severity_head", None) is not None)

    def forward(self, input_ids=None, attention_mask=None, **kwargs):
        binary_logits, kind_logits, scope_logits, severity_logits, _ = self.model.forward_all(
            input_ids, attention_mask)
        return SimpleNamespace(logits=binary_logits, kind_logits=kind_logits,
                               scope_logits=scope_logits, severity_logits=severity_logits)


def detect_precision(model_dir: Path) -> str:
    """优先读 PRECISION.txt（CI 注入），否则从权重 dtype 推断。"""
    marker = Path(model_dir) / "PRECISION.txt"
    if marker.is_file():
        return marker.read_text(encoding="utf-8").strip() or "unknown"
    sd = Path(model_dir) / "model.safetensors"
    dts = set()
    with safe_open(str(sd), framework="pt") as f:
        for k in f.keys():
            dts.add(f.get_slice(k).get_dtype())
    if "F8_E4M3" in dts:
        return "fp8"
    if "U8" in dts:
        return "int4"
    if "I8" in dts:
        return "int8"
    if "BF16" in dts:
        return "bf16"
    if "F16" in dts:
        return "fp16"
    return "fp32"


def load_ad_model(model_dir, device: str = "cpu", dtype: torch.dtype = torch.bfloat16):
    """加载广告分类模型（自动识别量化格式），返回 (model, tokenizer)，eval 模式。"""
    from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer

    model_dir = Path(model_dir)
    if (model_dir / "dual_config.json").is_file():
        # v2 双头模型：共享 encoder + binary/kind 两个头
        from transformers import AutoTokenizer

        from .dual_head_model import load_dual_head_any

        dual_cfg = json.loads((model_dir / "dual_config.json").read_text(encoding="utf-8"))
        kind_labels = [str(x) for x in dual_cfg.get("kind_labels") or []]
        scope_labels = [str(x) for x in dual_cfg.get("scope_labels") or []]
        severity_labels = [str(x) for x in dual_cfg.get("severity_labels") or []]
        dual = load_dual_head_any(str(model_dir), device=device, dtype=dtype)
        return (_DualHeadBinary(dual, kind_labels, scope_labels, severity_labels).to(device).eval(),
                AutoTokenizer.from_pretrained(str(model_dir)))

    sd_path = model_dir / "model.safetensors"
    dts = set()
    with safe_open(str(sd_path), framework="pt") as f:
        for k in f.keys():
            dts.add(f.get_slice(k).get_dtype())

    if not ({"I8", "U8", "F8_E4M3"} & dts):
        kw = {"dtype": dtype}
        model = AutoModelForSequenceClassification.from_pretrained(
            str(model_dir), attn_implementation="sdpa", **kw)
        tok = AutoTokenizer.from_pretrained(str(model_dir))
        return model.to(device).eval(), tok

    model = AutoModelForSequenceClassification.from_config(
        AutoConfig.from_pretrained(str(model_dir)))
    param_shapes = {k: v.shape for k, v in model.state_dict().items()}
    with safe_open(str(sd_path), framework="pt") as f:
        raw = {k: f.get_tensor(k) for k in f.keys()}

    state = {}
    for k, v in raw.items():
        if k.endswith(".__scale__") or k.endswith(".__pad__"):
            continue
        if v.dtype == torch.float8_e4m3fn:
            state[k] = v.to(dtype)
        elif v.dtype == torch.int8 and k + ".__scale__" in raw:
            scale = raw[k + ".__scale__"].float().view(-1, *([1] * (v.dim() - 1)))
            state[k] = (v.float() * scale).to(dtype)
        elif v.dtype == torch.uint8 and k + ".__scale__" in raw:
            hi = (v >> 4).to(torch.int16) - 8
            lo = (v & 0xF).to(torch.int16) - 8
            q = torch.stack([hi, lo], dim=-1).reshape(-1).float()
            scale = raw[k + ".__scale__"].float()
            s_exp = scale.repeat_interleave(INT4_GROUP)[:q.numel()]
            w = q * s_exp
            pad = int(raw[k + ".__pad__"].item()) if k + ".__pad__" in raw else 0
            if pad:
                w = w[:-pad]
            state[k] = w.reshape(param_shapes[k]).to(dtype)
        else:  # 未量化透传（1D bias/norm 等保留精度）
            state[k] = v.to(dtype)

    model.load_state_dict(state, strict=True)
    tok = AutoTokenizer.from_pretrained(str(model_dir))
    return model.to(device).eval(), tok