# echo_cog/index.py
"""
Model index loader for Echo Cog.

Loads the canonical model index (mistral_model_index.json) and provides
EchoCogConfig instances for vendor and regen models.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Tuple, Optional

from .runtime import EchoCogConfig


# Default paths
DEFAULT_INDEX_PATH = Path(__file__).resolve().parents[1] / "receipts" / "mistral_model_index.json"
DEFAULT_LLAMA_CLI = Path.home() / "llama.cpp" / "build" / "bin" / "llama-cli"


def load_index(path: Path = DEFAULT_INDEX_PATH) -> dict:
    """Load the model index JSON."""
    if not path.exists():
        raise FileNotFoundError(f"Model index not found: {path}")
    return json.loads(path.read_text())


def get_model_config(
    model_key: str,
    index_path: Path = DEFAULT_INDEX_PATH,
    llama_cli: Optional[Path] = None,
) -> EchoCogConfig:
    """
    Get an EchoCogConfig for a specific model from the index.

    Args:
        model_key: Key in the models dict (e.g., 'vendor_q4_k_m', 'helix_seed3_v01')
        index_path: Path to the model index JSON
        llama_cli: Optional override for llama-cli path

    Returns:
        EchoCogConfig for the specified model
    """
    data = load_index(index_path)
    models = data.get("models", {})

    if model_key not in models:
        available = list(models.keys())
        raise KeyError(f"Model '{model_key}' not in index. Available: {available}")

    model_info = models[model_key]
    model_path = Path(model_info["path"])

    if not model_path.exists():
        raise FileNotFoundError(f"Model file not found: {model_path}")

    cli = llama_cli or DEFAULT_LLAMA_CLI
    if not cli.exists():
        raise FileNotFoundError(f"llama-cli not found: {cli}")

    # Build system tag from model info
    kind = model_info.get("kind", "unknown")
    quant = model_info.get("quant", "unknown")
    if kind == "helix_regen":
        dna_seed = model_info.get("dna_seed", "")[:8]
        system_tag = f"Echo Cog Helix Regen ({quant}, seed={dna_seed})"
    else:
        system_tag = f"Echo Cog {kind.title()} Mistral ({quant})"

    return EchoCogConfig(
        llama_cli=cli,
        model_path=model_path,
        system_tag=system_tag,
    )


def get_canonical_pair(
    pair_name: str = "substrate_ab_test",
    index_path: Path = DEFAULT_INDEX_PATH,
    llama_cli: Optional[Path] = None,
) -> Tuple[EchoCogConfig, EchoCogConfig]:
    """
    Get a canonical vendor/regen pair from the index.

    Args:
        pair_name: Name of the pair in canonical_pairs dict
        index_path: Path to the model index JSON
        llama_cli: Optional override for llama-cli path

    Returns:
        (vendor_config, regen_config)
    """
    data = load_index(index_path)
    pairs = data.get("canonical_pairs", {})

    if pair_name not in pairs:
        available = list(pairs.keys())
        raise KeyError(f"Pair '{pair_name}' not in index. Available: {available}")

    pair = pairs[pair_name]
    vendor_key = pair["vendor"]
    regen_key = pair["regen"]

    vendor_cfg = get_model_config(vendor_key, index_path, llama_cli)
    regen_cfg = get_model_config(regen_key, index_path, llama_cli)

    return vendor_cfg, regen_cfg


def get_helix_regen_config(
    index_path: Path = DEFAULT_INDEX_PATH,
    llama_cli: Optional[Path] = None,
) -> EchoCogConfig:
    """
    Get the canonical Helix regen model config (helix_seed3_v01).

    This is the main entry point for users who just want
    'the' Helix-regenerated model with Echo Cog.
    """
    return get_model_config("helix_seed3_v01", index_path, llama_cli)


def get_vendor_config(
    index_path: Path = DEFAULT_INDEX_PATH,
    llama_cli: Optional[Path] = None,
) -> EchoCogConfig:
    """
    Get the canonical vendor model config (vendor_q4_k_m).

    Useful for A/B testing against the regen model.
    """
    return get_model_config("vendor_q4_k_m", index_path, llama_cli)
