#!/usr/bin/env python3
"""Cross-reference audit: disk ↔ llama-swap ↔ Hermes custom_providers.

Verifies every model in the pipeline exists at all three layers.
Run after any model download, rename, or directory reorg.
"""

import yaml
import re
import os
import sys
from pathlib import Path

LLAMA_SWAP_CONFIG = os.path.expanduser("~/.config/llama-swap/config.yaml")
HERMES_CONFIG = os.path.expanduser("~/.hermes/config.yaml")
MODELS_DIR = "/models/downloads"

def get_disk_models():
    """All .gguf files in the models directory."""
    models = set()
    for f in Path(MODELS_DIR).glob("*.gguf"):
        models.add(f.name)
    return models

def get_llama_swap_models():
    """Parse llama-swap config for model name → file path."""
    with open(LLAMA_SWAP_CONFIG) as f:
        config = yaml.safe_load(f)
    models = {}
    for name, cfg in config.get("models", {}).items():
        cmd = cfg.get("cmd", "")
        m = re.search(r'-m\s+(\S+)', cmd)
        if m:
            models[name] = m.group(1)
        else:
            models[name] = "NO -m FLAG"
    return models

def get_hermes_models():
    """Parse Hermes config for model IDs."""
    with open(HERMES_CONFIG) as f:
        config = yaml.safe_load(f)
    models = {}
    for prov in config.get("custom_providers", []):
        base_url = prov.get("base_url", "?")
        for m in prov.get("models", []):
            models[m["id"]] = base_url
    return models

def main():
    disk = get_disk_models()
    swap = get_llama_swap_models()
    hermes = get_hermes_models()

    errors = 0

    print("=== DISK MODELS ===")
    for m in sorted(disk):
        print(f"  📁 {m}")

    print("\n=== LLAMA-SWAP → DISK ===")
    for name, path in sorted(swap.items()):
        fname = os.path.basename(path)
        ok = fname in disk
        if ok:
            print(f"  ✓ {name} → {fname}")
        else:
            print(f"  ✗ {name} → {fname}  (MISSING ON DISK)")
            errors += 1

    print("\n=== HERMES → LLAMA-SWAP ===")
    swap_ids = set(swap.keys())
    for hid, url in sorted(hermes.items()):
        ok = hid in swap_ids
        if ok:
            print(f"  ✓ {hid} → {url}")
        else:
            print(f"  ✗ {hid} → {url}  (NOT IN LLAMA-SWAP)")
            errors += 1

    # Orphans
    hermes_ids = set(hermes.keys())
    orphans_hermes = hermes_ids - swap_ids
    orphans_swap = swap_ids - hermes_ids
    if orphans_hermes:
        print(f"\n⚠️  Hermes models NOT in llama-swap: {orphans_hermes}")
    if orphans_swap:
        print(f"\n⚠️  llama-swap models NOT in Hermes: {orphans_swap}")

    # Unused disk models
    used = set(os.path.basename(p) for p in swap.values())
    unused = disk - used
    if unused:
        print(f"\n💤 Models on disk NOT in any config: {unused}")

    if errors == 0:
        print("\n✅ ALL MATCH")
    else:
        print(f"\n❌ {errors} MISMATCHES")
        sys.exit(1)

if __name__ == "__main__":
    main()