from __future__ import annotations

from dataclasses import dataclass, replace
import hashlib
import json
from pathlib import Path
from types import MappingProxyType, SimpleNamespace
import threading

import numpy as np
import pytest

import glyph_hybrid_adapter as adapter_module
from glyph_assets import (
    AssetSegment,
    CarrierSegment,
    GenerationAttemptRecord,
    condition_carrier_v1,
)
from glyph_hybrid_adapter import (
    GlyphAdapterRuntime,
    HardenedGlyphHybridAdapter,
)
from glyph_hybrid_evidence import (
    GatePolicy,
    GlyphHybridEvidenceError,
)
from quality_runtime import (
    SQUIM_OBJECTIVE_WEIGHT_SHA256,
)


MODEL_REPO_ID = "OpenFormosa/BlueMagpie-TTS"
MODEL_REVISION = "a" * 40
ECAPA_REPO_ID = "speechbrain/spkrec-ecapa-voxceleb"
ECAPA_REVISION = "b" * 40
SAMPLE_RATE = 48_000


def _sha(value: bytes | str) -> str:
    if isinstance(value, str):
        value = value.encode("utf-8")
    return hashlib.sha256(value).hexdigest()


def _pcm(
    samples: int,
    amplitude: int,
    *,
    terminal_silence: int = 480,
) -> bytes:
    indexes = np.arange(samples, dtype=np.int64)
    values = np.where(indexes % 2, amplitude, -amplitude)
    values[-terminal_silence:] = 0
    return values.astype("<i2").tobytes()


def _waveform(pcm16_le: bytes) -> np.ndarray:
    return (
        np.frombuffer(pcm16_le, dtype="<i2").astype(np.float32)
        / np.float32(32768.0)
    )


def _pin(name: str, uri: str, revision: str, sha256: str | None = None):
    return SimpleNamespace(
        name=name,
        uri=f"hf://{uri}",
        revision=revision,
        sha256=sha256 or _sha(f"{name}:artifact"),
    )


@dataclass(frozen=True)
class _Policy:
    name: str = "quality-primary"
    cjk_cps: float = 4.2
    ascii_cps: float = 4.2
    hard_stop_margin_steps: int = 1


def _plan(raw_text: str = "甲a"):
    canonical = (
        "甲第一個英文字母"
        if raw_text == "甲a"
        else "第一個英文字母"
    )
    if raw_text == "甲a":
        segments = (
            SimpleNamespace(
                ordinal=0,
                kind="carrier",
                absolute_raw_start=0,
                absolute_raw_end=1,
                canonical_start=0,
                canonical_end=1,
                raw_text="甲",
                canonical_text="甲",
                asset_id=None,
            ),
            SimpleNamespace(
                ordinal=1,
                kind="asset_atom",
                absolute_raw_start=1,
                absolute_raw_end=2,
                canonical_start=1,
                canonical_end=8,
                raw_text="a",
                canonical_text="第一個英文字母",
                asset_id="letter-a",
            ),
        )
    else:
        segments = (
            SimpleNamespace(
                ordinal=0,
                kind="asset_atom",
                absolute_raw_start=0,
                absolute_raw_end=1,
                canonical_start=0,
                canonical_end=7,
                raw_text="a",
                canonical_text="第一個英文字母",
                asset_id="letter-a",
            ),
        )
    return SimpleNamespace(
        profile_id="glyph-hybrid-v5",
        grammar_id="ordinal-glyph-v1",
        normalization_id="zh-tw-v1",
        boundary_contract_id="exact-v1",
        canonical_target=canonical,
        segments=segments,
    )


def _bundle(anchor_sha256: str, asset_pcm: bytes):
    asset = SimpleNamespace(
        asset_id="letter-a",
        raw_byte="a",
        canonical_token="第一個英文字母",
        pcm16_le=asset_pcm,
    )
    entry = adapter_module.evidence.TrustedAsset(
        asset_id=asset.asset_id,
        raw_text=asset.raw_byte,
        canonical_text=asset.canonical_token,
        manifest_entry_sha256="0" * 64,
        pcm16_le=asset.pcm16_le,
    )
    entry_sha256 = (
        adapter_module.evidence.compute_trusted_asset_entry_sha256(entry)
    )
    bundle = SimpleNamespace(
        grammar_sha256=_sha("grammar"),
        speaker_anchor_sha256=anchor_sha256,
        join_profile=MappingProxyType(
            {
                "id": "bluemagpie.glyph-hybrid-join.v1",
                "pause_samples": 1_920,
            }
        ),
        source_pins=(
            _pin("asset-builder", "voidful/glyph-builder", "c" * 40),
        ),
        model_pins=(
            _pin("bluemagpie-checkpoint", MODEL_REPO_ID, MODEL_REVISION),
        ),
        evaluator_pins=(
            _pin(
                "breeze-asr-25",
                adapter_module.BREEZE25_MODEL_ID,
                adapter_module.BREEZE25_REVISION,
                adapter_module.BREEZE25_WEIGHT_SHA256,
            ),
            _pin("ecapa", ECAPA_REPO_ID, ECAPA_REVISION),
            _pin(
                "squim",
                "torchaudio/squim-objective",
                "dns2020",
                SQUIM_OBJECTIVE_WEIGHT_SHA256,
            ),
        ),
        assets=(asset,),
    )
    state = SimpleNamespace(
        bundle=bundle,
        asset_entry_sha256_by_asset_id=MappingProxyType(
            {"letter-a": entry_sha256}
        ),
    )
    return bundle, state


