#!/usr/bin/env python3
"""Focused integrity contract tests."""

import sys
import tempfile
import json
from pathlib import Path

import numpy as np
import soundfile as sf

sys.path.insert(0, str(Path(__file__).parent.parent))

from data_integrity import (initialize_transcript_provenance, sha256_file,
                            validate_export_integrity, validate_text_integrity)


def _row(wav_path: Path, text: str = "hello"):
    row = {"chunk_id": "clip_1", "text": text, "audio_path": str(wav_path)}
    initialize_transcript_provenance(row, transcript_verbatim="raw hello")
    row["audio_sha256"] = sha256_file(wav_path)
    row["audio_hash_scheme"] = "file_sha256_v1"
    row["promotion_status"] = "training_ready"
    return row


def _split_rows(root: Path):
    rows = []
    for index, split in enumerate(("train", "val", "test")):
        wav = root / f"clip_{split}.wav"
        sf.write(wav, np.full(240, index * 0.1, dtype=np.float32), 24000)
        row = _row(wav, text=f"hello {split}")
        row["chunk_id"] = f"clip_{split}"
        row["video_id"] = f"video_{split}"
        row["split"] = split
        row["split_policy"] = "identity_aware_component_hash_v4"
        row["split_component_id"] = f"component_{split}"
        row["split_near_text_audit_status"] = "complete"
        row["split_coverage_status"] = "complete"
        rows.append(row)
    return rows


def test_stale_text_or_audio_is_rejected_before_export():
    with tempfile.TemporaryDirectory() as tmp:
        wav = Path(tmp) / "clip.wav"
        sf.write(wav, np.zeros(240, dtype=np.float32), 24000)
        row = _row(wav)
        row["text"] = "changed"
        ok, reason = validate_text_integrity(row)
        assert not ok and reason == "training_text_alias_mismatch"
        try:
            validate_export_integrity([row])
        except ValueError as exc:
            assert "training_text_alias_mismatch" in str(exc)
        else:
            raise AssertionError("stale text must block export")


def test_export_requires_explicit_promotion():
    with tempfile.TemporaryDirectory() as tmp:
        wav = Path(tmp) / "clip.wav"
        sf.write(wav, np.zeros(240, dtype=np.float32), 24000)
        row = _row(wav)
        row.pop("promotion_status")
        try:
            validate_export_integrity([row])
        except ValueError as exc:
            assert "not_training_ready" in str(exc)
        else:
            raise AssertionError("unpromoted row must block export")


def test_f5_root_arrow_is_built_from_train_rows_only():
    import dataset_exporter

    with tempfile.TemporaryDirectory() as tmp:
        root = Path(tmp)
        rows = _split_rows(root)

        # A stale same-name destination must be replaced by the validated source.
        stale_dst = root / "integrity_custom" / "wavs" / "clip_train.wav"
        stale_dst.parent.mkdir(parents=True)
        sf.write(stale_dst, np.ones(240, dtype=np.float32), 24000)

        calls = []
        original = dataset_exporter.build_arrow_dataset
        dataset_exporter.build_arrow_dataset = lambda records, directory, name: calls.append(
            (directory.name, [record["chunk_id"] for record in records])
        )
        try:
            dataset_exporter.export_f5_tts_dataset(rows, "integrity", root, {})
        finally:
            dataset_exporter.build_arrow_dataset = original

        assert calls[0] == ("integrity_custom", ["clip_train"])
        assert ("val", ["clip_val"]) in calls
        assert ("test", ["clip_test"]) in calls
        assert sha256_file(stale_dst) == sha256_file(root / "clip_train.wav")

        duration_map = json.loads((root / "integrity_custom" / "duration.json").read_text(encoding="utf-8"))
        assert set(duration_map) == {"clip_train.wav"}

        root_manifest = (root / "integrity_custom" / "manifest.jsonl").read_text(encoding="utf-8").splitlines()
        all_manifest = (
            root / "integrity_custom" / "manifest_all_splits.jsonl"
        ).read_text(encoding="utf-8").splitlines()
        assert len(root_manifest) == 1
        assert len(all_manifest) == 3


