from __future__ import annotations

import os
from typing import Any, Dict, List, Tuple

import torch
import torch.nn as nn

from .config import QuantConfig
from .qmodules import QuantConv1d, QuantConv2d, QuantLinear


def _unwrap_state_dict(obj: Any) -> Dict[str, torch.Tensor]:
    if isinstance(obj, dict):
        for key in ("state_dict", "model", "net", "params"):
            value = obj.get(key)
            if isinstance(value, dict):
                return value
        return obj
    raise ValueError("Unsupported checkpoint format")


def load_state_dict_flexible(
    model: nn.Module,
    ckpt_path: str,
    device: torch.device,
) -> Tuple[List[str], List[str]]:
    ckpt = torch.load(ckpt_path, map_location="cpu")
    sd = _unwrap_state_dict(ckpt)

    model_sd = model.state_dict()
    if any(k.startswith("module.") for k in sd) and not any(k.startswith("module.") for k in model_sd):
        sd = {k[len("module."):]: v for k, v in sd.items()}

    missing, unexpected = model.load_state_dict(sd, strict=False)
    model.to(device)
    return list(missing), list(unexpected)


def save_quant_checkpoint(
    path: str,
    q_model: nn.Module,
    qcfg: QuantConfig,
    extra: Dict[str, Any] | None = None,
) -> None:
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    torch.save(
        {
            "state_dict": q_model.state_dict(),
            "quant_config": qcfg.to_dict(),
            "extra": extra or {},
        },
        path,
    )


def _get_parent(root: nn.Module, full_name: str) -> tuple[nn.Module, str]:
    parts = full_name.split(".")
    if len(parts) == 1:
        return root, parts[0]
    parent = root
    for part in parts[:-1]:
        parent = getattr(parent, part)
    return parent, parts[-1]


def _has_quant_weight(state_dict: Dict[str, torch.Tensor], module_name: str) -> bool:
    prefix = module_name + "."
    return any(k.startswith(prefix + "w_") for k in state_dict)


def _infer_weight_bits(state_dict: Dict[str, torch.Tensor], module_name: str, fallback: int) -> int:
    prefix = module_name + "."
    if (prefix + "w_q_packed") in state_dict:
        return 4
    if (prefix + "w_q_int8") in state_dict:
        return int(fallback)
    return int(fallback)


@torch.no_grad()
def rebuild_quant_model_from_checkpoint(
    fp_model: nn.Module,
    quant_ckpt_path: str,
    device: torch.device,
) -> nn.Module:
    payload = torch.load(quant_ckpt_path, map_location="cpu")
    if not (isinstance(payload, dict) and "quant_config" in payload and "state_dict" in payload):
        raise ValueError("Invalid quant checkpoint format; expected quant_config and state_dict")

    qcfg = QuantConfig.from_dict(payload["quant_config"])
    sd: Dict[str, torch.Tensor] = payload["state_dict"]

    for full_name, module in list(fp_model.named_modules()):
        if not isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv1d)):
            continue
        if not _has_quant_weight(sd, full_name):
            continue

        parent, child_name = _get_parent(fp_model, full_name)
        w_bits = _infer_weight_bits(sd, full_name, fallback=qcfg.w_bits)

        if isinstance(module, nn.Linear):
            replacement = QuantLinear(
                module,
                w_bits=w_bits,
                a_bits=qcfg.a_bits,
                w_per_channel=qcfg.w_per_channel,
                w_group_size=qcfg.w_group_size,
                w_clip_percentile=qcfg.w_clip_percentile,
                act_mode=qcfg.act_mode,
                act_group_size=qcfg.act_group_size,
                dstg_group_size=qcfg.dstg_group_size,
                act_per_channel=qcfg.a_per_channel,
                a_clip_percentile=qcfg.a_clip_percentile,
                dstg_clip_percentile=qcfg.dstg_clip_percentile,
                token_clip_percentile=qcfg.token_clip_percentile,
                token_outlier_bypass=qcfg.token_outlier_bypass,
                token_outlier_topk=qcfg.token_outlier_topk,
                dcrp_tau=qcfg.dcrp_tau,
                dcrp_zero_ratio=qcfg.dcrp_zero_ratio,
                token_step_over_std=qcfg.token_step_over_std,
                token_zero_ratio=qcfg.token_zero_ratio,
                compute_dtype=qcfg.compute_dtype,
                full_name=full_name,
            )
        elif isinstance(module, nn.Conv2d):
            replacement = QuantConv2d(
                module,
                w_bits=w_bits,
                a_bits=qcfg.a_bits,
                act_per_channel=qcfg.a_per_channel,
                w_per_channel=qcfg.w_per_channel,
                w_clip_percentile=qcfg.w_clip_percentile,
                a_clip_percentile=qcfg.a_clip_percentile,
                compute_dtype=qcfg.compute_dtype,
                full_name=full_name,
            )
        else:
            replacement = QuantConv1d(
                module,
                w_bits=w_bits,
                a_bits=qcfg.a_bits,
                act_per_channel=qcfg.a_per_channel,
                w_per_channel=qcfg.w_per_channel,
                w_clip_percentile=qcfg.w_clip_percentile,
                a_clip_percentile=qcfg.a_clip_percentile,
                compute_dtype=qcfg.compute_dtype,
                full_name=full_name,
            )

        setattr(parent, child_name, replacement)

    msg = fp_model.load_state_dict(sd, strict=False)
    fp_model.to(device)
    if str(qcfg.compute_dtype).lower() in ("fp16", "float16", "half") and device.type == "cuda":
        fp_model.half()

    print(
        "[Builder] Rebuilt quant model. "
        f"missing={len(getattr(msg, 'missing_keys', []))}, "
        f"unexpected={len(getattr(msg, 'unexpected_keys', []))}"
    )
    return fp_model