import math

import numpy as np
import pytest

from quality_runtime import (
    CandidateObservation,
    CandidateVerification,
    ChunkCandidateArtifact,
    NoQualifiedCandidateError,
    SEQUENCE_FALLBACK_MAX_LOCAL_BOUNDARY_SPEAKER_DROP,
    TrajectoryGateResult,
    run_coverage_adaptive_cascade,
    select_culprit_diverse_candidate_sequences,
    trajectory_gate_evidence,
    verify_candidate,
    verify_trajectory,
)


def _local_verification(chunks, passing=None, *, scores=None):
    if passing is None:
        passing = (True,) * len(chunks)
    if scores is None:
        scores = (0.8,) * len(chunks)
    observations = []
    artifacts = []
    for chunk, passed, similarity in zip(chunks, passing, scores, strict=True):
        observations.append(
            CandidateObservation(
                target_text=chunk,
                transcript_text=chunk if passed else "完全錯誤",
                audio_duration_seconds=2.0,
                speaker_similarity=similarity,
                begin_speaker_similarity=similarity,
                end_speaker_similarity=similarity,
                pace_cps=2.0,
            )
        )
        artifacts.append(
            ChunkCandidateArtifact(
                speaker_embedding=np.array([1.0, 0.0], dtype=np.float32),
                rms_db=-20.0,
            )
        )
    return verify_trajectory(
        observations,
        chunk_artifacts=artifacts,
        max_pace_cps=4.3,
    )


def _joined_rejected_local(chunks, passing=None):
    local = _local_verification(chunks, passing)
    return TrajectoryGateResult(
        passed=False,
        candidate_results=local.candidate_results,
        score=math.inf,
        rejection_reasons=("joined_output:semantic_gate",),
        chunk_artifacts=local.chunk_artifacts,
    )


def _exact_final(chunks, passed=True):
    target = "".join(chunks)
    transcript = target if passed else "錯誤內容"
    return verify_trajectory(
        [CandidateObservation(target, transcript, 1.0)],
        chunk_artifacts=[ChunkCandidateArtifact(rms_db=-20.0)],
    )


def test_whole_exact_winner_returns_before_any_refill_or_sequence_check():
    chunks = ("第一段", "第二段")
    generated = []
    refill_calls = []
    final_calls = []

    def generator(candidate_chunks, seed):
        generated.append((candidate_chunks, seed))
        return tuple(f"{seed}:{chunk}" for chunk in candidate_chunks)

    def whole_verifier(trajectory, candidate_chunks, seed):
        return CandidateVerification(
            _local_verification(candidate_chunks),
            joined_output=trajectory_gate_evidence(
                _exact_final(candidate_chunks)
            ),
        )

    result = run_coverage_adaptive_cascade(
        chunks,
        80,
        generator,
        whole_verifier,
        lambda *args: refill_calls.append(args),
        sequence_final_verifier=lambda *args: final_calls.append(args),
    )

    assert generated == [(chunks, 80)]
    assert refill_calls == []
    assert final_calls == []
    assert result.selection_mode == "whole_trajectory"
    assert result.chunk_candidate_indices == (0, 0)
    assert result.chunk_seeds == (80, 80)
    assert result.attempted_seeds == (80,)
    assert result.generated_chunk_count == 2
    assert result.generated_text_units == 6
    assert result.chunk_candidate_counts == (1, 1)
    assert len(result.diagnostics.attempts) == 1
    assert result.diagnostics.attempts[0].joined_output is not None
    assert result.diagnostics.attempts[0].joined_output.passed


def test_refill_budget_accounts_exact_generated_chunks_and_text_units():
    chunks = ("甲乙", "丙丁")
    generated = []

    def generator(candidate_chunks, seed):
        generated.append((candidate_chunks, seed))
        return tuple(f"{seed}:{chunk}" for chunk in candidate_chunks)

    result = run_coverage_adaptive_cascade(
        chunks,
        10,
        generator,
        lambda trajectory, candidate_chunks, seed: _joined_rejected_local(
            candidate_chunks
        ),
        lambda trajectory, candidate_chunks, seed: _local_verification(
            candidate_chunks
        ),
        sequence_final_verifier=lambda result, candidate_chunks: _exact_final(
            candidate_chunks
        ),
        max_generated_chunks=20,
        max_generated_text_units=6,
        max_sequence_paths=3,
    )

    assert generated == [
        (chunks, 10),
        (("甲乙",), 11),
    ]
    assert result.generated_chunk_count == 3
    assert result.generated_text_units == 6
    assert result.chunk_candidate_counts == (2, 1)
    assert result.attempted_seeds == (10, 11)
    assert result.chunk_candidate_indices == (1, 0)
    assert len(result.diagnostics.attempts) == 2
    assert result.diagnostics.sequence_search is not None
    assert result.diagnostics.sequence_search.eligible_candidate_counts == (2, 1)
    assert result.diagnostics.sequence_search.ranked_path_count >= 1
    assert [
        path.final_output.passed
        for path in result.diagnostics.sequence_search.checked_paths
    ] == [True]