def test_unverified_export_requires_a_reason():
    import dataset_exporter

    with tempfile.TemporaryDirectory() as tmp:
        root = Path(tmp)
        wav = root / "clip.wav"
        sf.write(wav, np.zeros(240, dtype=np.float32), 24000)
        row = _row(wav)
        row["split"] = "train"
        try:
            dataset_exporter.export_f5_tts_dataset(
                [row],
                "unverified",
                root,
                {"allow_unverified_export": True},
            )
        except ValueError as exc:
            assert "export_override_reason" in str(exc)
        else:
            raise AssertionError("an unverified export must record its reason")


def test_cross_split_identity_leak_is_rejected():
    with tempfile.TemporaryDirectory() as tmp:
        wav = Path(tmp) / "clip.wav"
        sf.write(wav, np.zeros(240, dtype=np.float32), 24000)
        train_row = _row(wav)
        train_row["split"] = "train"
        val_row = dict(train_row)
        val_row["chunk_id"] = "clip_2"
        val_row["split"] = "val"
        try:
            validate_export_integrity([train_row, val_row])
        except ValueError as exc:
            assert "cross_split_audio_sha256_leak" in str(exc)
            assert "cross_split_text_sha256_leak" in str(exc)
        else:
            raise AssertionError("cross-split payload identities must block export")


def test_verified_export_rejects_missing_validation_and_test_splits():
    with tempfile.TemporaryDirectory() as tmp:
        rows = _split_rows(Path(tmp))[:1]
        try:
            validate_export_integrity(rows)
        except ValueError as exc:
            assert "missing_required_splits:val,test" in str(exc)
        else:
            raise AssertionError("verified export must contain every required split")


def test_verified_export_accepts_complete_split_coverage():
    with tempfile.TemporaryDirectory() as tmp:
        validate_export_integrity(_split_rows(Path(tmp)))


def test_insufficient_component_coverage_metadata_is_rejected():
    with tempfile.TemporaryDirectory() as tmp:
        rows = _split_rows(Path(tmp))
        rows[0]["split_component_id"] = "component_train"
        rows[0]["split_coverage_status"] = "insufficient_components"
        try:
            validate_export_integrity(rows)
        except ValueError as exc:
            assert "clip_train:insufficient_split_coverage" in str(exc)
        else:
            raise AssertionError("insufficient component coverage must block verified export")


def test_unverified_export_bypasses_split_coverage_checks():
    with tempfile.TemporaryDirectory() as tmp:
        rows = _split_rows(Path(tmp))[:1]
        rows[0]["split_component_id"] = "component_train"
        rows[0]["split_coverage_status"] = "insufficient_components"
        validate_export_integrity(rows, allow_unverified=True)


def test_verified_export_requires_complete_identity_aware_near_text_audit():
    with tempfile.TemporaryDirectory() as tmp:
        rows = _split_rows(Path(tmp))
        rows[0]["split_near_text_audit_status"] = "candidate_limit_truncated"
        try:
            validate_export_integrity(rows)
        except ValueError as exc:
            assert "near_text_split_audit_incomplete" in str(exc)
        else:
            raise AssertionError("verified export must require a complete near-text split audit")


if __name__ == "__main__":
    test_stale_text_or_audio_is_rejected_before_export()
    test_export_requires_explicit_promotion()
    test_f5_root_arrow_is_built_from_train_rows_only()
    test_unverified_export_requires_a_reason()
    test_cross_split_identity_leak_is_rejected()
    test_verified_export_rejects_missing_validation_and_test_splits()
    test_verified_export_accepts_complete_split_coverage()
    test_insufficient_component_coverage_metadata_is_rejected()
    test_unverified_export_bypasses_split_coverage_checks()
    print("All data_integrity tests passed!")