from __future__ import annotations
import json
from pathlib import Path
import torch

DEFAULT_REPO_ID = "jackyoung27/matrix-sae"


def _root(root=None):
    return Path(root).resolve() if root else Path(__file__).resolve().parent


def load_manifest(root=None):
    return json.loads((_root(root) / "manifest.json").read_text())


def _find(spec, m):
    for e in m["checkpoints"]:
        if e["relative_path"] == spec or e["tag"] == spec:
            return e
    raise KeyError(spec)


def list_checkpoints(root=None, group=None, sae_type=None):
    es = load_manifest(root)["checkpoints"]
    if group: es = [e for e in es if e["group"] == group]
    if sae_type: es = [e for e in es if e["parsed"].get("sae_type") == sae_type]
    return es


def load_checkpoint(spec, device="cpu", root=None):
    r = _root(root)
    e = _find(spec, load_manifest(r))
    import sys
    if str(r) not in sys.path: sys.path.insert(0, str(r))
    from sae import build_sae_from_config
    d = r / e["relative_path"]
    cfg = json.loads((d / "config.json").read_text())
    ckpt = torch.load(d / "best.pt", map_location="cpu", weights_only=False)
    model = build_sae_from_config(cfg, state_dict=ckpt["model_state_dict"])
    model.load_state_dict(ckpt["model_state_dict"])
    return model.to(device).eval(), cfg, e, ckpt


def load_from_hub(spec, repo_id=DEFAULT_REPO_ID, device="cpu", revision=None, cache_dir=None):
    from huggingface_hub import hf_hub_download
    def dl(f): return hf_hub_download(repo_id=repo_id, filename=f, revision=revision, cache_dir=str(cache_dir) if cache_dir else None)
    m = json.loads(Path(dl("manifest.json")).read_text())
    e = _find(spec, m)
    rel = e["relative_path"]
    cfg_p = dl(f"{rel}/config.json")
    ckpt_p = dl(f"{rel}/best.pt")
    sae_p = dl("sae.py")
    import importlib.util, sys
    sae_dir = str(Path(sae_p).parent)
    if sae_dir not in sys.path: sys.path.insert(0, sae_dir)
    s = importlib.util.spec_from_file_location("sae", sae_p)
    assert s and s.loader
    mod = importlib.util.module_from_spec(s)
    s.loader.exec_module(mod)
    cfg = json.loads(Path(cfg_p).read_text())
    ckpt = torch.load(ckpt_p, map_location="cpu", weights_only=False)
    model = mod.build_sae_from_config(cfg, state_dict=ckpt["model_state_dict"])
    model.load_state_dict(ckpt["model_state_dict"])
    return model.to(device).eval(), cfg, e, ckpt