def _verification():
    row = SimpleNamespace(comparison=SimpleNamespace())
    return SimpleNamespace(passed=True, candidate_results=(row,))


def _runtime(
    *,
    anchor: np.ndarray,
    generate_chunk,
    transcript: str,
    counters: dict[str, int],
) -> GlyphAdapterRuntime:
    asr_languages: list[str | None] = []

    def transcriber(audio, sample_rate, **kwargs):
        assert sample_rate == SAMPLE_RATE
        assert kwargs["language"] in {
            adapter_module.BREEZE25_PRIMARY_LANGUAGE,
            adapter_module.BREEZE25_CONFIRMATION_LANGUAGE,
        }
        assert kwargs["task"] == "transcribe"
        assert kwargs["max_new_tokens"] == 160
        assert kwargs["max_verification_segments"] == 160
        asr_languages.append(kwargs["language"])
        counters["asr"] += 1
        return transcript

    transcriber.languages = asr_languages

    def speaker_measure(
        audio,
        sample_rate,
        encoder,
        exact_anchor,
        **kwargs,
    ):
        assert sample_rate == SAMPLE_RATE
        assert encoder == "encoder"
        assert np.array_equal(exact_anchor, anchor)
        counters["ecapa"] += 1
        return SimpleNamespace(
            similarity=0.55,
            begin_similarity=0.52,
            end_similarity=0.50,
        )

    def squim_measure(audio, sample_rate):
        assert sample_rate == SAMPLE_RATE
        counters["squim"] += 1
        return SimpleNamespace(stoi=0.91, pesq=1.55, si_sdr=4.0)

    return GlyphAdapterRuntime(
        sample_rate=SAMPLE_RATE,
        inference_steps=10,
        model_repo_id=MODEL_REPO_ID,
        model_revision=MODEL_REVISION,
        ecapa_repo_id=ECAPA_REPO_ID,
        ecapa_revision=ECAPA_REVISION,
        speaker_anchor_sha256=_sha(
            np.asarray(anchor, dtype="<f4").tobytes()
        ),
        speaker_centroid=object(),
        generation_lock=threading.Lock(),
        generate_chunk=generate_chunk,
        generation_policy_resolver=lambda ordinal: _Policy(),
        generation_cfg_resolver=lambda ordinal: 3.0,
        speaker_encoder_factory=lambda: "encoder",
        speaker_anchor_factory=lambda: anchor.copy(),
        gate_policy=GatePolicy(
            sample_rate=SAMPLE_RATE,
            max_cer=0.08,
            max_prefix_cer=0.0,
            max_suffix_cer=0.0,
            min_speaker_similarity=0.105,
            max_boundary_speaker_drop=0.095,
            min_squim_stoi=0.72,
            min_squim_pesq=1.20,
        ),
        asr_max_new_tokens=160,
        asr_max_verification_segments=160,
        renderer_builder=lambda raw, **kwargs: _plan(raw),
        assembly_verifier=lambda *args, **kwargs: counters.__setitem__(
            "assembly",
            counters["assembly"] + 1,
        ),
        breeze_primary_transcriber=transcriber,
        breeze_confirmation_transcriber=transcriber,
        speaker_measure=speaker_measure,
        squim_measure=squim_measure,
    )