def test_one_path_mode_still_creates_one_alternative_to_rejected_initial_audio():
    generated = []

    result = run_coverage_adaptive_cascade(
        ("完整內容",),
        10,
        lambda chunks, seed: (
            generated.append((chunks, seed))
            or tuple(f"{seed}:{chunk}" for chunk in chunks)
        ),
        lambda trajectory, chunks, seed: _joined_rejected_local(chunks),
        lambda trajectory, chunks, seed: _local_verification(chunks),
        sequence_final_verifier=lambda result, chunks: _exact_final(chunks),
        max_generated_chunks=2,
        max_generated_text_units=100,
        max_sequence_paths=1,
    )

    assert generated == [(("完整內容",), 10), (("完整內容",), 11)]
    assert result.chunk_candidate_indices == (1,)
    assert result.generated_chunk_count == 2


@pytest.mark.parametrize(
    ("budget_name", "value", "frozen_cap"),
    [
        ("max_generated_chunks", 21, 20),
        ("max_generated_text_units", 801, 800),
    ],
)
def test_generation_budgets_cannot_expand_beyond_frozen_caps(
    budget_name,
    value,
    frozen_cap,
):
    with pytest.raises(ValueError, match=f"frozen cap of {frozen_cap}"):
        run_coverage_adaptive_cascade(
            ("完整內容",),
            10,
            lambda chunks, seed: tuple(chunks),
            lambda trajectory, chunks, seed: _local_verification(chunks),
            lambda trajectory, chunks, seed: _local_verification(chunks),
            sequence_final_verifier=lambda result, chunks: _exact_final(chunks),
            **{budget_name: value},
        )


def test_refill_schedule_and_selected_path_are_deterministic():
    chunks = ("甲乙", "丙丁", "戊己")

    def run_once():
        generated = []
        checked = []

        def generator(candidate_chunks, seed):
            generated.append((candidate_chunks, seed))
            return tuple(f"{seed}:{chunk}" for chunk in candidate_chunks)

        result = run_coverage_adaptive_cascade(
            chunks,
            20,
            generator,
            lambda trajectory, candidate_chunks, seed: _joined_rejected_local(
                candidate_chunks
            ),
            lambda trajectory, candidate_chunks, seed: _local_verification(
                candidate_chunks
            ),
            sequence_final_verifier=lambda result, candidate_chunks: (
                checked.append(result.chunk_candidate_indices)
                or _exact_final(candidate_chunks)
            ),
            max_generated_chunks=8,
            max_generated_text_units=100,
            max_sequence_paths=2,
        )
        return generated, checked, result

    first_generated, first_checked, first = run_once()
    second_generated, second_checked, second = run_once()

    assert first_generated == second_generated
    assert first_checked == second_checked
    assert first.chunk_candidate_indices == second.chunk_candidate_indices
    assert first.chunk_seeds == second.chunk_seeds
    assert first.generated_chunk_count == second.generated_chunk_count == 6
    assert first.chunk_candidate_counts == second.chunk_candidate_counts == (2, 2, 2)


def test_zero_coverage_rows_are_refilled_before_low_coverage_rows_round_robin():
    chunks = ("第一段", "第二段", "第三段")
    generated = []

    def generator(candidate_chunks, seed):
        generated.append((candidate_chunks, seed))
        return tuple(f"{seed}:{chunk}" for chunk in candidate_chunks)

    def refill_verifier(trajectory, candidate_chunks, seed):
        # The first attempt for the second chunk fails.  The zero-coverage
        # third chunk must get the next seed before the second chunk retries.
        return _local_verification(candidate_chunks, passing=(seed != 101,))

    result = run_coverage_adaptive_cascade(
        chunks,
        100,
        generator,
        lambda trajectory, candidate_chunks, seed: _joined_rejected_local(
            candidate_chunks,
            passing=(True, False, False),
        ),
        refill_verifier,
        sequence_final_verifier=lambda result, candidate_chunks: _exact_final(
            candidate_chunks
        ),
        max_generated_chunks=6,
        max_generated_text_units=100,
        max_sequence_paths=1,
    )

    assert generated[1:] == [
        (("第二段",), 101),
        (("第三段",), 102),
        (("第二段",), 103),
    ]
    assert result.chunk_candidate_indices == (0, 3, 2)
    assert result.chunk_candidate_counts == (1, 1, 1)


