import ast
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pytest
import torch

import latency_timing
from quality_runtime import (
    Breeze25Runtime,
    transcribe_breeze25,
    WhisperRuntime,
    transcribe_verification_whisper,
    transcribe_whisper,
)


ROOT = Path(__file__).resolve().parents[1]


def _clock(monkeypatch, values):
    ticks = iter(values)
    monkeypatch.setattr(latency_timing, "perf_counter", lambda: next(ticks))


def _fields(line):
    assert line.startswith(latency_timing.LATENCY_LOG_PREFIX)
    return dict(
        field.split("=", 1)
        for field in line.removeprefix(latency_timing.LATENCY_LOG_PREFIX).split()
    )


def test_request_timer_emits_fixed_content_free_cumulative_schema(monkeypatch):
    _clock(monkeypatch, (10.0, 11.0, 13.5, 14.0, 15.0, 20.0))
    emitted = []

    with latency_timing.latency_request(emit=emitted.append) as totals:
        with latency_timing.latency_stage("generation"):
            pass
        with latency_timing.latency_stage("breeze25_asr"):
            pass

    assert len(emitted) == 1
    fields = _fields(emitted[0])
    assert tuple(fields) == tuple(
        item
        for stage in latency_timing.LATENCY_STAGE_NAMES
        for item in (f"{stage}_seconds", f"{stage}_calls")
    )
    assert fields["request_total_seconds"] == "10.000000"
    assert fields["request_total_calls"] == "1"
    assert fields["generation_seconds"] == "2.500000"
    assert fields["generation_calls"] == "1"
    assert fields["breeze25_asr_seconds"] == "1.000000"
    assert fields["breeze25_asr_calls"] == "1"
    assert totals.active_depths == {}


def test_nested_request_and_same_stage_emit_and_count_only_once(monkeypatch):
    _clock(monkeypatch, (0.0, 1.0, 5.0, 6.0))
    emitted = []

    with latency_timing.latency_request(emit=emitted.append):
        with latency_timing.latency_request(emit=emitted.append):
            with latency_timing.latency_stage("ecapa"):
                with latency_timing.latency_stage("ecapa"):
                    pass

    fields = _fields(emitted[0])
    assert len(emitted) == 1
    assert fields["request_total_seconds"] == "6.000000"
    assert fields["request_total_calls"] == "1"
    assert fields["ecapa_seconds"] == "4.000000"
    assert fields["ecapa_calls"] == "1"


def test_public_decorator_logs_on_failure_without_request_arguments(
    monkeypatch,
    capsys,
):
    _clock(monkeypatch, (2.0, 3.0))

    @latency_timing.timed_latency_request
    def fail(private_value):
        raise RuntimeError(private_value)

    with pytest.raises(RuntimeError, match="private-payload"):
        fail("private-payload")

    line = capsys.readouterr().out.strip()
    assert line.startswith(latency_timing.LATENCY_LOG_PREFIX)
    assert "private-payload" not in line
    assert _fields(line)["request_total_seconds"] == "1.000000"


class _Processor:
    def __call__(self, audio, **_kwargs):
        rows = len(audio) if isinstance(audio, list) else 1
        return SimpleNamespace(
            input_features=torch.ones((rows, 4, 4)),
            attention_mask=torch.ones((rows, 4), dtype=torch.long),
        )

    def batch_decode(self, token_ids, **_kwargs):
        return ["內容完整"] * int(token_ids.shape[0])


class _Model:
    def generate(self, features, **_kwargs):
        return torch.ones((features.shape[0], 2), dtype=torch.long)


def test_compatibility_asr_aliases_share_one_stage_without_nested_timing():
    assert WhisperRuntime is Breeze25Runtime
    assert transcribe_whisper is transcribe_breeze25
    assert transcribe_verification_whisper is transcribe_breeze25
    runtime = Breeze25Runtime(
        processor=_Processor(),
        model=_Model(),
        device=torch.device("cpu"),
        dtype=torch.float32,
    )
    emitted = []
    waveform = np.ones(4_000, dtype=np.float32)

    with latency_timing.latency_request(emit=emitted.append):
        assert transcribe_whisper(waveform, 16_000, runtime=runtime) == "內容完整"
        assert (
            transcribe_verification_whisper(
                waveform,
                16_000,
                runtime=runtime,
            )
            == "內容完整"
        )

    fields = _fields(emitted[0])
    assert fields["breeze25_asr_calls"] == "2"
    assert float(fields["breeze25_asr_seconds"]) >= 0.0
    assert "turbo_asr_calls" not in fields
    assert "large_v3_asr_calls" not in fields


def _decorator_stage(node):
    for decorator in node.decorator_list:
        if (
            isinstance(decorator, ast.Call)
            and isinstance(decorator.func, ast.Name)
            and decorator.func.id == "timed_latency_stage"
            and len(decorator.args) == 1
            and isinstance(decorator.args[0], ast.Constant)
        ):
            return decorator.args[0].value
    return None


def test_space_wires_every_requested_latency_stage_without_gate_changes():
    app_tree = ast.parse((ROOT / "app.py").read_text(encoding="utf-8"))
    quality_tree = ast.parse(
        (ROOT / "quality_runtime.py").read_text(encoding="utf-8")
    )
    echo_tree = ast.parse(
        (ROOT / "audio_artifact_gate.py").read_text(encoding="utf-8")
    )
    app_functions = {
        node.name: node for node in app_tree.body if isinstance(node, ast.FunctionDef)
    }
    quality_functions = {
        node.name: node
        for node in quality_tree.body
        if isinstance(node, ast.FunctionDef)
    }
    echo_functions = {
        node.name: node
        for node in echo_tree.body
        if isinstance(node, ast.FunctionDef)
    }

    assert _decorator_stage(app_functions["_assemble_trajectory_audio"]) == "assemble"
    generate_source = ast.get_source_segment(
        (ROOT / "app.py").read_text(encoding="utf-8"),
        app_functions["_generate_chunk"],
    )
    assert generate_source is not None
    assert 'with latency_stage("generation"):' in generate_source
    assert (
        _decorator_stage(quality_functions["transcribe_breeze25"])
        == "breeze25_asr"
    )
    assert "transcribe_whisper" not in quality_functions
    assert "transcribe_verification_whisper" not in quality_functions
    assert _decorator_stage(quality_functions["speaker_evidence_from_audio"]) == "ecapa"
    assert (
        _decorator_stage(quality_functions["release_speaker_evidence_from_audio"])
        == "ecapa"
    )
    assert (
        _decorator_stage(quality_functions["squim_objective_evidence_from_audio"])
        == "squim"
    )
    assert _decorator_stage(quality_functions["active_audio_median_f0_hz"]) == "f0"
    assert _decorator_stage(echo_functions["evaluate_echo_smearing"]) == "echo"

    expected_gpu_decorators = {
        "tts_speaker": "speaker_gpu",
        "tts_reference": "reference_gpu",
    }
    for name, gpu_decorator in expected_gpu_decorators.items():
        decorators = {
            decorator.id
            for decorator in app_functions[name].decorator_list
            if isinstance(decorator, ast.Name)
        }
        assert decorators == {gpu_decorator, "timed_latency_request"}