#!/usr/bin/env python3
"""
GGUF v3 Metadata Scanner

Reads architecture, context length, layer count, head counts, quantization, 
HF source, and file size from GGUF files without requiring the gguf Python package.
Handles GGUF v3 (u64 key lengths), big tokenizer arrays, and partial buffers.

Usage:
    python3 gguf-metadata-scan.py [path/*.gguf]

Output: JSON lines, one per model file.
"""

import struct, os, json, sys

INTERESTING_KEYS = {
    "general.name", "general.architecture", "general.file_type",
    "general.size_label", "llama.context_length", "general.base_model",
    "general.parameter_count", "llama.block_count", "llama.attention.layer_count",
    "llama.attention.head_count", "llama.attention.head_count_kv",
    "general.quantization_version", "tokenizer.ggml.model",
    "general.base_model.0.name", "general.base_model.0.repo_url",
    "general.base_model.0.organization", "general.quantized_by",
    "general.basename",
}

# Key prefixes to always collect (architecture-specific keys)
INTERESTING_PREFIXES = (
    "general.", "llama.", "tokenizer.ggml.model",
    "qwen", "gemma", "deepseek", "gpt-oss", "nemotron",
)

FILE_TYPE_MAP = {
    0: "F32", 1: "F16", 2: "Q4_0", 3: "Q4_1", 5: "Q5_0", 6: "Q5_1",
    7: "Q8_0", 10: "Q2_K", 12: "Q4_K", 13: "Q5_K", 14: "Q6_K", 15: "Q8_K",
    30: "Q4_K_M", 31: "Q4_K_S", 32: "Q5_K_M", 33: "Q5_K_S", 34: "Q6_K",
}

def scan_gguf(path, buf_size=2*1024*1024):
    """Extract metadata from a GGUF v3 file by reading the header section."""
    fname = os.path.basename(path)
    try:
        size = os.path.getsize(path)
        with open(path, 'rb') as f:
            magic = f.read(4)
            if magic != b'GGUF':
                return {"file": fname, "disk_gb": round(size/1e9, 1), "error": "not GGUF"}
            ver = struct.unpack('<I', f.read(4))[0]
            tc = struct.unpack('<Q', f.read(8))[0]
            kc = struct.unpack('<Q', f.read(8))[0]
            hdr = f.read(min(buf_size - f.tell(), size - f.tell()))
    except (OSError, IOError) as e:
        return {"file": fname, "error": str(e)}

    meta = {
        "file": fname,
        "disk_gb": round(size / 1e9, 1),
        "version": ver,
        "tensors": tc,
        "kv_entries": kc,
    }

    pos = 0
    n = len(hdr)

    for i in range(kc):
        if pos + 8 > n:
            break
        kl = struct.unpack('<Q', hdr[pos:pos+8])[0]
        pos += 8
        if pos + kl > n:
            break
        key = hdr[pos:pos+kl].decode('utf-8', errors='replace')
        pos += kl
        if pos + 4 > n:
            break
        vt = struct.unpack('<I', hdr[pos:pos+4])[0]
        pos += 4

        # Only parse keys we care about (saves time on large tokenizer arrays)
        should_parse = (
            key in INTERESTING_KEYS
            or key.startswith(INTERESTING_PREFIXES)
            and not key.startswith(("tokenizer.", "general.unknown", "transformer."))
        )

        if vt == 8:  # string
            if pos + 8 > n:
                break
            sl = struct.unpack('<Q', hdr[pos:pos+8])[0]
            pos += 8
            if pos + sl > n:
                break
            if should_parse:
                meta[key] = hdr[pos:pos+sl].decode(errors='replace')
            pos += sl
        elif vt == 10:  # uint64
            if pos + 8 > n:
                break
            if should_parse:
                meta[key] = struct.unpack('<Q', hdr[pos:pos+8])[0]
            pos += 8
        elif vt == 11:  # int64
            if pos + 8 > n:
                break
            if should_parse:
                meta[key] = struct.unpack('<q', hdr[pos:pos+8])[0]
            pos += 8
        elif vt in (4, 5):  # uint32/int32
            if pos + 4 > n:
                break
            if should_parse:
                meta[key] = struct.unpack('<I', hdr[pos:pos+4])[0]
            pos += 4
        elif vt == 6:  # float32
            if pos + 4 > n:
                break
            if should_parse:
                meta[key] = round(struct.unpack('<f', hdr[pos:pos+4])[0], 2)
            pos += 4
        elif vt == 7:  # bool
            if pos + 1 > n:
                break
            if should_parse:
                meta[key] = bool(struct.unpack('B', hdr[pos:pos+1])[0])
            pos += 1
        elif vt == 0:  # uint8
            if pos + 1 > n:
                break
            if should_parse:
                meta[key] = struct.unpack('B', hdr[pos:pos+1])[0]
            pos += 1
        elif vt == 9:  # array
            if pos + 12 > n:
                break
            arr_t = struct.unpack('<I', hdr[pos:pos+4])[0]
            arr_len = struct.unpack('<Q', hdr[pos+4:pos+12])[0]
            pos += 12
            if should_parse:
                meta[key] = f"<array[{arr_len}]>"
            # Skip elements
            if arr_t == 8:
                for _ in range(arr_len):
                    if pos + 8 > n:
                        break
                    sl = struct.unpack('<Q', hdr[pos:pos+8])[0]
                    pos += 8
                    if pos + sl > n:
                        break
                    pos += sl
            else:
                elem_size = {0:1, 1:1, 2:2, 3:2, 4:4, 5:4, 6:4, 7:1, 10:8, 11:8, 12:8}.get(arr_t, 4)
                pos += min(arr_len * elem_size, n - pos)
        else:
            # Unknown type — skip
            pass

    # Architecture-specific keys
    for key in list(meta.keys()):
        if 'context_length' in key and isinstance(meta[key], (int, float)):
            meta['context_length'] = meta[key]
        if key.startswith('general.'):
            short = key.replace('general.', '', 1)
            if short not in meta:
                meta[short] = meta[key]

    # Map file_type code to name
    if 'file_type' in meta and isinstance(meta['file_type'], int):
        meta['quant'] = FILE_TYPE_MAP.get(meta['file_type'], f"Q_type{meta['file_type']}")

    meta['_kv_read'] = i + 1
    return meta


