from pathlib import Path

BASE_DIR = Path.home() / "inspection-backend"

MODEL_CONFIGS = {
    "llama-3.2-1b": {
        "hf_name":    "meta-llama/Llama-3.2-1B-Instruct",
        "d_model":    2048,
        "n_layers":   16,
        "sae_layer":  8,
        "sae_path":   "sae/sae_layer8.pt",
        "norm_path":  "sae/acts_norm.pt",
        "acts_path":  "sae/acts_cache.pt",
        "edit_layer": None,
        "sae_source": "native",
        "neuronpedia_slug": None,
        "supports_inspection":     True,
        "supports_weight_editing": False,
    },
    "pythia-2.8b": {
        "hf_name":    "EleutherAI/pythia-2.8b",
        "d_model":    2560,
        "n_layers":   32,
        "sae_layer":  16,
        "sae_path":   "sae/sae_layer16_pythia.pt",
        "norm_path":  "sae/acts_norm_pythia.pt",
        "acts_path":  "sae/acts_cache_pythia.pt",
        "edit_layer": 16,
        "sae_source": "native",
        "neuronpedia_slug": None,
        "supports_inspection":     False,
        "supports_weight_editing": True,
    },
    # ── Neuronpedia-backed: weights from HF, labels from Neuronpedia API ──────
    "gpt2-small-np": {
        "hf_name":    "gpt2",
        "d_model":    768,
        "n_layers":   12,
        "sae_layer":  8,
        "sae_path":   "sae/sae_gpt2small_layer8_np.safetensors",
        "norm_path":  None,
        "acts_path":  None,
        "edit_layer": None,
        "sae_source": "neuronpedia",
        # HF repo + file path for the SAE weights
        "neuronpedia_hf_repo":   "jbloom/GPT2-Small-SAEs-Reformatted",
        "neuronpedia_hf_file":   "blocks.8.hook_resid_pre/sae_weights.safetensors",
        # Neuronpedia model/layer slug for the explanation API
        "neuronpedia_slug":      "gpt2-small/8-res-jb",
        "supports_inspection":     True,
        "supports_weight_editing": False,
    },
    # ── HuggingFace-backed: weights from HF (SAELens), labels via GPT-4o-mini ─
    "pythia-70m-hf": {
        "hf_name":    "EleutherAI/pythia-70m-deduped",
        "d_model":    512,
        "n_layers":   6,
        "sae_layer":  3,
        "sae_path":   "sae/sae_pythia70m_layer3_hf.safetensors",
        "norm_path":  None,
        "acts_path":  None,
        "edit_layer": None,
        "sae_source": "huggingface",
        "hf_repo":    "EleutherAI/sae-pythia-70m-32k",
        "hf_file":    "layers.3/sae.safetensors",
        "neuronpedia_slug": None,
        "supports_inspection":     True,
        "supports_weight_editing": False,
    },
}


def get_config(model_id: str) -> dict:
    if model_id not in MODEL_CONFIGS:
        raise ValueError(f"Unknown model_id '{model_id}'. Available: {list(MODEL_CONFIGS.keys())}")
    return MODEL_CONFIGS[model_id]


def get_sae_layer(model_id: str) -> int:
    return get_config(model_id)["sae_layer"]


def get_edit_layer(model_id: str) -> int:
    cfg = get_config(model_id)
    if cfg["edit_layer"] is None:
        raise ValueError(f"Model '{model_id}' does not support weight editing.")
    return cfg["edit_layer"]


def supports_inspection(model_id: str) -> bool:
    return get_config(model_id)["supports_inspection"]


def supports_weight_editing(model_id: str) -> bool:
    return get_config(model_id)["supports_weight_editing"]


def get_sae_source(model_id: str) -> str:
    return get_config(model_id).get("sae_source", "native")


def get_sae_path(model_id: str) -> Path:
    return BASE_DIR / get_config(model_id)["sae_path"]


def get_norm_path(model_id: str) -> Path | None:
    p = get_config(model_id).get("norm_path")
    return BASE_DIR / p if p else None


def get_acts_path(model_id: str) -> Path | None:
    p = get_config(model_id).get("acts_path")
    return BASE_DIR / p if p else None
