"""支援量化 draft 的載入器（原 local_load 只支援 bf16）。"""
import json
from pathlib import Path
import mlx.core as mx
import mlx.nn as nn
from dflash_port.model_mlx import DFlash2DraftModel, DFlashDraftModel, DFlashConfig
from dflash_port import local_load as _ll


def load_draft_any(path):
    path = Path(path)
    cfg = json.loads((path / "config.json").read_text())
    weights = {k: v for f in path.glob("*.safetensors") for k, v in mx.load(str(f)).items()}

    # codebook 鍵名在兩種發佈裡不一致，統一成 .weight 結尾
    for name in ("predecessor_codebook", "successor_codebook"):
        bare = f"candidate_selector.{name}"
        if bare in weights and f"{bare}.weight" not in weights:
            weights[f"{bare}.weight"] = weights.pop(bare)

    q = cfg.get("quantization") or {}
    if not q:                                   # 非量化 → 走原本的路徑
        return _ll.load_draft_from_dir(path)

    # 重用原載入器的 config 組裝（複製其邏輯的關鍵欄位）
    d = cfg.get("dflash_config", {})
    rope = cfg.get("rope_parameters") or cfg.get("rope_scaling")
    layer_types = tuple(cfg.get("layer_types") or ["full_attention"] * cfg["num_hidden_layers"])
    config = DFlashConfig(
        hidden_size=cfg["hidden_size"], num_hidden_layers=cfg["num_hidden_layers"],
        num_attention_heads=cfg["num_attention_heads"], num_key_value_heads=cfg["num_key_value_heads"],
        head_dim=cfg["head_dim"], intermediate_size=cfg["intermediate_size"],
        vocab_size=cfg["vocab_size"], rms_norm_eps=cfg["rms_norm_eps"],
        rope_theta=cfg.get("rope_theta", (rope or {}).get("rope_theta", 10000.0)),
        max_position_embeddings=cfg["max_position_embeddings"],
        block_size=int(d.get("block_size", 16)), target_layer_ids=tuple(d["target_layer_ids"]),
        num_target_layers=cfg["num_target_layers"], mask_token_id=d["mask_token_id"],
        rope_scaling=rope, layer_types=layer_types, sliding_window=cfg.get("sliding_window"),
        final_logit_softcapping=d.get("final_logit_softcapping", cfg.get("final_logit_softcapping")),
        input_embedding_scale=float(d.get("input_embedding_scale", 1.0)),
        output_multiplier=float(d.get("output_multiplier", 1.0)),
        conv_kernel_size=int(d.get("conv_kernel_size", 0)), conv_group_size=int(d.get("conv_group_size", 0)),
        selector_rank=int(d.get("selector_rank", 0)), selector_top_k=int(d.get("selector_top_k", 0)),
        is_causal=cfg.get("is_causal"),
    )
    klass = DFlash2DraftModel if "DFlash2DraftModel" in (cfg.get("architectures") or []) else DFlashDraftModel
    model = klass(config)

    # 只量化「權重檔裡真的有 scales」的模組 —— 與 mlx-lm 的判斷方式一致
    qset = {k[:-len(".scales")] for k in weights if k.endswith(".scales")}
    nn.quantize(model, group_size=int(q["group_size"]), bits=int(q["bits"]),
                class_predicate=lambda p, m: p in qset)
    model.eval()
    model.load_weights(list(weights.items()))
    mx.eval(model.parameters())
    return model