from __future__ import annotations
import json
import sys
from pathlib import Path
import torch
from safetensors.torch import load_model as load_safetensors_model
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer


def _backbone(model):
    base = getattr(model, "model", model)
    return getattr(base, "language_model", base)


def load_model(repo_dir=".", device=None):
    root = Path(repo_dir).resolve()
    if str(root) not in sys.path:
        sys.path.insert(0, str(root))
    meta = json.loads((root / "standalone_config.json").read_text())
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    config = AutoConfig.from_pretrained(root, trust_remote_code=True, local_files_only=True)
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
    variant = meta["variant"]

    if variant == "memory_fusion":
        from tinycenn_lm.qwen3_5_memory_fusion import Qwen35MemoryFusionConfig, replace_attention_layers
        replace_attention_layers(
            model,
            Qwen35MemoryFusionConfig.from_dict(meta["variant_config"]),
            meta["layers"],
        )
        model.config.use_cache = False
    elif variant == "cenn_integrated":
        from tinycenn_lm.optimized_memory import OptimizedMemory
        from tinycenn_lm.qwen35_integrated_memory import Qwen35IntegratedAttention
        bb = _backbone(model)
        for raw in meta["layers"]:
            i = int(raw)
            original = bb.layers[i].self_attn
            core = OptimizedMemory(**meta["core_configs"][str(i)])
            bb.layers[i].self_attn = Qwen35IntegratedAttention(original, core, i)
    elif variant == "pdelta3_clvr":
        from tinycenn_lm.qwen35_pdelta3_runtime import QwenPDelta3CLVRConfig, replace_full_attention_layers
        replace_full_attention_layers(
            model,
            QwenPDelta3CLVRConfig.from_dict(meta["variant_config"]),
            meta["layers"],
        )
        model.config.use_cache = False
    else:
        raise ValueError(f"Unknown TinyCeNN variant: {variant}")

    missing, unexpected = load_safetensors_model(
        model,
        str(root / "model.safetensors"),
        strict=True,
        device="cpu",
    )
    if missing or unexpected:
        raise RuntimeError(
            f"Strict weight load failed: missing={list(missing)[:8]} "
            f"unexpected={list(unexpected)[:8]}"
        )
    tokenizer = AutoTokenizer.from_pretrained(root, local_files_only=True, use_fast=True)
    return model.to(device).eval(), tokenizer