"""
Quick connectivity + load test for external SAE sources:
  1. Neuronpedia REST API  — feature metadata + activations endpoint
  2. SAELens via HuggingFace Hub — download a tiny SAE checkpoint

Run:
    python inspection-backend/test_sae_sources.py

No GPU required. Adds ~200 MB to HF cache on first run (SAELens GPT-2 SAE).
"""

import sys
import json
import struct
import tempfile
import traceback
from pathlib import Path

# ── helpers ──────────────────────────────────────────────────────────────────

PASS = "  [PASS]"
FAIL = "  [FAIL]"
SKIP = "  [SKIP]"


def section(title: str):
    print(f"\n{'='*60}")
    print(f"  {title}".encode("ascii", errors="replace").decode("ascii"))
    print(f"{'='*60}")


def result(ok: bool, msg: str):
    print(f"{PASS if ok else FAIL}  {msg}")
    return ok


# ── 1. Neuronpedia ────────────────────────────────────────────────────────────

def test_neuronpedia():
    section("Neuronpedia REST API")
    try:
        import requests
    except ImportError:
        print(f"{SKIP}  'requests' not installed — pip install requests")
        return False

    base = "https://www.neuronpedia.org/api"

    # 1a. Ping: fetch feature 0 from GPT-2 layer 0 (always public, tiny response)
    url = f"{base}/feature/gpt2-small/0-res-jb/0"
    try:
        r = requests.get(url, timeout=15)
        ok = r.status_code == 200
        result(ok, f"GET feature metadata  [{r.status_code}]  {url}")
        if ok:
            data = r.json()
            has_index = "index" in data or "feature_index" in data or isinstance(data, dict)
            result(has_index, f"Response is JSON with keys: {list(data.keys())[:6]}")
    except Exception as e:
        result(False, f"Request failed: {e}")
        return False

    # 1b. Search endpoint (model list) — checks the API is live beyond one feature
    url2 = f"{base}/models"
    try:
        r2 = requests.get(url2, timeout=10)
        ok2 = r2.status_code == 200
        result(ok2, f"GET /models            [{r2.status_code}]")
        if ok2 and r2.text.strip():
            try:
                models = r2.json()
                if isinstance(models, list) and models:
                    result(True, f"Got {len(models)} model entries; first id: {models[0].get('id','?')}")
                else:
                    result(True, f"Response parsed (type={type(models).__name__})")
            except Exception:
                result(True, f"Endpoint live (body not JSON, {len(r2.text)} chars)")
        elif ok2:
            result(True, "Endpoint live (empty body)")
    except Exception as e:
        result(False, f"/models request failed: {e}")

    # 1c. Try the search / explain endpoint for a feature label
    url3 = f"{base}/feature/gpt2-small/0-res-jb/0/explanation"
    try:
        r3 = requests.get(url3, timeout=10)
        # 404 is fine — it means the endpoint path changed; we just want to confirm network
        result(r3.status_code in (200, 404, 422),
               f"Explanation endpoint responded  [{r3.status_code}]")
    except Exception as e:
        result(False, f"Explanation endpoint: {e}")

    return True


# ── 2. SAELens / HuggingFace Hub ─────────────────────────────────────────────