def _carrier_case(*, logged_pcm: bytes | None = None):
    anchor = np.asarray([1.0, 0.0], dtype=np.float32)
    # The expanded seven-unit letter token needs enough active duration to
    # remain below the unchanged 4.30 units/s production pace gate.
    asset_pcm = _pcm(48_000, 3_000)
    bundle, state = _bundle(
        _sha(anchor.astype("<f4").tobytes()),
        asset_pcm,
    )
    generated_pcm = _pcm(48_000, 4_000)
    generated_audio = _waveform(generated_pcm)
    logged_audio = _waveform(logged_pcm or generated_pcm)
    outcome = SimpleNamespace(
        audio=logged_audio,
        stop_reason="stop_threshold",
        endpoint_energy_ratio=0.0,
        generated_steps=4,
        hard_stop_steps=5,
    )
    replay_outcome = SimpleNamespace(
        audio=generated_audio,
        stop_reason="stop_threshold",
        endpoint_energy_ratio=0.0,
        generated_steps=4,
        hard_stop_steps=5,
    )
    record = GenerationAttemptRecord(
        ordinal=0,
        attempt_id="model-attempt-0",
        seed=12_345,
        text_sha256=_sha("甲"),
        endpoint_evidence_sha256=_sha("endpoint"),
        generated_chunk_count=1,
        generated_text_units=1,
    )
    reservation = SimpleNamespace(
        carrier_segment_id="glyph-plan-segment-0",
        candidate_index=0,
        candidate_ordinal=0,
        network_conditioned=False,
        attempt_id=record.attempt_id,
        seed=record.seed,
        text="甲",
    )
    completed = SimpleNamespace(
        reservation=reservation,
        endpoint_outcome=outcome,
        record=record,
    )
    carrier = CarrierSegment(
        segment_id=reservation.carrier_segment_id,
        audio=logged_audio,
        selected_attempt_id=record.attempt_id,
        selected_seed=record.seed,
        selected_text_sha256=record.text_sha256,
        selected_endpoint_evidence_sha256=record.endpoint_evidence_sha256,
    )
    asset = AssetSegment(
        segment_id="glyph-plan-segment-1",
        asset_id="letter-a",
    )
    conditioned = condition_carrier_v1(carrier)
    pause = bytes(1_920 * 2)
    output = conditioned + pause + asset_pcm
    carrier_samples = len(conditioned) // 2
    spans = (
        SimpleNamespace(
            kind="carrier",
            segment_id=carrier.segment_id,
            start_sample=0,
            end_sample=carrier_samples,
        ),
        SimpleNamespace(
            kind="pause",
            segment_id="pause-1",
            start_sample=carrier_samples,
            end_sample=carrier_samples + 1_920,
        ),
        SimpleNamespace(
            kind="asset",
            segment_id=asset.segment_id,
            start_sample=carrier_samples + 1_920,
            end_sample=len(output) // 2,
        ),
    )
    assembly = SimpleNamespace(
        sample_rate=SAMPLE_RATE,
        pcm16_le=output,
        public_pcm16_le=output,
        ledger=SimpleNamespace(spans=spans),
    )
    ledger = SimpleNamespace(
        records=(record,),
        completed_attempts=(completed,),
    )
    counters = {
        "generation": 0,
        "asr": 0,
        "ecapa": 0,
        "squim": 0,
        "assembly": 0,
    }

    def generate_chunk(*args, **kwargs):
        counters["generation"] += 1
        assert args[0] == "甲"
        assert kwargs["request_seed"] == record.seed
        assert kwargs["network_conditioned"] is False
        return replay_outcome

    runtime = _runtime(
        anchor=anchor,
        generate_chunk=generate_chunk,
        transcript="甲第一個英文字母",
        counters=counters,
    )
    assert runtime.speaker_anchor_sha256 == bundle.speaker_anchor_sha256
    return SimpleNamespace(
        runtime=runtime,
        raw_text="甲a",
        plan=_plan(),
        profile_state=state,
        segments=(carrier, asset),
        assembly=assembly,
        semantic_commitment=object(),
        request_ledger=ledger,
        final_verification=_verification(),
        independent_final_verification=_verification(),
        counters=counters,
    )


def _call(adapter: HardenedGlyphHybridAdapter, case) -> bytes:
    return adapter(
        raw_text=case.raw_text,
        plan=case.plan,
        profile_state=case.profile_state,
        segments=case.segments,
        assembly=case.assembly,
        semantic_commitment=case.semantic_commitment,
        request_ledger=case.request_ledger,
        final_verification=case.final_verification,
        independent_final_verification=(
            case.independent_final_verification
        ),
    )


