from pathlib import Path

import pytest
import torch
from pydantic import ValidationError

from style_bert_vits2.models.utils.checkpoints import load_checkpoint, save_checkpoint
from style_bert_vits2.nlp.japanese.user_dict.word_model import UserDictWord


def _user_dict_word(**overrides: object) -> UserDictWord:
    values = {
        "surface": "test",
        "priority": 5,
        "part_of_speech": "名詞",
        "part_of_speech_detail_1": "一般",
        "part_of_speech_detail_2": "*",
        "part_of_speech_detail_3": "*",
        "inflectional_type": "*",
        "inflectional_form": "*",
        "stem": "*",
        "yomi": "テスト",
        "pronunciation": "テスト",
        "accent_type": 1,
        "accent_associative_rule": "*",
    }
    values.update(overrides)
    return UserDictWord.model_validate(values)


def test_user_dict_word_uses_pydantic_v2_validation() -> None:
    word = _user_dict_word()

    assert word.surface == "ｔｅｓｔ"
    assert word.mora_count == 3

    with pytest.raises(ValidationError):
        _user_dict_word(pronunciation="てすと")


def test_checkpoint_round_trip_with_safe_torch_load(tmp_path: Path) -> None:
    checkpoint_path = tmp_path / "checkpoint.pth"
    source = torch.nn.Linear(2, 1)
    optimizer = torch.optim.AdamW(source.parameters())
    save_checkpoint(source, optimizer, 1e-4, 12, checkpoint_path)

    destination = torch.nn.Linear(2, 1)
    destination_optimizer = torch.optim.AdamW(destination.parameters())
    _, _, learning_rate, iteration = load_checkpoint(
        checkpoint_path, destination, destination_optimizer
    )

    assert learning_rate == 1e-4
    assert iteration == 12
    for source_parameter, destination_parameter in zip(
        source.parameters(), destination.parameters()
    ):
        assert torch.equal(source_parameter, destination_parameter)