def test_saelens_hf():
    section("SAELens via HuggingFace Hub")

    # 2a. Can we import huggingface_hub?
    try:
        from huggingface_hub import hf_hub_download, list_repo_files
    except ImportError:
        print(f"{SKIP}  'huggingface_hub' not installed — pip install huggingface-hub")
        return False

    result(True, "huggingface_hub importable")

    # SAELens publishes a collection of SAEs at:
    #   EleutherAI/sae-pythia-70m-32k  (smallest public repo, no auth required)
    repo = "EleutherAI/sae-pythia-70m-32k"

    # 2b. List files in the repo (no download yet — just metadata)
    try:
        files = list(list_repo_files(repo, repo_type="model"))
        result(True, f"Listed {len(files)} files in {repo}")
        cfg_files = [f for f in files if f.endswith(".json")]
        pt_files  = [f for f in files if f.endswith(".pt") or f.endswith(".safetensors")]
        result(bool(cfg_files), f"Config files found: {cfg_files[:3]}")
        result(bool(pt_files),  f"Checkpoint files found: {pt_files[:3]}")
    except Exception as e:
        result(False, f"list_repo_files failed: {e}")
        traceback.print_exc()
        return False

    # 2c. Download the config for layer 0 (tiny JSON, < 1 KB)
    #     SAELens layout: "layers.0/cfg.json"
    cfg_candidates = [f for f in cfg_files if "layer" in f.lower() or "cfg" in f.lower()]
    cfg_target = cfg_candidates[0] if cfg_candidates else (cfg_files[0] if cfg_files else None)

    if cfg_target:
        try:
            local = hf_hub_download(repo_id=repo, filename=cfg_target, repo_type="model")
            with open(local) as fh:
                cfg = json.load(fh)
            result(True, f"Downloaded + parsed {cfg_target}")
            result(True, f"Config keys: {list(cfg.keys())[:8]}")
        except Exception as e:
            result(False, f"Config download failed: {e}")
    else:
        print(f"{SKIP}  No JSON config found to download")

    # 2d. Download the smallest checkpoint file (first .pt or .safetensors)
    ckpt_target = pt_files[0] if pt_files else None
    if ckpt_target:
        try:
            import io, os, contextlib
            print(f"\n  Downloading checkpoint: {ckpt_target}  (may take a moment)...")
            # Suppress tqdm progress bar — it emits unicode arrows that break cp1252 terminals
            with contextlib.redirect_stderr(io.StringIO()):
                local_ckpt = hf_hub_download(repo_id=repo, filename=ckpt_target, repo_type="model")
            size_mb = Path(local_ckpt).stat().st_size / 1e6
            result(True, f"Downloaded {ckpt_target}  ({size_mb:.1f} MB)  ->  {local_ckpt}")
        except Exception as e:
            result(False, f"Checkpoint download failed: {e}")
            return False

        # 2e. Load the checkpoint and verify tensor shapes
        try:
            if local_ckpt.endswith(".safetensors"):
                from safetensors import safe_open
                tensors = {}
                with safe_open(local_ckpt, framework="pt", device="cpu") as f:
                    for k in f.keys():
                        tensors[k] = f.get_tensor(k)
                keys = list(tensors.keys())
                result(True, f"safetensors load succeeded -- keys: {keys[:6]}")
                if keys:
                    sample = tensors[keys[0]]
                    result(True, f"Sample tensor '{keys[0]}': shape={tuple(sample.shape)} dtype={sample.dtype}")
            else:
                import torch
                ckpt = torch.load(local_ckpt, map_location="cpu", weights_only=False)
                keys = list(ckpt.keys()) if isinstance(ckpt, dict) else ["(not a dict)"]
                result(True, f"torch.load succeeded -- top-level keys: {keys[:6]}")
        except ImportError as e:
            print(f"{SKIP}  tensor load skipped: {e}")
        except Exception as e:
            result(False, f"tensor load failed: {e}")
    else:
        print(f"{SKIP}  No checkpoint files found to download")

    return True


# ── 3. SAELens Python package (optional) ─────────────────────────────────────

def test_saelens_package():
    section("SAELens Python package (optional)")
    try:
        import sae_lens  # noqa: F401
        result(True, f"sae_lens importable  (version: {getattr(sae_lens, '__version__', 'unknown')})")
        return True
    except ImportError:
        print(f"{SKIP}  sae_lens not installed — pip install sae-lens")
        print(        "          (not required; HF Hub download above is sufficient)")
        return False


# ── 4. Compatibility with Aquin's SparseAutoencoder ──────────────────────────

def test_aquin_sae_compat():
    section("Aquin SAE ↔ SAELens weight compatibility probe")
    try:
        import torch
        import sys, os
        sys.path.insert(0, str(Path(__file__).parent))
        from sae import SparseAutoencoder
    except ImportError as e:
        print(f"{SKIP}  torch or local sae.py not available: {e}")
        return False

    # SAELens GPT-2 SAEs use d_model=768, expansion=8 → n_features=24576 for some layers
    # Pythia-70M uses d_model=512
    d_model, n_features = 512, 4096
    sae = SparseAutoencoder(d_model=d_model, n_features=n_features)
    dummy = torch.randn(1, d_model)
    feats, recon = sae(dummy)
    result(feats.shape == (1, n_features), f"Encode shape: {feats.shape}")
    result(recon.shape == (1, d_model),    f"Decode shape: {recon.shape}")

    # Simulate what loading a SAELens checkpoint would look like:
    # SAELens keys: W_enc, b_enc, W_dec, b_dec, b_pre  (same as Aquin's SAE)
    saelens_like_keys = {"W_enc", "b_enc", "W_dec", "b_dec", "b_pre"}
    aquin_keys = set(k.split(".")[-1] for k in sae.state_dict().keys())
    overlap = saelens_like_keys & aquin_keys
    result(len(overlap) == len(saelens_like_keys),
           f"Key overlap with SAELens format: {sorted(overlap)}")

    return True


# ── main ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    results = {}
    results["neuronpedia"]       = test_neuronpedia()
    results["saelens_hf"]        = test_saelens_hf()
    results["saelens_package"]   = test_saelens_package()
    results["aquin_sae_compat"]  = test_aquin_sae_compat()

    section("Summary")
    all_passed = True
    for name, ok in results.items():
        # saelens_package is optional — SKIP doesn't count as failure
        is_optional = name == "saelens_package"
        status = PASS if ok else (SKIP if (ok is None or is_optional) else FAIL)
        print(f"{status}  {name}")
        if ok is False and not is_optional:
            all_passed = False

    print()
    if all_passed:
        print("  All checks passed — SAE sources look usable.")
    else:
        print("  Some checks failed. See details above.")
    sys.exit(0 if all_passed else 1)