#!/usr/bin/env python3
"""
Repository model inspector.

This script is designed to work in `--config-only` mode without importing
PyTorch/Diffusers/Transformers. It reads JSON configs from a local Diffusers
repository layout and prints a summary.

With `--params`, it can also compute parameter counts by scanning
`*.safetensors` headers (without loading tensor data into RAM).
"""

import argparse
import json
import math
from pathlib import Path
from typing import Any, Dict, Iterable, Optional


def load_json(path: Path) -> Dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def human_params(value: Optional[int]) -> str:
    if value is None:
        return "n/a"
    if value >= 1_000_000_000:
        return f"{value/1e9:.2f}B"
    return f"{value/1e6:.2f}M"


def read_model_index(model_dir: Path) -> Dict[str, Any]:
    idx_path = model_dir / "model_index.json"
    if not idx_path.exists():
        return {}
    return load_json(idx_path)


def describe_model_index(model_index: Dict[str, Any]) -> None:
    if not model_index:
        return
    print("Pipeline pieces (model_index.json):")
    for key, val in model_index.items():
        if key.startswith("_"):
            continue
        print(f"  {key:14s} -> {val}")
    print()


def detect_pipeline_kind(model_index: Dict[str, Any]) -> str:
    cls = str(model_index.get("_class_name", "")).lower()
    if "zimage" in cls or ("transformer" in model_index and "unet" not in model_index):
        return "zimage"
    if "stable" in cls or "unet" in model_index:
        return "sdxl_like"
    return "unknown"


def iter_safetensors_files(directory: Path) -> Iterable[Path]:
    if not directory.exists():
        return []
    return sorted(p for p in directory.iterdir() if p.is_file() and p.suffix == ".safetensors")


def count_params_from_safetensors(files: Iterable[Path]) -> int:
    from safetensors import safe_open

    total = 0
    for file in files:
        with safe_open(str(file), framework="np") as f:
            for key in f.keys():
                shape = f.get_slice(key).get_shape()
                total += math.prod(shape)
    return int(total)


