"""Load TinyModel v11 from an artefact directory.

An artefact dir (e.g. `tiny-model/model/v11/`) contains:
    config.json                # arch + metadata
    artifacts/model_compiled.pt # frozen-FFN attention-retrained weights
    artifacts/model_full.pt     # end of Phase 2 training
    artifacts/train_mask.pt     # optional — training-token mask
    vindex/                     # optional — extracted vindex

Most callers want `load_from_artifacts(path)`, which constructs the
model from `config.json` and loads the checkpoint of their choice.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import torch

from .model import TinyModel


@dataclass
class ModelConfig:
    vocab_size: int
    dim: int
    n_layers: int
    ffn_dim: int
    n_heads: int
    n_kv_heads: int
    max_seq: int
    version: str = "v11"
    notes: dict = field(default_factory=dict)

    @classmethod
    def from_dict(cls, d: dict) -> "ModelConfig":
        known = {f for f in cls.__dataclass_fields__}
        kwargs = {k: v for k, v in d.items() if k in known}
        kwargs["notes"] = {k: v for k, v in d.items() if k not in known}
        return cls(**kwargs)

    def to_dict(self) -> dict:
        out = {
            "version": self.version,
            "vocab_size": self.vocab_size,
            "dim": self.dim,
            "n_layers": self.n_layers,
            "ffn_dim": self.ffn_dim,
            "n_heads": self.n_heads,
            "n_kv_heads": self.n_kv_heads,
            "max_seq": self.max_seq,
        }
        out.update(self.notes)
        return out


def load_config(artefact_dir: str | Path) -> ModelConfig:
    path = Path(artefact_dir) / "config.json"
    with open(path) as f:
        return ModelConfig.from_dict(json.load(f))


def load_from_artifacts(
    artefact_dir: str | Path,
    checkpoint: str = "model_compiled.pt",
    device: Optional[str | torch.device] = None,
    strict: bool = True,
) -> tuple[TinyModel, ModelConfig]:
    """Instantiate TinyModel from a v11-style artefact dir.

    Returns (model, config). The model is set to eval() and moved to
    `device` (auto-selects MPS → CUDA → CPU if None).
    """
    root = Path(artefact_dir)
    config = load_config(root)

    if device is None:
        if torch.backends.mps.is_available():
            device = torch.device("mps")
        elif torch.cuda.is_available():
            device = torch.device("cuda")
        else:
            device = torch.device("cpu")
    elif isinstance(device, str):
        device = torch.device(device)

    model = TinyModel(
        vocab_size=config.vocab_size,
        dim=config.dim,
        n_layers=config.n_layers,
        ffn_dim=config.ffn_dim,
        n_heads=config.n_heads,
        n_kv_heads=config.n_kv_heads,
        max_seq=config.max_seq,
    ).to(device)

    ckpt_path = root / "artifacts" / checkpoint
    state = torch.load(str(ckpt_path), map_location="cpu", weights_only=True)
    model.load_state_dict(state, strict=strict)
    model.eval()
    return model, config