"""
Crystal Vault Manifest Loader

Loads per-model vault manifests that define tensor slot layouts.
"""

import json
from pathlib import Path
from typing import Dict, List, Any, Iterator, Tuple


def load_vault_manifest(manifest_path: str) -> Dict[str, Any]:
    """
    Load a vault manifest JSON file.

    Args:
        manifest_path: Path to vault manifest JSON

    Returns:
        Parsed manifest dict
    """
    path = Path(manifest_path)
    if not path.exists():
        raise FileNotFoundError(f"Vault manifest not found: {manifest_path}")

    with open(path) as f:
        manifest = json.load(f)

    return manifest


def expand_manifest_tensors(manifest: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Expand manifest to full tensor list (including layer tensors).

    Args:
        manifest: Loaded vault manifest

    Returns:
        List of all tensor dicts with name, shape, elements
    """
    tensors = []

    # Add global tensors
    for t in manifest.get("tensors", []):
        tensors.append({
            "name": t["name"],
            "shape": t["shape"],
            "elements": t["elements"],
            "vault_slot": t.get("vault_slot"),
            "role": t.get("role"),
        })

    # Expand layer template
    layer_template = manifest.get("layer_template", {})
    num_layers = manifest.get("architecture", {}).get("num_hidden_layers", 0)
    template = layer_template.get("template", [])

    for layer_idx in range(num_layers):
        for t in template:
            name = f"model.layers.{layer_idx}.{t['suffix']}"
            tensors.append({
                "name": name,
                "shape": t["shape"],
                "elements": t["elements"],
                "vault_slot": f"SEG{3 + layer_idx * len(template) + template.index(t):04d}",
                "role": f"layer_{layer_idx}_{t['suffix'].split('.')[-1]}",
            })

    return tensors


def iter_vault_slots(manifest: Dict[str, Any]) -> Iterator[Tuple[str, List[int], int]]:
    """
    Iterate over all vault slots in order.

    Yields:
        (tensor_name, shape, num_elements)
    """
    for t in expand_manifest_tensors(manifest):
        yield t["name"], t["shape"], t["elements"]


def compute_vault_totals(manifest: Dict[str, Any]) -> Dict[str, Any]:
    """
    Compute totals for a vault manifest.

    Returns:
        Dict with total_tensors, total_elements, total_bytes (float32)
    """
    tensors = expand_manifest_tensors(manifest)
    total_elements = sum(t["elements"] for t in tensors)

    return {
        "total_tensors": len(tensors),
        "total_elements": total_elements,
        "total_bytes_fp32": total_elements * 4,
        "total_bytes_fp16": total_elements * 2,
    }


def get_default_mistral_vault() -> str:
    """Get path to default Mistral-7B vault manifest."""
    return str(Path(__file__).parent / "mistral_7b_vault_v01.json")


__all__ = [
    "load_vault_manifest",
    "expand_manifest_tensors",
    "iter_vault_slots",
    "compute_vault_totals",
    "get_default_mistral_vault",
]