def test_ragged_dp_accepts_unequal_row_widths_and_stable_cost_order():
    ranked = select_culprit_diverse_candidate_sequences(
        [[0.0], [0.3, 0.1, 0.2], [0.4, 0.0]],
        [
            [[0.0, 0.0, 0.0]],
            [
                [0.0, 0.0],
                [0.0, 0.0],
                [0.0, 0.0],
            ],
        ],
        max_paths=3,
    )

    assert [selection.candidate_indices for selection in ranked] == [
        (0, 1, 1),
        (0, 2, 1),
        (0, 0, 1),
    ]
    assert [selection.total_score for selection in ranked] == pytest.approx(
        [0.1, 0.2, 0.3]
    )


def test_culprit_diversity_uses_final_checks_on_distinct_culprit_assignments():
    ranked = select_culprit_diverse_candidate_sequences(
        [[0.0, 0.001, 0.002], [0.0, 10.0, 20.0], [0.0]],
        [
            [[0.0, 0.0, 0.0]] * 3,
            [[0.0], [0.0], [0.0]],
        ],
        culprit_indices=(1,),
        max_paths=3,
    )

    assert [selection.candidate_indices for selection in ranked] == [
        (0, 0, 0),
        (0, 1, 0),
        (0, 2, 0),
    ]
    assert len({selection.candidate_indices[1] for selection in ranked}) == 3


@pytest.mark.parametrize(
    ("generator", "whole_verifier", "message"),
    [
        (
            lambda chunks, seed: np.zeros(8, dtype=np.float32),
            lambda trajectory, chunks, seed: _local_verification(chunks),
            "malformed trajectory",
        ),
        (
            lambda chunks, seed: tuple(chunks),
            lambda trajectory, chunks, seed: None,
            "invalid result",
        ),
        (
            lambda chunks, seed: tuple(chunks),
            lambda trajectory, chunks, seed: (_ for _ in ()).throw(
                RuntimeError("ASR unavailable")
            ),
            "whole trajectory verification failed",
        ),
    ],
)
def test_malformed_or_asr_initial_evidence_aborts_fail_closed(
    generator,
    whole_verifier,
    message,
):
    with pytest.raises(RuntimeError, match=message):
        run_coverage_adaptive_cascade(
            ("完整內容",),
            10,
            generator,
            whole_verifier,
            lambda *args: _local_verification(("完整內容",)),
            sequence_final_verifier=lambda result, chunks: _exact_final(chunks),
        )


def test_passing_speaker_gate_without_ecapa_artifact_aborts_fail_closed():
    local = _local_verification(("完整內容",))
    malformed = TrajectoryGateResult(
        passed=True,
        candidate_results=local.candidate_results,
        score=local.score,
        rejection_reasons=(),
        chunk_artifacts=(ChunkCandidateArtifact(rms_db=-20.0),),
    )

    with pytest.raises(RuntimeError, match="acoustic evidence"):
        run_coverage_adaptive_cascade(
            ("完整內容",),
            10,
            lambda chunks, seed: tuple(chunks),
            lambda *args: malformed,
            lambda *args: local,
            sequence_final_verifier=lambda result, chunks: _exact_final(chunks),
        )


def test_final_verifier_exception_invalid_evidence_and_all_reject_fail_closed():
    common = dict(
        chunks=("甲乙", "丙丁"),
        root_seed=10,
        candidate_generator=lambda chunks, seed: tuple(
            f"{seed}:{chunk}" for chunk in chunks
        ),
        whole_candidate_verifier=lambda trajectory, chunks, seed: (
            _joined_rejected_local(chunks)
        ),
        refill_candidate_verifier=lambda trajectory, chunks, seed: (
            _local_verification(chunks)
        ),
        max_generated_chunks=6,
        max_generated_text_units=100,
        max_sequence_paths=3,
    )

    with pytest.raises(RuntimeError, match="refusing unverified audio"):
        run_coverage_adaptive_cascade(
            **common,
            sequence_final_verifier=lambda *args: (_ for _ in ()).throw(
                RuntimeError("large-v3 OOM")
            ),
        )
    with pytest.raises(RuntimeError, match="invalid result"):
        run_coverage_adaptive_cascade(
            **common,
            sequence_final_verifier=lambda *args: None,
        )

    checked = []
    with pytest.raises(
        NoQualifiedCandidateError,
        match="3 assembled paths",
    ) as caught:
        run_coverage_adaptive_cascade(
            **common,
            sequence_final_verifier=lambda result, chunks: (
                checked.append(result.chunk_candidate_indices)
                or _exact_final(chunks, passed=False)
            ),
        )
    assert len(checked) == 3
    assert len(set(checked)) == 3
    search = caught.value.diagnostics.sequence_search
    assert search is not None
    assert search.ranked_path_count == 3
    assert len(search.checked_paths) == 3
    assert all(not path.final_output.passed for path in search.checked_paths)


