import ast
import hashlib
import os
from pathlib import Path

import numpy as np
import soundfile as sf
import torch


ROOT = Path(__file__).resolve().parents[1]
APP_PATH = ROOT / "app.py"


def _function_node(name: str) -> ast.FunctionDef:
    tree = ast.parse(APP_PATH.read_text(encoding="utf-8"))
    return next(
        node
        for node in tree.body
        if isinstance(node, ast.FunctionDef) and node.name == name
    )


def _isolated_function(name: str, namespace: dict):
    node = _function_node(name)
    module = ast.Module(body=[node], type_ignores=[])
    ast.fix_missing_locations(module)
    exec(compile(module, APP_PATH, "exec"), namespace)
    return namespace[name]


def test_fast_text_chunks_uses_one_bounded_pass_per_chunk():
    split = _isolated_function(
        "_fast_text_chunks",
        {"CHUNK_CHARS": 8},
    )

    assert split("") == ()
    assert split("甲乙丙丁，戊己庚辛壬癸") == (
        "甲乙丙丁，",
        "戊己庚辛壬癸",
    )


def test_fast_assembler_does_not_analyze_or_rerank_audio():
    finish_calls = []

    def finish(waveform, sample_rate, *, fade_ms):
        finish_calls.append((sample_rate, fade_ms))
        return np.asarray(waveform, dtype=np.float32)

    assemble = _isolated_function(
        "_assemble_fast_audio",
        {
            "np": np,
            "SR": 100,
            "FINAL_ENDPOINT_FADE_MS": 3.0,
            "finish_audio": finish,
        },
    )
    output = assemble(
        (
            np.ones(3, dtype=np.float32),
            np.full(2, 2.0, dtype=np.float32),
        )
    )

    assert np.array_equal(output[:3], np.ones(3, dtype=np.float32))
    assert np.array_equal(output[3:15], np.zeros(12, dtype=np.float32))
    assert np.array_equal(output[15:], np.full(2, 2.0, dtype=np.float32))
    assert finish_calls == [(100, 3.0)]


def test_public_hot_path_contains_no_quality_gate_or_retry_cascade():
    forbidden = {
        "run_coverage_adaptive_cascade",
        "verify_candidate",
        "verify_trajectory",
        "require_verified_final_output",
        "transcribe_breeze25",
        "speaker_evidence_from_audio",
        "release_speaker_evidence_from_audio",
        "squim_objective_evidence_from_audio",
        "evaluate_echo_smearing",
        "_maybe_synthesize_glyph_speaker_request",
    }
    for function_name in (
        "_synthesize",
        "_generate_fast_chunk",
        "_assemble_fast_audio",
        "tts_speaker",
        "tts_reference",
    ):
        called = {
            node.func.id
            for node in ast.walk(_function_node(function_name))
            if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
        }
        assert not called.intersection(forbidden)

    fast_source = ast.get_source_segment(
        APP_PATH.read_text(encoding="utf-8"),
        _function_node("_generate_fast_chunk"),
    )
    assert fast_source is not None
    assert '"retry_badcase": False' in fast_source
    assert "retry_badcase_max_times" not in fast_source


def test_startup_does_not_download_or_preload_verifiers():
    tree = ast.parse(APP_PATH.read_text(encoding="utf-8"))
    top_level_calls = [
        node
        for statement in tree.body
        for node in ast.walk(statement)
        if not isinstance(statement, (ast.FunctionDef, ast.ClassDef))
        and isinstance(node, ast.Call)
    ]
    source = APP_PATH.read_text(encoding="utf-8")

    assert "BREEZE25_ASR_DIR = snapshot_download" not in source
    assert not any(
        isinstance(call.func, ast.Name)
        and call.func.id == "_preload_release_cpu_runtime"
        for call in top_level_calls
    )


def test_builtin_speaker_a_is_the_default():
    source = APP_PATH.read_text(encoding="utf-8")
    load_speakers = ast.get_source_segment(source, _function_node("_load_speakers"))
    speaker_source = ast.get_source_segment(source, _function_node("tts_speaker"))

    assert load_speakers is not None
    assert "default_index = 0" in load_speakers
    assert "SPEAKER_GENERATION_SEED = 56_789" in source
    assert 'SPEAKER_B_LABEL = "內建語者 B"' in source
    assert "SPEAKER_B_PROJECTOR_SCALE = 1.125" in source
    assert "SPEAKER_B_EMBEDDING_SHA256 = (" in source
    assert '"speaker_b_embedding.pt"' in load_speakers
    assert "embedding_sha256 != SPEAKER_B_EMBEDDING_SHA256" in load_speakers
    assert 'embedding_payload.get("speaker_id") != "female_voice"' in load_speakers
    assert "centroids[1] = torch.nn.functional.normalize(" in load_speakers
    assert speaker_source is not None
    assert "seed: int = SPEAKER_GENERATION_SEED" in speaker_source
    assert "request_seed=int(seed)" in speaker_source
    assert "speaker_projector_scale=(" in speaker_source
    assert "if speaker == SPEAKER_B_LABEL" in speaker_source


