"""Tests for the Gold Corpus split management."""

from __future__ import annotations

import hashlib
import json
import uuid

import pytest

from services.intelligence_pipeline_v3.gold_corpus.sampler import (
    CompanyCountBucket,
    Difficulty,
    DocumentMetadata,
    LengthBucket,
    SourceType,
)
from services.intelligence_pipeline_v3.gold_corpus.splits import (
    CorpusSplit,
    SplitConfig,
    SplitManifest,
    create_splits,
    freeze_holdout,
    select_hard_cases,
)

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


def _make_doc(difficulty: Difficulty = Difficulty.EASY) -> DocumentMetadata:
    return DocumentMetadata(
        document_id=str(uuid.uuid4()),
        document_type="article",
        event_class="earnings_beat",
        length_bucket=LengthBucket.MEDIUM,
        source_type=SourceType.NEWS,
        company_count_bucket=CompanyCountBucket.SINGLE,
        difficulty=difficulty,
    )


def _make_corpus(size: int = 200) -> list[DocumentMetadata]:
    """Generate a mixed-difficulty corpus."""
    import random

    rng = random.Random(42)
    difficulties = [Difficulty.EASY, Difficulty.MEDIUM, Difficulty.HARD]
    return [_make_doc(rng.choice(difficulties)) for _ in range(size)]


# ---------------------------------------------------------------------------
# Tests — SplitConfig
# ---------------------------------------------------------------------------


class TestSplitConfig:
    def test_default_ratios_sum_to_one(self) -> None:
        config = SplitConfig()
        assert config.validate_ratios()

    def test_invalid_ratios_detected(self) -> None:
        config = SplitConfig(
            train_ratio=0.5,
            calibration_ratio=0.5,
            holdout_ratio=0.5,
            agreement_ratio=0.1,
        )
        assert not config.validate_ratios()

    def test_custom_ratios(self) -> None:
        config = SplitConfig(
            train_ratio=0.70,
            calibration_ratio=0.10,
            holdout_ratio=0.15,
            agreement_ratio=0.05,
        )
        assert config.validate_ratios()


# ---------------------------------------------------------------------------
# Tests — create_splits
# ---------------------------------------------------------------------------


class TestCreateSplits:
    def test_all_splits_present(self) -> None:
        corpus = _make_corpus(200)
        splits = create_splits(corpus)
        assert CorpusSplit.TRAIN in splits
        assert CorpusSplit.CALIBRATION in splits
        assert CorpusSplit.HOLDOUT in splits
        assert CorpusSplit.ANNOTATOR_AGREEMENT in splits

    def test_all_documents_assigned(self) -> None:
        corpus = _make_corpus(200)
        splits = create_splits(corpus)
        total = sum(s.total_count for s in splits.values())
        assert total == len(corpus)

    def test_no_overlap_between_splits(self) -> None:
        corpus = _make_corpus(200)
        splits = create_splits(corpus)
        all_ids: list[str] = []
        for manifest in splits.values():
            all_ids.extend(manifest.document_ids)
        assert len(all_ids) == len(set(all_ids))

    def test_holdout_is_frozen(self) -> None:
        corpus = _make_corpus(200)
        splits = create_splits(corpus)
        holdout = splits[CorpusSplit.HOLDOUT]
        assert holdout.frozen is True
        assert holdout.frozen_at is not None
        assert "prompt_tuning" in holdout.restricted_uses
        assert "model_training" in holdout.restricted_uses

    def test_approximate_split_ratios(self) -> None:
        corpus = _make_corpus(1000)
        config = SplitConfig()
        splits = create_splits(corpus, config)

        total = len(corpus)
        # Allow 5% tolerance on ratios
        train_ratio = splits[CorpusSplit.TRAIN].total_count / total
        holdout_ratio = splits[CorpusSplit.HOLDOUT].total_count / total

        assert 0.50 < train_ratio < 0.70
        assert 0.15 < holdout_ratio < 0.25

    def test_agreement_subset_prefers_hard_cases(self) -> None:
        # Create corpus with known difficulty distribution
        easy = [_make_doc(Difficulty.EASY) for _ in range(150)]
        hard = [_make_doc(Difficulty.HARD) for _ in range(50)]
        corpus = easy + hard

        config = SplitConfig(hard_case_priority_for_agreement=True)
        splits = create_splits(corpus, config)
        agreement = splits[CorpusSplit.ANNOTATOR_AGREEMENT]

        # The agreement subset should contain hard cases
        agreement_ids = set(agreement.document_ids)
        hard_ids = {d.document_id for d in hard}
        hard_in_agreement = agreement_ids & hard_ids
        # Most of the agreement subset should be hard cases
        assert len(hard_in_agreement) > 0

    def test_deterministic_splits(self) -> None:
        corpus = _make_corpus(200)
        config = SplitConfig(random_seed=42)
        splits1 = create_splits(corpus, config)
        splits2 = create_splits(corpus, config)

        for split in CorpusSplit:
            assert splits1[split].document_ids == splits2[split].document_ids

    def test_empty_corpus_raises(self) -> None:
        with pytest.raises(ValueError, match="empty corpus"):
            create_splits([])

    def test_invalid_ratios_raises(self) -> None:
        corpus = _make_corpus(100)
        config = SplitConfig(
            train_ratio=0.5,
            calibration_ratio=0.5,
            holdout_ratio=0.5,
            agreement_ratio=0.5,
        )
        with pytest.raises(ValueError, match="sum to 1.0"):
            create_splits(corpus, config)