def test_adapter_replays_generation_and_reruns_pinned_evaluators(
    monkeypatch,
):
    key_sizes: list[int] = []

    def token_bytes(size: int) -> bytes:
        key_sizes.append(size)
        return b"k" * size

    monkeypatch.setattr(adapter_module.secrets, "token_bytes", token_bytes)
    case = _carrier_case()
    serialized = _call(
        HardenedGlyphHybridAdapter(case.runtime),
        case,
    )
    payload = json.loads(serialized)

    assert key_sizes == [32]
    assert case.counters == {
        "generation": 1,
        "asr": 2,
        "ecapa": 1,
        "squim": 1,
        "assembly": 1,
    }
    assert case.runtime.breeze_primary_transcriber.languages == [
        adapter_module.BREEZE25_PRIMARY_LANGUAGE,
        adapter_module.BREEZE25_CONFIRMATION_LANGUAGE,
    ]
    assert payload["request_proof"]["content_commitment_scheme"] == (
        "hmac-sha256-v1"
    )
    assert payload["request_proof"]["final_gate"]["passed"] is True
    assert payload["request_proof"]["independent_gate"]["passed"] is True
    segments = {
        row["segment_id"]: row
        for row in payload["request_proof"]["segments"]
    }
    calls = {
        row["call_id"]: row
        for row in payload["runtime_call_ledger"]
    }
    for gate_name in ("final_gate", "independent_gate"):
        gate = payload["request_proof"][gate_name]
        assert (
            calls[gate["asr_call_id"]]["target_sha256"]
            == payload["request_proof"]["canonical_target_sha256"]
        )
    assert (
        calls[
            payload["request_proof"]["final_gate"]["asr_call_id"]
        ]["evaluator_pin_name"]
        != calls[
            payload["request_proof"]["independent_gate"]["asr_call_id"]
        ]["evaluator_pin_name"]
    )
    evaluator_pins = {
        pin["name"]: pin
        for pin in payload["renderer_contract"]["evaluator_pins"]
    }
    primary_pin = evaluator_pins[
        adapter_module.BREEZE25_PRIMARY_PIN_NAME
    ]
    confirmation_pin = evaluator_pins[
        adapter_module.BREEZE25_CONFIRMATION_PIN_NAME
    ]
    assert (
        primary_pin["uri"],
        primary_pin["revision"],
        primary_pin["sha256"],
    ) == (
        confirmation_pin["uri"],
        confirmation_pin["revision"],
        confirmation_pin["sha256"],
    )
    assert primary_pin["config_sha256"] != confirmation_pin["config_sha256"]
    asset_segment = next(
        row for row in segments.values() if row["kind"] == "asset"
    )
    asset_load = next(
        row
        for row in calls.values()
        if row["kind"] == "asset_load"
    )
    assert (
        asset_load["asset_id"],
        asset_load["manifest_entry_sha256"],
        asset_load["loaded_pcm_sha256"],
        asset_load["loaded_sample_count"],
    ) == (
        asset_segment["asset_id"],
        asset_segment["asset_manifest_entry_sha256"],
        asset_segment["realized_pcm_sha256"],
        asset_segment["output_sample_count"],
    )
    selection = payload["carrier_selections"][0]
    carrier_segment = segments[selection["segment_id"]]
    selected_model = calls[selection["selected_call_id"]]
    assert (
        selected_model["text_sha256"],
        selected_model["output_pcm_sha256"],
        selected_model["output_sample_count"],
    ) == (
        carrier_segment["carrier_text_sha256"],
        carrier_segment["realized_pcm_sha256"],
        carrier_segment["output_sample_count"],
    )
    assert "letter-a" not in serialized.decode("ascii")
    assert "\\u7532" not in serialized.decode("ascii")
    assert (
        payload["renderer_contract_sha256"]
        == _sha(
            json.dumps(
                payload["renderer_contract"],
                ensure_ascii=True,
                allow_nan=False,
                separators=(",", ":"),
                sort_keys=True,
            ).encode("ascii")
        )
    )
    source_pins = {
        pin["name"]: pin
        for pin in payload["renderer_contract"]["source_pins"]
    }
    exact_runtime_sources = {
        "glyph-ascii-runtime-source": (
            adapter_module.glyph_ascii_module.__file__
        ),
        "glyph-hybrid-plan-source": adapter_module.glyph_plan_module.__file__,
        "glyph-assets-runtime-source": (
            adapter_module.glyph_assets_module.__file__
        ),
        "glyph-hybrid-evidence-source": adapter_module.evidence.__file__,
        "glyph-hybrid-adapter-source": adapter_module.__file__,
        "glyph-generation-runtime-source": (
            case.runtime.generate_chunk.__code__.co_filename
        ),
    }
    assert set(exact_runtime_sources).issubset(source_pins)
    for name, path in exact_runtime_sources.items():
        assert source_pins[name]["sha256"] == _sha(Path(path).read_bytes())


