"""Basic smoke test for Inflect-Nano-v1: model loading, text-to-tokens, and synthesis."""
from __future__ import annotations

import sys
from pathlib import Path

# Ensure repo root and vendored frontend are on sys.path
REPO_ROOT = Path(__file__).resolve().parent
VENDORED_FRONTEND = REPO_ROOT / "third_party" / "tiny_tts_frontend"
sys.path.insert(0, str(REPO_ROOT))
sys.path.insert(0, str(VENDORED_FRONTEND))

import numpy as np
import torch

from tiny_tts.nn import commons
from tiny_tts.text import phonemes_to_ids
from tiny_tts.text.english import grapheme_to_phoneme, normalize_text
from tiny_tts.utils import ADD_BLANK

from inflect_nano.text_cleaning import clean_tinytts_text
from inflect_nano.vocoder import HifiGanGenerator, make_config
from inflect_nano.acoustic import MicroFastSpeech, MicroFastSpeechConfig


def test_text_to_tokens():
    """Verify the text -> token pipeline runs end-to-end."""
    text = "Hello world, this is a test."
    cleaned = clean_tinytts_text(text)
    normalized = normalize_text(cleaned)
    phones, tones, _ = grapheme_to_phoneme(normalized)
    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
    if ADD_BLANK:
        phone_ids = commons.insert_blanks(phone_ids, 0)
        tone_ids = commons.insert_blanks(tone_ids, 0)
        lang_ids = commons.insert_blanks(lang_ids, 0)

    assert len(phone_ids) > 0, "Should produce non-empty phone IDs"
    assert len(phone_ids) == len(tone_ids) == len(lang_ids), "All token sequences should have same length"
    print(f"  [PASS] text_to_tokens: {len(phone_ids)} tokens from '{text}'")


def test_load_acoustic():
    """Verify the acoustic model loads without error."""
    device = torch.device("cpu")
    ckpt_path = REPO_ROOT / "weights" / "inflect_nano_v1_acoustic.pt"
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    cfg = MicroFastSpeechConfig(**ckpt["config"])
    model = MicroFastSpeech(cfg).to(device)
    model.load_state_dict(ckpt["model"])
    model.eval()

    params = sum(p.numel() for p in model.parameters())
    print(f"  [PASS] load_acoustic: {params:,} parameters loaded")


def test_load_vocoder():
    """Verify the vocoder loads without error."""
    device = torch.device("cpu")
    ckpt_path = REPO_ROOT / "weights" / "inflect_nano_v1_vocoder.pt"
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    cfg = make_config((ckpt.get("config") or {}).get("variant", "snake_v2mid"))
    model = HifiGanGenerator(cfg).to(device)
    model.load_state_dict(ckpt["generator"])
    model.remove_weight_norm()
    model.eval()

    params = sum(p.numel() for p in model.parameters())
    print(f"  [PASS] load_vocoder: {params:,} parameters loaded")


def test_synthesize():
    """End-to-end: produce a waveform from text."""
    device = torch.device("cpu")

    # Load acoustic
    acoustic_ckpt = torch.load(REPO_ROOT / "weights" / "inflect_nano_v1_acoustic.pt", map_location=device, weights_only=False)
    acoustic_cfg = MicroFastSpeechConfig(**acoustic_ckpt["config"])
    acoustic = MicroFastSpeech(acoustic_cfg).to(device)
    acoustic.load_state_dict(acoustic_ckpt["model"])
    acoustic.eval()
    speakers = acoustic_ckpt.get("speakers") or {"mark": 0}

    # Load vocoder
    vocoder_ckpt = torch.load(REPO_ROOT / "weights" / "inflect_nano_v1_vocoder.pt", map_location=device, weights_only=False)
    vocoder_cfg = make_config((vocoder_ckpt.get("config") or {}).get("variant", "snake_v2mid"))
    vocoder = HifiGanGenerator(vocoder_cfg).to(device)
    vocoder.load_state_dict(vocoder_ckpt["generator"])
    vocoder.remove_weight_norm()
    vocoder.eval()

    # Tokenize
    text = "This is a quick test."
    cleaned = clean_tinytts_text(text)
    normalized = normalize_text(cleaned)
    phones, tones, _ = grapheme_to_phoneme(normalized)
    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
    if ADD_BLANK:
        phone_ids = commons.insert_blanks(phone_ids, 0)
        tone_ids = commons.insert_blanks(tone_ids, 0)
        lang_ids = commons.insert_blanks(lang_ids, 0)

    phone = torch.LongTensor(phone_ids).unsqueeze(0).to(device)
    tone = torch.LongTensor(tone_ids).unsqueeze(0).to(device)
    lang = torch.LongTensor(lang_ids).unsqueeze(0).to(device)
    speaker = torch.LongTensor([int(speakers.get("mark", 0))]).to(device)

    # Synthesize
    with torch.inference_mode():
        mel = acoustic.infer(phone, tone, lang, speaker)
        wav = vocoder(mel).squeeze().cpu().numpy()

    assert isinstance(wav, np.ndarray), "Output should be a numpy array"
    assert wav.size > 0, "Waveform should not be empty"
    assert np.abs(wav).max() <= 1.0, "Waveform should be normalized to [-1, 1]"

    print(f"  [PASS] synthesize: {wav.size} samples ({wav.size / 24000:.2f}s) from '{text}'")


if __name__ == "__main__":
    print("Inflect-Nano-v1 smoke tests\n")
    test_text_to_tokens()
    test_load_acoustic()
    test_load_vocoder()
    test_synthesize()
    print("\nAll tests passed!")