# ---------------------------------------------------------------------------
# Tests — select_hard_cases
# ---------------------------------------------------------------------------


class TestSelectHardCases:
    def test_returns_hard_difficulty_docs(self) -> None:
        easy = [_make_doc(Difficulty.EASY) for _ in range(50)]
        hard = [_make_doc(Difficulty.HARD) for _ in range(20)]
        corpus = easy + hard
        result = select_hard_cases(corpus)
        assert len(result) == 20
        assert all(d.difficulty == Difficulty.HARD for d in result)

    def test_respects_max_count(self) -> None:
        hard = [_make_doc(Difficulty.HARD) for _ in range(50)]
        result = select_hard_cases(hard, max_count=10)
        assert len(result) == 10

    def test_empty_if_no_hard_cases(self) -> None:
        easy = [_make_doc(Difficulty.EASY) for _ in range(50)]
        result = select_hard_cases(easy)
        assert len(result) == 0


# ---------------------------------------------------------------------------
# Tests — freeze_holdout
# ---------------------------------------------------------------------------


class TestFreezeHoldout:
    def test_produces_valid_json(self) -> None:
        manifest = SplitManifest(
            split=CorpusSplit.HOLDOUT,
            document_ids=["doc-1", "doc-2", "doc-3"],
            total_count=3,
        )
        frozen_json = freeze_holdout(manifest)
        parsed = json.loads(frozen_json)
        assert parsed["frozen"] is True
        assert parsed["total_count"] == 3
        assert len(parsed["document_id_hashes"]) == 3

    def test_hashes_are_sha256(self) -> None:
        doc_ids = ["test-doc-1", "test-doc-2"]
        manifest = SplitManifest(
            split=CorpusSplit.HOLDOUT,
            document_ids=doc_ids,
            total_count=2,
        )
        frozen_json = freeze_holdout(manifest)
        parsed = json.loads(frozen_json)

        for doc_id, stored_hash in zip(doc_ids, parsed["document_id_hashes"]):
            expected = hashlib.sha256(doc_id.encode()).hexdigest()
            assert stored_hash == expected

    def test_manifest_has_checksum(self) -> None:
        manifest = SplitManifest(
            split=CorpusSplit.HOLDOUT,
            document_ids=["doc-1"],
            total_count=1,
        )
        frozen_json = freeze_holdout(manifest)
        parsed = json.loads(frozen_json)
        assert "manifest_checksum" in parsed
        assert len(parsed["manifest_checksum"]) == 64  # SHA-256 hex

    def test_restricted_uses_set(self) -> None:
        manifest = SplitManifest(
            split=CorpusSplit.HOLDOUT,
            document_ids=["doc-1"],
            total_count=1,
        )
        frozen_json = freeze_holdout(manifest)
        parsed = json.loads(frozen_json)
        assert "prompt_tuning" in parsed["restricted_uses"]
        assert "model_training" in parsed["restricted_uses"]
        assert "hyperparameter_search" in parsed["restricted_uses"]

    def test_rejects_non_holdout_split(self) -> None:
        manifest = SplitManifest(
            split=CorpusSplit.TRAIN,
            document_ids=["doc-1"],
            total_count=1,
        )
        with pytest.raises(ValueError, match="holdout"):
            freeze_holdout(manifest)


# ---------------------------------------------------------------------------
# Tests — SplitManifest integrity
# ---------------------------------------------------------------------------


class TestSplitManifest:
    def test_verify_integrity_passes(self) -> None:
        doc_ids = ["doc-a", "doc-b", "doc-c"]
        hashes = [hashlib.sha256(d.encode()).hexdigest() for d in doc_ids]
        manifest = SplitManifest(
            split=CorpusSplit.HOLDOUT,
            document_ids=doc_ids,
            document_id_hashes=hashes,
            total_count=3,
        )
        assert manifest.verify_integrity()

    def test_verify_integrity_fails_on_tampered_hash(self) -> None:
        doc_ids = ["doc-a", "doc-b"]
        hashes = [hashlib.sha256(d.encode()).hexdigest() for d in doc_ids]
        hashes[1] = "0" * 64  # Tampered
        manifest = SplitManifest(
            split=CorpusSplit.HOLDOUT,
            document_ids=doc_ids,
            document_id_hashes=hashes,
            total_count=2,
        )
        assert not manifest.verify_integrity()

    def test_verify_integrity_fails_on_length_mismatch(self) -> None:
        manifest = SplitManifest(
            split=CorpusSplit.HOLDOUT,
            document_ids=["doc-a", "doc-b"],
            document_id_hashes=["hash-a"],
            total_count=2,
        )
        assert not manifest.verify_integrity()
