"""Attach / detach TENP collectors on SparseMoEMLP layers of a DiT."""

from __future__ import annotations

import json
import os
from typing import Iterable

import torch

from inference.model.dit.dit_module import SparseMoEMLP

from .collector import (
    DEFAULT_PERTURB_RATIO,
    DEFAULT_PERTURB_REPEATS,
    DEFAULT_TOKEN_CACHE_PER_BIN,
    LayerCollector,
)
from .temporal import DEFAULT_NUM_BINS

COLLECT_DIR_ENV = "MAGI_TENP_COLLECT_DIR"
COLLECT_LAYERS_ENV = "MAGI_TENP_LAYERS"
COLLECT_BINS_ENV = "MAGI_TENP_NUM_BINS"
COLLECT_SENS_ENV = "MAGI_TENP_SENSITIVITY"
COLLECT_RHO_ENV = "MAGI_TENP_PERTURB_RATIO"
COLLECT_REPEAT_ENV = "MAGI_TENP_PERTURB_REPEATS"
COLLECT_GRAM_ENV = "MAGI_TENP_GRAM"
COLLECT_CACHE_ENV = "MAGI_TENP_TOKEN_CACHE_PER_BIN"


def iter_sparse_moe(model) -> Iterable[tuple[int, SparseMoEMLP]]:
    layers = model.block.layers
    for layer in layers:
        mlp = getattr(layer, "mlp", None)
        if isinstance(mlp, SparseMoEMLP):
            yield int(mlp.layer_idx), mlp


def _env_int(name: str, default: int) -> int:
    raw = os.getenv(name, "").strip()
    if not raw:
        return int(default)
    return int(raw)


def _env_float(name: str, default: float) -> float:
    raw = os.getenv(name, "").strip()
    if not raw:
        return float(default)
    return float(raw)


def _env_bool(name: str, default: bool) -> bool:
    raw = os.getenv(name, "").strip().lower()
    if not raw:
        return bool(default)
    return raw not in ("0", "false", "no", "off")


def attach_collectors(
    model,
    layers: set[int] | None = None,
    device: torch.device | None = None,
    num_bins: int | None = None,
    collect_sensitivity: bool | None = None,
    perturb_ratio: float | None = None,
    perturb_repeats: int | None = None,
    collect_gram: bool | None = None,
    token_cache_per_bin: int | None = None,
) -> dict[int, LayerCollector]:
    collectors: dict[int, LayerCollector] = {}
    bins = DEFAULT_NUM_BINS if num_bins is None else int(num_bins)
    sens = True if collect_sensitivity is None else bool(collect_sensitivity)
    rho = DEFAULT_PERTURB_RATIO if perturb_ratio is None else float(perturb_ratio)
    repeats = DEFAULT_PERTURB_REPEATS if perturb_repeats is None else int(perturb_repeats)
    gram = True if collect_gram is None else bool(collect_gram)
    cache_n = (
        DEFAULT_TOKEN_CACHE_PER_BIN
        if token_cache_per_bin is None
        else max(int(token_cache_per_bin), 0)
    )
    for layer_idx, mlp in iter_sparse_moe(model):
        if layers is not None and int(layer_idx) not in layers:
            continue
        d_int = int(mlp.experts[0].intermediate_size)
        dev = device or mlp.experts[0].down_proj.weight.device
        col = LayerCollector(
            d_int,
            mlp.num_experts,
            torch.device(dev),
            mlp=mlp,
            num_bins=bins,
            collect_sensitivity=sens,
            perturb_ratio=rho,
            perturb_repeats=repeats,
            collect_gram=gram,
            token_cache_per_bin=cache_n,
        )
        mlp.set_mid_observer(col)
        collectors[int(layer_idx)] = col
    return collectors


def set_collectors_timestep(
    model,
    timestep,
    step_index: int | None = None,
    num_steps: int | None = None,
) -> None:
    from inference.model.dit.dit_module import set_sparse_moe_diffusion_timestep

    set_sparse_moe_diffusion_timestep(model, timestep, step_index=step_index, num_steps=num_steps)
    collectors = getattr(model, "_tenp_collectors", None)
    if not collectors:
        return
    if torch.is_tensor(timestep):
        t = float(timestep.detach().float().reshape(-1)[0].item())
    else:
        t = float(timestep)
    for col in collectors.values():
        col.current_timestep = t
        col.current_step = None if step_index is None else int(step_index)
        col.num_steps = None if num_steps is None else int(num_steps)