def zimage_config_only_summary(model_dir: Path, include_params: bool) -> Dict[str, Any]:
    model_index = read_model_index(model_dir)

    te_cfg_path = model_dir / "text_encoder" / "config.json"
    transformer_cfg_path = model_dir / "transformer" / "config.json"
    vae_cfg_path = model_dir / "vae" / "config.json"
    scheduler_cfg_path = model_dir / "scheduler" / "scheduler_config.json"

    te_cfg = load_json(te_cfg_path) if te_cfg_path.exists() else {}
    transformer_cfg = load_json(transformer_cfg_path) if transformer_cfg_path.exists() else {}
    vae_cfg = load_json(vae_cfg_path) if vae_cfg_path.exists() else {}
    scheduler_cfg = load_json(scheduler_cfg_path) if scheduler_cfg_path.exists() else {}

    text_encoder_params = None
    transformer_params = None
    vae_params = None
    if include_params:
        text_encoder_params = count_params_from_safetensors(iter_safetensors_files(model_dir / "text_encoder"))
        transformer_params = count_params_from_safetensors(iter_safetensors_files(model_dir / "transformer"))
        vae_params = count_params_from_safetensors(iter_safetensors_files(model_dir / "vae"))

    print("[Text encoder]")
    if te_cfg:
        arch = te_cfg.get("architectures", [])
        arch_name = arch[0] if isinstance(arch, list) and arch else "n/a"
        print(f"  architecture={arch_name}")
        print(
            "  "
            f"layers={te_cfg.get('num_hidden_layers', 'n/a')}, "
            f"hidden={te_cfg.get('hidden_size', 'n/a')}, "
            f"heads={te_cfg.get('num_attention_heads', 'n/a')}, "
            f"intermediate={te_cfg.get('intermediate_size', 'n/a')}"
        )
        print(f"  vocab={te_cfg.get('vocab_size', 'n/a')}, max_positions={te_cfg.get('max_position_embeddings', 'n/a')}")
    else:
        print("  [warn] missing text_encoder/config.json")
    print(f"  params={human_params(text_encoder_params)}")
    print()

    print("[Transformer]")
    if transformer_cfg:
        print(f"  class={transformer_cfg.get('_class_name', 'n/a')}")
        print(
            "  "
            f"dim={transformer_cfg.get('dim', 'n/a')}, "
            f"layers={transformer_cfg.get('n_layers', 'n/a')}, "
            f"heads={transformer_cfg.get('n_heads', 'n/a')}"
        )
        print(f"  in_channels={transformer_cfg.get('in_channels', 'n/a')}, cap_feat_dim={transformer_cfg.get('cap_feat_dim', 'n/a')}")
        print(f"  patch_size={transformer_cfg.get('all_patch_size', 'n/a')}, f_patch_size={transformer_cfg.get('all_f_patch_size', 'n/a')}")
    else:
        print("  [warn] missing transformer/config.json")
    print(f"  params={human_params(transformer_params)}")
    print()

    print("[VAE]")
    if vae_cfg:
        print(f"  class={vae_cfg.get('_class_name', 'n/a')}")
        print(
            "  "
            f"sample_size={vae_cfg.get('sample_size', 'n/a')}, "
            f"in_channels={vae_cfg.get('in_channels', 'n/a')}, "
            f"latent_channels={vae_cfg.get('latent_channels', 'n/a')}, "
            f"out_channels={vae_cfg.get('out_channels', 'n/a')}"
        )
        print(f"  block_out_channels={vae_cfg.get('block_out_channels', 'n/a')}, scaling_factor={vae_cfg.get('scaling_factor', 'n/a')}")
    else:
        print("  [warn] missing vae/config.json")
    print(f"  params={human_params(vae_params)}")
    print()

    print("[Scheduler]")
    if scheduler_cfg:
        print(
            "  "
            f"class={scheduler_cfg.get('_class_name', 'n/a')}, "
            f"timesteps={scheduler_cfg.get('num_train_timesteps', 'n/a')}, "
            f"shift={scheduler_cfg.get('shift', 'n/a')}"
        )
    else:
        print("  [warn] missing scheduler/scheduler_config.json")
    print()

    return {
        "kind": "zimage",
        "pipeline": model_index,
        "text_encoder": {"config": te_cfg, "params": text_encoder_params},
        "transformer": {"config": transformer_cfg, "params": transformer_params},
        "vae": {"config": vae_cfg, "params": vae_params},
        "scheduler": {"config": scheduler_cfg},
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Inspect a local Diffusers-style repository layout.")
    parser.add_argument("--model-dir", type=Path, default=Path(".."), help="Path to the diffusers pipeline directory.")
    parser.add_argument("--device", default="cpu", help="Unused (kept for CLI compatibility).")
    parser.add_argument("--fp16", action="store_true", help="Unused (kept for CLI compatibility).")
    parser.add_argument("--config-only", action="store_true", help="Read JSON configs and print a summary.")
    parser.add_argument("--params", action="store_true", help="Count parameters from *.safetensors headers (no tensor loading).")
    parser.add_argument("--json-out", type=Path, default=None, help="Write a JSON summary to this path.")
    args = parser.parse_args()

    model_index = read_model_index(args.model_dir)
    if not model_index:
        raise SystemExit(f"model_index.json not found under {args.model_dir}")

    describe_model_index(model_index)
    kind = detect_pipeline_kind(model_index)

    if not args.config_only:
        raise SystemExit("Only --config-only mode is supported by this inspector.")

    if kind != "zimage":
        raise SystemExit(f"Unsupported pipeline kind: {kind} (expected ZImagePipeline-style layout)")

    summary = zimage_config_only_summary(args.model_dir, include_params=args.params)

    if args.json_out is not None:
        args.json_out.parent.mkdir(parents=True, exist_ok=True)
        args.json_out.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
        print(f"[info] wrote JSON summary to {args.json_out}")


if __name__ == "__main__":
    main()