#!/usr/bin/env python3
"""Validate remote-code loading and checkpoint shapes without allocating weights."""

from __future__ import annotations

from pathlib import Path

from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModel, AutoTokenizer
from transformers.dynamic_module_utils import get_class_from_dynamic_module

from validate_repo import EXPECTED_SPECIAL_TOKENS, load_json, read_safetensors_header


ROOT = Path(__file__).resolve().parent


def main() -> None:
    config = AutoConfig.from_pretrained(ROOT, trust_remote_code=True)
    assert type(config).__name__ == "RaonConfig"
    assert type(config.text_model_config).__name__ == "AXK1Config"
    assert config.max_position_embeddings == config.text_model_config.max_position_embeddings

    tokenizer = AutoTokenizer.from_pretrained(ROOT, trust_remote_code=True)
    assert len(tokenizer) == config.text_model_config.vocab_size
    for token_id, token_text in EXPECTED_SPECIAL_TOKENS.items():
        assert tokenizer.encode(token_text, add_special_tokens=False) == [token_id]

    processor_class = get_class_from_dynamic_module("modeling_raon.RaonProcessor", ROOT)
    processor = processor_class.from_pretrained(ROOT)
    assert processor.sampling_rate == 24_000

    with init_empty_weights(include_buffers=True):
        model = AutoModel.from_config(config, trust_remote_code=True)

    assert type(model).__name__ == "RaonModel"
    assert type(model.text_model).__name__ == "AXK1Model"
    assert next(model.parameters()).device.type == "meta"

    model_shapes = {name: tuple(tensor.shape) for name, tensor in model.state_dict().items()}
    index = load_json("model.safetensors.index.json")
    checkpoint_shapes: dict[str, tuple[int, ...]] = {}
    for shard_name in sorted(set(index["weight_map"].values())):
        header = read_safetensors_header(ROOT / shard_name)
        checkpoint_shapes.update(
            {
                name: tuple(info["shape"])
                for name, info in header.items()
                if name != "__metadata__"
            }
        )

    model_keys = set(model_shapes)
    checkpoint_keys = set(checkpoint_shapes)
    missing = sorted(model_keys - checkpoint_keys)
    unexpected = sorted(checkpoint_keys - model_keys)
    mismatched = sorted(
        name
        for name in model_keys & checkpoint_keys
        if model_shapes[name] != checkpoint_shapes[name]
    )
    assert not missing, f"Checkpoint is missing {len(missing)} model tensors: {missing[:20]}"
    assert not unexpected, f"Checkpoint has {len(unexpected)} unexpected tensors: {unexpected[:20]}"
    assert not mismatched, (
        f"Checkpoint has {len(mismatched)} shape mismatches: "
        f"{[(name, model_shapes[name], checkpoint_shapes[name]) for name in mismatched[:20]]}"
    )

    parameter_count = sum(parameter.numel() for parameter in model.parameters())
    assert parameter_count == index["metadata"]["total_parameters"]
    print(
        "Model validation passed: "
        f"{type(model).__name__} with {type(model.text_model).__name__}, "
        f"{len(model_shapes):,} state entries, "
        f"{parameter_count:,} parameters, "
        "zero missing/unexpected/mismatched tensors."
    )


if __name__ == "__main__":
    main()