def test_boundary_only_relaxation_is_capped_and_still_requires_exact_final_pass():
    chunks = ("第一段完整", "第二段完整")
    chunk = chunks[0]
    rejected = verify_candidate(
        CandidateObservation(
            target_text=chunk,
            transcript_text=chunk,
            audio_duration_seconds=2.0,
            speaker_similarity=0.8,
            begin_speaker_similarity=0.8,
            end_speaker_similarity=0.651,
        ),
        max_boundary_speaker_drop=0.10,
    )
    safe_second = _local_verification((chunks[1],))
    local = TrajectoryGateResult(
        passed=False,
        candidate_results=(rejected, safe_second.candidate_results[0]),
        score=math.inf,
        rejection_reasons=("chunk_0:boundary_speaker_drop",),
        chunk_artifacts=(
            ChunkCandidateArtifact(
                speaker_embedding=np.array([1.0, 0.0], dtype=np.float32),
                rms_db=-20.0,
            ),
            safe_second.chunk_artifacts[0],
        ),
    )
    callbacks = []
    result = run_coverage_adaptive_cascade(
        chunks,
        10,
        lambda chunks, seed: tuple(chunks),
        lambda *args: local,
        lambda *args: local,
        sequence_final_verifier=lambda result, chunks: (
            callbacks.append(result.chunk_candidate_indices)
            or _exact_final(chunks)
        ),
        max_generated_chunks=2,
        max_generated_text_units=100,
        max_sequence_paths=1,
        sequence_fallback_max_local_boundary_speaker_drop=(
            SEQUENCE_FALLBACK_MAX_LOCAL_BOUNDARY_SPEAKER_DROP
        ),
    )

    assert callbacks == [(0, 0)]
    assert result.selection_mode == "coverage_sequence_dp"
    assert result.verification.candidate_results[0].boundary_speaker_drop == pytest.approx(
        0.149
    )

    with pytest.raises(NoQualifiedCandidateError):
        run_coverage_adaptive_cascade(
            chunks,
            10,
            lambda chunks, seed: tuple(chunks),
            lambda *args: local,
            lambda *args: local,
            sequence_final_verifier=lambda result, chunks: _exact_final(
                chunks,
                passed=False,
            ),
            max_generated_chunks=2,
            max_generated_text_units=100,
            max_sequence_paths=1,
            sequence_fallback_max_local_boundary_speaker_drop=0.15,
        )
    with pytest.raises(ValueError, match="no greater than 0.15"):
        run_coverage_adaptive_cascade(
            chunks,
            10,
            lambda chunks, seed: tuple(chunks),
            lambda *args: local,
            lambda *args: local,
            sequence_final_verifier=lambda result, chunks: _exact_final(chunks),
            sequence_fallback_max_local_boundary_speaker_drop=0.151,
        )

    single_local = TrajectoryGateResult(
        passed=False,
        candidate_results=(rejected,),
        score=math.inf,
        rejection_reasons=("chunk_0:boundary_speaker_drop",),
        chunk_artifacts=(local.chunk_artifacts[0],),
    )
    with pytest.raises(NoQualifiedCandidateError, match="zero safe coverage"):
        run_coverage_adaptive_cascade(
            (chunk,),
            10,
            lambda candidate_chunks, seed: tuple(candidate_chunks),
            lambda *args: single_local,
            lambda *args: single_local,
            sequence_final_verifier=lambda result, candidate_chunks: _exact_final(
                candidate_chunks
            ),
            max_generated_chunks=1,
            max_generated_text_units=100,
            max_sequence_paths=1,
            sequence_fallback_max_local_boundary_speaker_drop=0.15,
        )