def detach_collectors(model) -> None:
    for _layer_idx, mlp in iter_sparse_moe(model):
        mlp.set_mid_observer(None)


def parse_layer_set(raw: str | None) -> set[int] | None:
    if raw is None or not str(raw).strip():
        return None
    return {int(x) for x in str(raw).split(",") if x.strip()}


def save_collectors(path: str, collectors: dict[int, LayerCollector]) -> None:
    payload = {
        "version": 2,
        "num_bins": next(iter(collectors.values())).num_bins if collectors else DEFAULT_NUM_BINS,
        "layers": {str(k): col.state_dict() for k, col in collectors.items()},
    }
    torch.save(payload, path)


def load_collector_payload(path: str, map_location="cpu") -> dict:
    return torch.load(path, map_location=map_location, weights_only=False)


def collector_token_summary(collectors: dict[int, LayerCollector]) -> dict:
    payload = {
        "layers": {str(k): col.state_dict() for k, col in collectors.items()},
    }
    return summary_from_payload(payload)


def summary_from_payload(payload: dict) -> dict:
    layers = []
    for layer_s in sorted(payload["layers"], key=lambda x: int(x)):
        layer_sd = payload["layers"][layer_s]
        experts = []
        audio_primary = False
        for eid_s in sorted(layer_sd["stats"], key=lambda x: int(x)):
            st = layer_sd["stats"][eid_s]
            n_mod = [int(x) for x in st["n_tokens_mod"].tolist()]
            total = int(sum(n_mod))
            p_a = (n_mod[1] / total) if total else 0.0
            if p_a > 0.5:
                audio_primary = True
            rec = {
                    "expert_id": int(eid_s),
                    "n_tokens": int(st["n_tokens"]),
                    "n_mod": {"V": n_mod[0], "A": n_mod[1], "T": n_mod[2]},
                    "p_A": round(p_a, 4),
            }
            n_bin = st.get("n_tokens_bin")
            if n_bin is not None:
                rec["n_tokens_bin"] = [int(x) for x in n_bin.tolist()]
            experts.append(rec)
        layers.append(
            {
                "layer": int(layer_s),
                "experts": experts,
                "audio_primary_any": audio_primary,
            }
        )
    return {
        "audio_primary_discount": False,
        "note": "Audio-primary p_mid discount stays off. n_A in text-only infer is not real audio evidence.",
        "layers": layers,
    }


def maybe_attach_from_env(model) -> dict[int, LayerCollector] | None:
    from .unstructured import maybe_attach_capture

    maybe_attach_capture(model)
    out_dir = os.getenv(COLLECT_DIR_ENV, "").strip()
    if not out_dir:
        return None
    layers = parse_layer_set(os.getenv(COLLECT_LAYERS_ENV, ""))
    collectors = attach_collectors(
        model,
        layers=layers,
        num_bins=_env_int(COLLECT_BINS_ENV, DEFAULT_NUM_BINS),
        collect_sensitivity=_env_bool(COLLECT_SENS_ENV, True),
        perturb_ratio=_env_float(COLLECT_RHO_ENV, DEFAULT_PERTURB_RATIO),
        perturb_repeats=_env_int(COLLECT_REPEAT_ENV, DEFAULT_PERTURB_REPEATS),
        collect_gram=_env_bool(COLLECT_GRAM_ENV, False),
        token_cache_per_bin=_env_int(COLLECT_CACHE_ENV, 0),
    )
    model._tenp_collectors = collectors
    model._tenp_collect_dir = out_dir
    layer_ids = sorted(collectors)
    print(
        f"TENP collect dir={out_dir} layers={layer_ids or 'none'} "
        f"bins={_env_int(COLLECT_BINS_ENV, DEFAULT_NUM_BINS)}"
    )
    return collectors


def maybe_save_from_env(model) -> str | None:
    from .unstructured import maybe_prune_captured

    maybe_prune_captured(model)
    collectors = getattr(model, "_tenp_collectors", None)
    out_dir = getattr(model, "_tenp_collect_dir", "") or os.getenv(COLLECT_DIR_ENV, "").strip()
    if not collectors or not out_dir:
        return None
    os.makedirs(out_dir, exist_ok=True)
    stats_path = os.path.join(out_dir, "tenp_stats.pt")
    save_collectors(stats_path, collectors)
    summary = collector_token_summary(collectors)
    summary_path = os.path.join(out_dir, "tenp_collect_summary.json")
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)
        f.write("\n")
    detach_collectors(model)
    print(f"TENP wrote {stats_path} and {summary_path}")
    return stats_path