def test_adapter_replays_every_completed_carrier_attempt():
    case = _carrier_case()
    first_completed = case.request_ledger.completed_attempts[0]
    second_pcm = _pcm(48_000, 4_500)
    second_outcome = SimpleNamespace(
        audio=_waveform(second_pcm),
        stop_reason="stop_threshold",
        endpoint_energy_ratio=0.0,
        generated_steps=4,
        hard_stop_steps=5,
    )
    second_record = GenerationAttemptRecord(
        ordinal=1,
        attempt_id="model-attempt-1",
        seed=54_321,
        text_sha256=_sha("甲"),
        endpoint_evidence_sha256=_sha("endpoint-1"),
        generated_chunk_count=1,
        generated_text_units=1,
    )
    second_reservation = SimpleNamespace(
        carrier_segment_id="glyph-plan-segment-0",
        candidate_index=1,
        candidate_ordinal=1,
        network_conditioned=False,
        attempt_id=second_record.attempt_id,
        seed=second_record.seed,
        text="甲",
    )
    second_completed = SimpleNamespace(
        reservation=second_reservation,
        endpoint_outcome=second_outcome,
        record=second_record,
    )
    replay_by_seed = {
        first_completed.record.seed: first_completed.endpoint_outcome,
        second_record.seed: second_outcome,
    }

    def generate_chunk(*args, **kwargs):
        case.counters["generation"] += 1
        assert args[0] == "甲"
        return replay_by_seed[kwargs["request_seed"]]

    case.runtime = replace(
        case.runtime,
        generate_chunk=generate_chunk,
    )
    case.request_ledger = SimpleNamespace(
        records=(first_completed.record, second_record),
        completed_attempts=(first_completed, second_completed),
    )

    payload = json.loads(
        _call(HardenedGlyphHybridAdapter(case.runtime), case)
    )

    model_calls = [
        row
        for row in payload["runtime_call_ledger"]
        if row["kind"] == "model_generation"
    ]
    assert case.counters["generation"] == 2
    assert [row["attempt_ordinal"] for row in model_calls] == [0, 1]
    assert payload["carrier_selections"][0]["attempted_call_ids"] == [
        row["call_id"] for row in model_calls
    ]
    assert (
        payload["carrier_selections"][0]["selected_call_id"]
        == model_calls[0]["call_id"]
    )


def test_adapter_rejects_logged_generation_that_exact_replay_disagrees_with():
    case = _carrier_case(logged_pcm=_pcm(48_000, 5_000))

    with pytest.raises(
        GlyphHybridEvidenceError,
        match="does not match pinned replay",
    ):
        _call(HardenedGlyphHybridAdapter(case.runtime), case)


def test_pure_asset_request_never_invokes_generation_evaluator():
    anchor = np.asarray([1.0, 0.0], dtype=np.float32)
    # A pure expanded letter token carries seven speech units.
    asset_pcm = _pcm(96_000, 3_000)
    _, state = _bundle(
        _sha(anchor.astype("<f4").tobytes()),
        asset_pcm,
    )
    asset = AssetSegment(
        segment_id="glyph-plan-segment-0",
        asset_id="letter-a",
    )
    assembly = SimpleNamespace(
        sample_rate=SAMPLE_RATE,
        pcm16_le=asset_pcm,
        public_pcm16_le=asset_pcm,
        ledger=SimpleNamespace(
            spans=(
                SimpleNamespace(
                    kind="asset",
                    segment_id=asset.segment_id,
                    start_sample=0,
                    end_sample=len(asset_pcm) // 2,
                ),
            )
        ),
    )
    counters = {
        "generation": 0,
        "asr": 0,
        "ecapa": 0,
        "squim": 0,
        "assembly": 0,
    }

    def forbidden_generation(*args, **kwargs):
        counters["generation"] += 1
        raise AssertionError("pure asset request replayed generation")

    runtime = _runtime(
        anchor=anchor,
        generate_chunk=forbidden_generation,
        transcript="第一個英文字母",
        counters=counters,
    )
    case = SimpleNamespace(
        runtime=runtime,
        raw_text="a",
        plan=_plan("a"),
        profile_state=state,
        segments=(asset,),
        assembly=assembly,
        semantic_commitment=object(),
        request_ledger=SimpleNamespace(
            records=(),
            completed_attempts=(),
        ),
        final_verification=_verification(),
        independent_final_verification=_verification(),
    )

    serialized = _call(HardenedGlyphHybridAdapter(runtime), case)
    payload = json.loads(serialized)

    assert counters["generation"] == 0
    assert payload["renderer_contract"]["generation_policy_ids"] == []
    assert payload["request_proof"]["final_gate"][
        "endpoint_evidence_complete"
    ] is True