def test_speaker_and_reference_modes_automatically_reserve_long_text_lease():
    source = APP_PATH.read_text(encoding="utf-8")
    lease = _isolated_function(
        "_gpu_lease_seconds",
        {
            "AUTO_LONG_TEXT_CHARS": 80,
            "INTERACTIVE_GPU_LEASE_SECONDS": 60,
            "LONGFORM_GPU_LEASE_SECONDS": 120,
        },
    )

    assert lease("甲" * 80) == 60
    assert lease("甲" * 81) == 120
    assert "def tts_longform(" not in source
    assert 'with gr.Tab("長文")' not in source
    assert "def _speaker_gpu_duration(text, speaker, cfg, steps, speed, seed)" in source
    assert "def _reference_gpu_duration(text, reference_wav, cfg, steps, speed, seed)" in source
    assert 'label="內建語者 Seed"' in source
    assert 'label="參考音色 Seed"' in source


def test_builtin_speaker_b_uses_the_pinned_checkpoint_embedding(tmp_path):
    checkpoint_dir = tmp_path / "checkpoints"
    checkpoint_dir.mkdir()
    base_centroids = torch.tensor(
        [
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
        ]
    )
    speaker_b = torch.tensor([1.0, 2.0, 2.0])
    torch.save(
        {
            "speaker_ids": ["hung_yi_lee", "female_voice"],
            "centroids": base_centroids,
        },
        checkpoint_dir / "speaker_centroids.pt",
    )
    torch.save(
        {
            "speaker_id": "female_voice",
            "embedding": speaker_b,
            "generation_seed": 56_789,
            "speaker_projector_scale": 1.125,
        },
        checkpoint_dir / "speaker_b_embedding.pt",
    )
    embedding_sha256 = hashlib.sha256(
        (checkpoint_dir / "speaker_b_embedding.pt").read_bytes()
    ).hexdigest()
    load_speakers = _isolated_function(
        "_load_speakers",
        {
            "hashlib": hashlib,
            "os": os,
            "torch": torch,
            "MODEL_DIR": str(tmp_path),
            "SPEAKER_GENERATION_SEED": 56_789,
            "SPEAKER_B_PROJECTOR_SCALE": 1.125,
            "SPEAKER_B_EMBEDDING_SHA256": embedding_sha256,
        },
    )

    speakers, default = load_speakers()

    assert default == "內建語者 A"
    torch.testing.assert_close(speakers["內建語者 A"], base_centroids[0])
    torch.testing.assert_close(
        speakers["內建語者 B"],
        torch.nn.functional.normalize(speaker_b, dim=0),
    )


def test_reference_conditioning_uses_raw_reference_audio_once():
    source = APP_PATH.read_text(encoding="utf-8")
    reference_source = ast.get_source_segment(source, _function_node("tts_reference"))
    fast_source = ast.get_source_segment(source, _function_node("_generate_fast_chunk"))
    synthesize_source = ast.get_source_segment(source, _function_node("_synthesize"))

    assert "REFERENCE_GENERATION_SEED = 1_337" in source
    assert "REFERENCE_CFG = 2.0" in source
    assert "REFERENCE_AUDIO_SECONDS = 6.0" in source
    assert reference_source is not None
    assert "extract_windowed_speaker_embedding(" not in reference_source
    assert "_get_ecapa_encoder(" not in reference_source
    assert "None," in reference_source
    assert "seed: int = REFERENCE_GENERATION_SEED" in reference_source
    assert "request_seed=int(seed)" in reference_source
    assert "cfg: float = REFERENCE_CFG" in reference_source
    assert "reference_wav_path=prepared_reference" in reference_source
    assert "short_text_min_cfg=REFERENCE_CFG" in reference_source
    assert "retry" not in reference_source
    assert fast_source is not None
    assert 'kwargs["reference_wav_path"] = reference_wav_path' in fast_source
    assert synthesize_source is not None
    assert "reference_wav_path=reference_wav_path" in synthesize_source


def test_reference_audio_center_crop_matches_the_trained_six_second_cap(tmp_path):
    prepare = _isolated_function(
        "_prepare_reference_audio",
        {
            "os": os,
            "sf": sf,
            "tempfile": __import__("tempfile"),
            "REFERENCE_AUDIO_SECONDS": 6.0,
        },
    )
    sample_rate = 10
    waveform = np.linspace(-0.5, 0.5, 100, dtype=np.float32)
    source_path = tmp_path / "reference.wav"
    sf.write(source_path, waveform, sample_rate, subtype="FLOAT")

    prepared_path, temporary_path = prepare(str(source_path))
    try:
        prepared, prepared_rate = sf.read(prepared_path, dtype="float32")
        assert prepared_rate == sample_rate
        assert prepared.shape == (60,)
        np.testing.assert_allclose(prepared, waveform[20:80], atol=1.0 / 32768.0)
    finally:
        assert temporary_path is not None
        os.unlink(temporary_path)

    short_path = tmp_path / "short.wav"
    sf.write(short_path, waveform[:40], sample_rate, subtype="FLOAT")
    unchanged_path, unchanged_temporary = prepare(str(short_path))
    assert unchanged_path == str(short_path)
    assert unchanged_temporary is None


def test_reference_projector_gain_is_scoped_and_removed():
    source = APP_PATH.read_text(encoding="utf-8")
    synthesize_source = ast.get_source_segment(source, _function_node("_synthesize"))

    assert synthesize_source is not None
    assert "model.speaker_projector.register_forward_hook(" in synthesize_source
    assert "projector_hook.remove()" in synthesize_source