SPEED_CACHE = {
    # Model filename fragment -> (t/s, vram_gb)
    "qwen36-35b": ("~60", "~9.2"),
    "Qwen3.6-35B-A3B-MTP": ("~60", "~9.2"),
    "gemma-4-26B": ("~45", "~7.7"),
    "gemma-4-12B": ("~123", "~7.7"),
    "Nemotron-Terminal-14B": ("~90", "~6.5"),
    "GLM-4.7-Flash-Q4_K_M": ("~33", "~6.0"),
    "GLM-4.7-Flash-REAP": ("~43", "~5.8"),
    "gpt-oss-20b": ("N/A", "N/A"),
    "mtp-gemma-4-12B": ("—", "—"),
    "mtp-gemma-4-26B": ("—", "—"),
}

def lookup_speed(fname):
    """Look up known speed/VRAM from benchmark cache by filename fragment."""
    for key, (speed, vram) in SPEED_CACHE.items():
        if key in fname:
            return speed, vram
    return None, None


def format_table(results, markdown=False):
    """Pretty-print a Markdown table from scan results. Includes speed/VRAM if known."""
    rows = []
    for r in results:
        fname = r.get('file', '?')
        params = r.get('size_label') or r.get('general.size_label', '?')
        arch = r.get('architecture') or r.get('general.architecture', '?')
        quant = r.get('quant') or '?'
        ctx_raw = r.get('context_length', '?')
        if isinstance(ctx_raw, int):
            ctx = f"{ctx_raw//1000}K" if ctx_raw >= 1000 else str(ctx_raw)
        else:
            ctx = str(ctx_raw)
        disk = f"{r.get('disk_gb', '?'):.1f}" if isinstance(r.get('disk_gb'), (int, float)) else '?'
        speed = r.get('speed', 'N/A')
        vram = r.get('vram', '?')
        hf = r.get('base_model.0.repo_url', '') or r.get('general.base_model.0.repo_url', '')
        hf_short = hf.replace('https://huggingface.co/', '') if hf else '?'
        rows.append((fname, params, quant, disk, vram, ctx, speed, arch, hf_short))

    sep = "|" if markdown else " "
    header = f"{sep} Model {sep} Params {sep} Quant {sep} Disk {sep} VRAM {sep} Ctx {sep} t/s {sep} HF Source"
    ruler = f"{sep}{'─'*48}{sep}{'─'*14}{sep}{'─'*10}{sep}{'─'*8}{sep}{'─'*6}{sep}{'─'*8}{sep}{'─'*6}{sep}{'─'*40}"
    print(header)
    if markdown:
        print(f"{sep}:--{sep}:--{sep}:--{sep}:--{sep}:--{sep}:--{sep}:--{sep}:--")
    else:
        print(ruler)

    for r in rows:
        cols = f"{sep} {r[0]:<46} {sep} {r[1]:<12} {sep} {r[2]:<8} {sep} {r[3]:<6} {sep} {r[4]:<4} {sep} {r[5]:<6} {sep} {r[6]:<4} {sep} {r[8]}"
        print(cols)
    print()
    print(f"Total: {len(rows)} models")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Scan GGUF model metadata')
    parser.add_argument('paths', nargs='*', help='GGUF files to scan (default: /models/downloads/*.gguf)')
    parser.add_argument('--speed-data', action='store_true', help='Include speed/VRAM from benchmark cache')
    parser.add_argument('--markdown', action='store_true', help='Output markdown table')
    args = parser.parse_args()

    paths = args.paths if args.paths else []
    if not paths:
        paths = sorted([
            p for p in os.listdir('/models/downloads/')
            if p.endswith('.gguf')
        ])
        paths = [os.path.join('/models/downloads/', p) for p in paths]

    results = []
    for p in paths:
        meta = scan_gguf(p)
        if args.speed_data:
            speed, vram = lookup_speed(os.path.basename(p))
            if speed:
                meta['speed'] = speed
            if vram:
                meta['vram'] = vram
        results.append(meta)

    if sys.stdout.isatty() or args.markdown:
        format_table(results, markdown=args.markdown)
    else:
        print(json.dumps(results, indent=2, default=str))