#!/usr/bin/env python3
"""
inspect_model.py — Inspect a Hugging Face model config BEFORE downloading.

Queries the HF Hub API to read config.json without downloading weights.
Reports:
    - Architecture, hidden size, layers, heads
    - Vocabulary size
    - Dtype, rope config
    - Quantization config (bits, group_size, desc_act)
    - Estimated VRAM for different quant levels on RTX 3090

Usage:
    python scripts/inspect_model.py QuantTrio/Qwen3.6-27B-AWQ
    python scripts/inspect_model.py Qwen/Qwen3.6-27B-FP8 --compare
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import httpx

HF_API = "https://huggingface.co/api"
HF_RAW = "https://huggingface.co"


def fetch_json(repo_id: str, filename: str) -> dict | None:
    """Fetch a JSON file from a HF repo without downloading weights."""
    url = f"{HF_RAW}/{repo_id}/raw/main/{filename}"
    try:
        resp = httpx.get(url, timeout=15, follow_redirects=True)
        if resp.status_code == 200:
            return resp.json()
        # Try resolving the default branch
        model_api = f"{HF_API}/models/{repo_id}"
        resp2 = httpx.get(model_api, timeout=10)
        if resp2.status_code == 200:
            branch = resp2.json().get("defaultBranch", "main")
            url = f"{HF_RAW}/{repo_id}/raw/{branch}/{filename}"
            resp3 = httpx.get(url, timeout=15)
            if resp3.status_code == 200:
                return resp3.json()
        return None
    except Exception as e:
        print(f"  ⚠️  Error fetching {filename}: {e}")
        return None


def estimate_vram(
    hidden_size: int,
    num_layers: int,
    num_heads: int,
    num_kv_heads: int,
    vocab_size: int,
    max_seq_len: int = 8192,
    dtype_bytes: int = 2,        # bf16
    kv_dtype_bytes: int = 2,     # KV cache dtype
) -> dict:
    """Estimate VRAM usage for a model at various quant levels.

    Uses the standard formula:
        weights = params * bytes_per_param
        kv_cache = 2 * num_layers * num_kv_heads * head_dim * max_seq_len * kv_dtype_bytes
        activation = ~2-4 GB overhead
    """
    head_dim = hidden_size // num_heads
    kv_head_dim = hidden_size // num_heads  # GQA: kv_heads share head_dim

    # Parameters (approximate)
    # Q, K, V: 3 * hidden_size * hidden_size (but K,V use kv_heads)
    q_size = hidden_size * hidden_size
    k_size = num_kv_heads * kv_head_dim * hidden_size
    v_size = num_kv_heads * kv_head_dim * hidden_size
    # O: hidden_size * hidden_size
    o_size = hidden_size * hidden_size
    # MLP: 2 * hidden_size * intermediate_size + intermediate_size * hidden_size
    # Qwen uses SwiGLU with 3 gates → intermediate ~= hidden_size * 8/3
    intermediate_size = int(hidden_size * 8 / 3)
    mlp_size = 3 * hidden_size * intermediate_size
    # Per-layer total
    per_layer_params = q_size + k_size + v_size + o_size + mlp_size
    # Embedding + LM head
    embedding_params = vocab_size * hidden_size * 2  # tied in most models
    # Layer norm etc (negligible)
    total_params = per_layer_params * num_layers + embedding_params

    # KV cache per layer
    kv_cache_per_layer = 2 * num_kv_heads * kv_head_dim * max_seq_len * kv_dtype_bytes
    kv_cache_total = kv_cache_per_layer * num_layers

    # Overhead (activations, workspace, CUDA context)
    overhead = 2 * 1024**3  # 2 GB

    results = {}
    for name, bpw in [("bf16", 2), ("fp8", 1), ("int4", 0.5), ("int4_g128", 0.55)]:
        weight_bytes = total_params * bpw
        if "fp8" in name:
            kv_bytes = kv_cache_total  # same dtype
        else:
            kv_bytes = kv_cache_total
        total_vram = weight_bytes + kv_bytes + overhead
        results[name] = total_vram / (1024**3)

    return results


def inspect_model(repo_id: str, max_seq_len: int = 8192):
    """Fetch and analyze model config."""
    print(f"\n{'='*70}")
    print(f" MODEL INSPECTOR — {repo_id}")
    print(f"{'='*70}\n")

    config = fetch_json(repo_id, "config.json")
    quant_config = fetch_json(repo_id, "quantize_config.json") or fetch_json(repo_id, "quant_config.json")

    if not config:
        print(f"❌ Could not fetch config.json for {repo_id}")
        print(f"   URL tried: {HF_RAW}/{repo_id}/raw/main/config.json")
        sys.exit(1)

    # Architecture
    arch = config.get("architectures", ["unknown"])[0]
    model_type = config.get("model_type", "unknown")
    hidden_size = config.get("hidden_size", 0)
    num_layers = config.get("num_hidden_layers", config.get("num_layers", 0))
    num_heads = config.get("num_attention_heads", config.get("num_heads", 0))
    num_kv_heads = config.get("num_key_value_heads", num_heads)
    vocab_size = config.get("vocab_size", 0)
    intermediate_size = config.get("intermediate_size", 0)
    rope_theta = config.get("rope_theta", 0)
    max_position = config.get("max_position_embeddings", 0)
    rope_scaling = config.get("rope_scaling", {})
    torch_dtype = config.get("torch_dtype", "auto")

    # Quant info
    quant_method = ""
    quant_bits = 0
    quant_group_size = 0
    quant_desc_act = False
    quant_sym = False

    if quant_config:
        quant_method = quant_config.get("quant_method", quant_config.get("quantization", ""))
        quant_bits = quant_config.get("bits", quant_config.get("w_bit", 0))
        quant_group_size = quant_config.get("group_size", quant_config.get("q_group_size", 0))
        quant_desc_act = quant_config.get("desc_act", False)
        quant_sym = quant_config.get("sym", False)

    # Print
    print("┌─────────────────────────────────────────────────────┐")
    print(f"│ Architecture                                         │")
    print(f"├─────────────────────────────────────────────────────┤")
    print(f"│ Type:        {model_type:<40s} │")
    print(f"│ Hidden:      {hidden_size:<40d} │")
    print(f"│ Layers:      {num_layers:<40d} │")
    print(f"│ Heads:       {num_heads:<40d} │")
    print(f"│ KV Heads:    {num_kv_heads:<40d} │")
    print(f"│ Head dim:    {hidden_size // num_heads if num_heads else 0:<40d} │")
    print(f"│ Vocab:       {vocab_size:<40d} │")
    print(f"│ Intermed:    {intermediate_size:<40d} │")
    print(f"│ Torch dtype: {str(torch_dtype):<40s} │")
    print(f"│ Max pos:     {max_position:<40d} │")
    print(f"│ Rope theta:  {rope_theta:<40.0f} │")
    if rope_scaling:
        print(f"│ Rope scale:  {str(rope_scaling):<40s} │")
    print(f"└─────────────────────────────────────────────────────┘")

    if quant_config:
        print()
        print("┌─────────────────────────────────────────────────────┐")
        print(f"│ Quantization Config                                  │")
        print(f"├─────────────────────────────────────────────────────┤")
        print(f"│ Method:      {quant_method:<40s} │")
        print(f"│ Bits:        {quant_bits:<40d} │")
        print(f"│ Group size:  {quant_group_size:<40d} │")
        print(f"│ Desc act:    {str(quant_desc_act):<40s} │")
        print(f"│ Symmetric:   {str(quant_sym):<40s} │")
        print(f"└─────────────────────────────────────────────────────┘")

    # VRAM estimation
    print()
    print("┌─────────────────────────────────────────────────────┐")
    print(f"│ Estimated VRAM (max_seq_len={max_seq_len})                 │")
    print(f"├─────────────────────────────────────────────────────┤")
    vram = estimate_vram(hidden_size, num_layers, num_heads, num_kv_heads, vocab_size, max_seq_len)
    for name, gb in vram.items():
        bar = "█" * int(gb / 2) + "░" * (12 - int(gb / 2))
        fits = "✅ FITS" if gb <= 23 else "❌ OOM"
        print(f"│ {name:10s}: {gb:5.1f} GB {bar} {fits:<10s} │")
    print(f"└─────────────────────────────────────────────────────┘")

    # Summary
    print()
    print("💡 Recommendations:")
    if hidden_size and num_layers:
        total_params_b = hidden_size * hidden_size * 8 * num_layers / 1e9  # rough
        print(f"   Estimated params: ~{hidden_size / 1000:.0f}B-class model")

    if quant_method:
        if quant_method in ("awq", "gptq"):
            print(f"   {quant_method.upper()} 4-bit — best for RTX 3090 (Ampere)")
            print(f"   Use: --dtype auto (vLLM auto-detects quant)")
        elif quant_method == "fp8":
            print(f"   FP8 — tight on 24GB. Use --kv-cache-dtype fp8 + --gpu-memory-utilization 0.85")

    if num_kv_heads < num_heads and num_heads:
        ratio = num_heads / num_kv_heads
        print(f"   GQA ratio: {ratio:.0f}:1 — KV cache is {ratio:.0f}x smaller than MHA")

    print()
    print(f"   Recommended vLLM flags for RTX 3090:")
    print(f"   --max-model-len {max_seq_len} --gpu-memory-utilization 0.90 --enable-prefix-caching")


def compare_models(models: list[str], max_seq_len: int = 8192):
    """Compare multiple model variants side by side."""
    print(f"\n{'='*70}")
    print(f" MODEL COMPARISON — {len(models)} variants")
    print(f"{'='*70}\n")

    rows = []
    for repo_id in models:
        config = fetch_json(repo_id, "config.json")
        if not config:
            print(f"  ❌ {repo_id}: config.json not available")
            continue

        quant_config = fetch_json(repo_id, "quantize_config.json") or fetch_json(repo_id, "quant_config.json")
        quant_method = ""
        quant_bits = 0
        if quant_config:
            quant_method = quant_config.get("quant_method", "")
            quant_bits = quant_config.get("bits", 0)

        hidden_size = config.get("hidden_size", 0)
        num_layers = config.get("num_hidden_layers", 0)
        num_heads = config.get("num_attention_heads", 0)
        num_kv_heads = config.get("num_key_value_heads", num_heads)
        vocab_size = config.get("vocab_size", 0)

        vram = estimate_vram(hidden_size, num_layers, num_heads, num_kv_heads, vocab_size, max_seq_len)

        rows.append({
            "model": repo_id.split("/")[-1],
            "quant": quant_method or "none",
            "bits": quant_bits or 16,
            "vram_bf16": vram.get("bf16", 0),
            "vram_fp8": vram.get("fp8", 0),
            "vram_int4": vram.get("int4", 0),
            "fits_24gb": vram.get("fp8", 999) <= 23 or vram.get("int4", 999) <= 23,
        })

    # Print table
    print(f"│ {'Model':<35s} │ {'Quant':>8s} │ {'Bits':>5s} │ {'VRAM(bf16)':>10s} │ {'VRAM(fp8)':>10s} │ {'VRAM(int4)':>10s} │ {'Fits 24GB':>10s} │")
    print(f"│{'-'*37}│{'-'*10}│{'-'*7}│{'-'*12}│{'-'*12}│{'-'*12}│{'-'*12}│")
    for r in rows:
        fits = "✅" if r["fits_24gb"] else "❌"
        print(f"│ {r['model']:<35s} │ {r['quant']:>8s} │ {r['bits']:>4d}  │ {r['vram_bf16']:>8.1f} GB │ {r['vram_fp8']:>8.1f} GB │ {r['vram_int4']:>8.1f} GB │ {fits:>10s} │")


def main():
    parser = argparse.ArgumentParser(description="Inspect HF model config without downloading")
    parser.add_argument("model", nargs="?", help="Model repo id (e.g. QuantTrio/Qwen3.6-27B-AWQ)")
    parser.add_argument("--compare", nargs="*", help="Compare multiple models")
    parser.add_argument("--max-seq-len", type=int, default=8192, help="Context length for VRAM estimate")
    parser.add_argument("--all-variants", action="store_true", help="Compare all Qwen 27B quant variants")
    args = parser.parse_args()

    if args.all_variants:
        models = [
            "Qwen/Qwen3.6-27B",
            "Qwen/Qwen3.6-27B-FP8",
            "QuantTrio/Qwen3.6-27B-AWQ",
            "groxaxo/Qwen3.6-27B-GPTQ-Pro-4bit",
        ]
        compare_models(models, args.max_seq_len)
    elif args.compare is not None:
        models = args.compare if args.compare else [
            "QuantTrio/Qwen3.6-27B-AWQ",
            "Qwen/Qwen3.6-27B-FP8",
            "groxaxo/Qwen3.6-27B-GPTQ-Pro-4bit",
        ]
        compare_models(models, args.max_seq_len)
    elif args.model:
        inspect_model(args.model, args.max_seq_len)
    else:
        # Default: compare all Qwen 27B variants
        models = [
            "Qwen/Qwen3.6-27B",
            "Qwen/Qwen3.6-27B-FP8",
            "QuantTrio/Qwen3.6-27B-AWQ",
            "groxaxo/Qwen3.6-27B-GPTQ-Pro-4bit",
        ]
        compare_models(models, args.max_seq_len)


if __name__ == "__main__":
    main()