from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Iterable

import torch
from transformers import AutoConfig, AutoModel

from sensenova_u1 import check_checkpoint_compatibility


@dataclass(frozen=True)
class GroupRule:
    name: str
    prefixes: tuple[str, ...] = ()
    contains: tuple[str, ...] = ()
    excludes: tuple[str, ...] = ()


@dataclass(frozen=True)
class ParamEntry:
    name: str
    numel: int
    dtype: str
    bytes: int


@dataclass(frozen=True)
class ParamGroupStat:
    name: str
    params: int
    trainable_params: int
    bytes: int
    entries: tuple[ParamEntry, ...]


@dataclass(frozen=True)
class ParamCountResult:
    model_path: str
    total_params: int
    trainable_params: int
    total_bytes: int
    groups: tuple[ParamGroupStat, ...]


# NOTE on architecture (SenseNova-U1, MoT):
#   * vision_model.*                       -> visual und.
#   * fm_modules.*                         -> generation-only (visual gen., fm_head, timestep/noise embedders)
#   * language_model.* w/ "_mot_gen"       -> generation expert inside the LLM backbone
#   * language_model.* w/o "_mot_gen"      -> understanding expert inside the LLM backbone
#   * language_model.model.embed_tokens.*  -> token input embedding, used by every text token in both
#                                             pathways (image-gen still embeds the text prompt)
#   * language_model.lm_head.*             -> text-token output projection. Also exercised by the
#                                             generation pathway because t2i-reasoning runs a
#                                             thinking phase that emits text tokens before image
#                                             tokens. Hence both belong to the "shared" group.
DEFAULT_GROUPS: tuple[GroupRule, ...] = (
    GroupRule("generation_transformer", prefixes=("fm_modules",)),
    GroupRule(
        "generation_transformer",
        prefixes=("language_model",),
        contains=("_mot_gen",),
    ),
    GroupRule(
        "shared",
        prefixes=(
            "language_model.model.embed_tokens",
            "language_model.lm_head",
        ),
    ),
    GroupRule("understanding_transformer", prefixes=("vision_model",)),
    GroupRule("understanding_transformer", prefixes=("language_model",)),
)


def format_param_count(n: int) -> str:
    """Format a parameter count using SI suffixes (B = Billion = 1e9)."""
    units = (("B", 1_000_000_000), ("M", 1_000_000), ("K", 1_000))
    for suffix, base in units:
        if abs(n) >= base:
            return f"{n / base:.3f}{suffix}"
    return str(n)


def format_bytes(n: int) -> str:
    """Format a byte count in decimal units (GB = 1e9 bytes), matching SI."""
    units = (("GB", 1_000_000_000), ("MB", 1_000_000), ("KB", 1_000))
    for suffix, base in units:
        if abs(n) >= base:
            return f"{n / base:.3f}{suffix}"
    return f"{n}B"


def build_rules(custom_groups_json: str | None = None) -> tuple[GroupRule, ...]:
    if not custom_groups_json:
        return DEFAULT_GROUPS

    with open(custom_groups_json, "r", encoding="utf-8") as f:
        raw = json.load(f)
    if not isinstance(raw, dict):
        raise ValueError("custom group config must be a JSON object")

    rules: list[GroupRule] = []
    for group_name, prefixes in raw.items():
        if not isinstance(group_name, str):
            raise ValueError("group name must be string")
        if not isinstance(prefixes, list) or not all(isinstance(x, str) for x in prefixes):
            raise ValueError(f"group '{group_name}' prefixes must be list[str]")
        rules.append(GroupRule(group_name, tuple(prefixes)))
    return tuple(rules)


def _rule_matches(rule: GroupRule, param_name: str) -> bool:
    if rule.prefixes and not any(param_name.startswith(p) for p in rule.prefixes):
        return False
    if rule.contains and not any(c in param_name for c in rule.contains):
        return False
    if rule.excludes and any(e in param_name for e in rule.excludes):
        return False
    return bool(rule.prefixes or rule.contains)


def infer_group(param_name: str, rules: Iterable[GroupRule]) -> str:
    for rule in rules:
        if _rule_matches(rule, param_name):
            return rule.name

    lowered = param_name.lower()
    if "embed" in lowered or "embedding" in lowered:
        return "embedding_misc"
    return "other"


class ModelParamInspector:
    def __init__(
        self,
        model_path: str,
        *,
        dtype: torch.dtype = torch.bfloat16,
    ) -> None:
        self.model_path = model_path
        config = AutoConfig.from_pretrained(model_path)
        check_checkpoint_compatibility(config)
        self.model = AutoModel.from_pretrained(model_path, config=config, torch_dtype=dtype)

    def count(self, rules: Iterable[GroupRule]) -> ParamCountResult:
        total_params = 0
        trainable_params = 0
        total_bytes = 0
        group_to_params: dict[str, int] = {}
        group_to_trainable: dict[str, int] = {}
        group_to_bytes: dict[str, int] = {}
        group_to_entries: dict[str, list[ParamEntry]] = {}
        seen_param_ids: set[int] = set()

        for name, param in self.model.named_parameters():
            param_id = id(param)
            if param_id in seen_param_ids:
                continue
            seen_param_ids.add(param_id)

            numel = int(param.numel())
            # element_size() reflects the actual per-element byte width of this
            # parameter, which is robust to mixed-dtype checkpoints (e.g. norms
            # forced to fp32 even when the rest is loaded as bf16).
            nbytes = numel * param.element_size()
            total_params += numel
            total_bytes += nbytes
            if param.requires_grad:
                trainable_params += numel

            group = infer_group(name, rules)
            group_to_params[group] = group_to_params.get(group, 0) + numel
            group_to_bytes[group] = group_to_bytes.get(group, 0) + nbytes
            group_to_entries.setdefault(group, []).append(
                ParamEntry(
                    name=name,
                    numel=numel,
                    dtype=str(param.dtype).replace("torch.", ""),
                    bytes=nbytes,
                )
            )
            if param.requires_grad:
                group_to_trainable[group] = group_to_trainable.get(group, 0) + numel

        groups = tuple(
            ParamGroupStat(
                name=k,
                params=v,
                trainable_params=group_to_trainable.get(k, 0),
                bytes=group_to_bytes.get(k, 0),
                entries=tuple(sorted(group_to_entries.get(k, []), key=lambda e: e.numel, reverse=True)),
            )
            for k, v in sorted(group_to_params.items(), key=lambda x: x[1], reverse=True)
        )
        return ParamCountResult(
            model_path=self.model_path,
            total_params=total_params,
            trainable_params=trainable_params,
            total_bytes=total_bytes,
            groups=groups,
        )
