"""Tests for the Gold Corpus sampling framework."""

from __future__ import annotations

import uuid

import pytest

from services.intelligence_pipeline_v3.gold_corpus.sampler import (
    CompanyCountBucket,
    CorpusSamplingConfig,
    CoverageReport,
    Difficulty,
    DiversityRequirements,
    DiversityTag,
    DocumentMetadata,
    LengthBucket,
    SourceType,
    StratificationDimensions,
    sample_corpus,
    validate_corpus_coverage,
)

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


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


def _generate_diverse_pool(size: int = 3000) -> list[DocumentMetadata]:
    """Generate a large diverse pool covering all strata."""
    import random

    rng = random.Random(123)
    pool: list[DocumentMetadata] = []

    doc_types = ["article", "filing", "transcript", "press_release", "macro_event"]
    event_classes = [
        "earnings_beat", "earnings_miss", "guidance_raise", "guidance_cut",
        "ma_announcement", "legal_regulatory", "product_launch", "supply_chain",
        "rating_change", "management_change", "macro_event", "dividend_change", "buyback",
    ]
    lengths = list(LengthBucket)
    sources = list(SourceType)
    company_counts = list(CompanyCountBucket)
    difficulties = list(Difficulty)
    tags = list(DiversityTag)

    for _ in range(size):
        doc_type = rng.choice(doc_types)
        source = SourceType(doc_type) if doc_type in [s.value for s in SourceType] else rng.choice(sources)

        doc_tags: list[DiversityTag] = []
        if rng.random() < 0.15:
            doc_tags.append(rng.choice(tags))
        if doc_type == "transcript":
            doc_tags.append(DiversityTag.TRANSCRIPT)
        if doc_type == "macro_event":
            doc_tags.append(DiversityTag.MACRO_EVENT)

        pool.append(
            DocumentMetadata(
                document_id=str(uuid.uuid4()),
                document_type=doc_type,
                event_class=rng.choice(event_classes),
                length_bucket=rng.choice(lengths),
                source_type=source,
                company_count_bucket=rng.choice(company_counts),
                difficulty=rng.choice(difficulties),
                diversity_tags=doc_tags,
            )
        )

    # Ensure we have enough documents with specific diversity tags
    for tag in tags:
        for _ in range(60):
            pool.append(
                DocumentMetadata(
                    document_id=str(uuid.uuid4()),
                    document_type=rng.choice(doc_types),
                    event_class=rng.choice(event_classes),
                    length_bucket=rng.choice(lengths),
                    source_type=rng.choice(sources),
                    company_count_bucket=rng.choice(company_counts),
                    difficulty=rng.choice(difficulties),
                    diversity_tags=[tag],
                )
            )

    return pool


# ---------------------------------------------------------------------------
# Tests — CorpusSamplingConfig
# ---------------------------------------------------------------------------


class TestCorpusSamplingConfig:
    def test_default_config_valid(self) -> None:
        config = CorpusSamplingConfig()
        assert config.target_size == 1000
        assert config.random_seed == 42

    def test_custom_target_size(self) -> None:
        config = CorpusSamplingConfig(target_size=500)
        assert config.target_size == 500

    def test_minimum_target_size(self) -> None:
        with pytest.raises(Exception):
            CorpusSamplingConfig(target_size=50)


# ---------------------------------------------------------------------------
# Tests — sample_corpus
# ---------------------------------------------------------------------------


class TestSampleCorpus:
    def test_returns_at_least_target_size(self) -> None:
        pool = _generate_diverse_pool(3000)
        config = CorpusSamplingConfig(target_size=1000)
        result = sample_corpus(pool, config)
        assert len(result) >= 1000

    def test_no_duplicate_documents(self) -> None:
        pool = _generate_diverse_pool(3000)
        result = sample_corpus(pool)
        ids = [d.document_id for d in result]
        assert len(ids) == len(set(ids))

    def test_deterministic_with_same_seed(self) -> None:
        pool = _generate_diverse_pool(2000)
        config = CorpusSamplingConfig(random_seed=99)
        result1 = sample_corpus(pool, config)
        result2 = sample_corpus(pool, config)
        assert [d.document_id for d in result1] == [d.document_id for d in result2]

    def test_different_seed_gives_different_sample(self) -> None:
        pool = _generate_diverse_pool(2000)
        result1 = sample_corpus(pool, CorpusSamplingConfig(random_seed=1))
        result2 = sample_corpus(pool, CorpusSamplingConfig(random_seed=2))
        ids1 = set(d.document_id for d in result1)
        ids2 = set(d.document_id for d in result2)
        # They should differ (not guaranteed to be entirely different, but should overlap less than 100%)
        assert ids1 != ids2

    def test_diversity_tags_represented(self) -> None:
        pool = _generate_diverse_pool(3000)
        config = CorpusSamplingConfig(target_size=1000)
        result = sample_corpus(pool, config)

        # Check diversity tags are present
        all_tags: set[DiversityTag] = set()
        for doc in result:
            all_tags.update(doc.diversity_tags)

        # All diversity tag types should be represented
        for tag in DiversityTag:
            assert tag in all_tags, f"Diversity tag {tag} not represented in sample"

    def test_small_pool_returns_all(self) -> None:
        pool = [_make_doc() for _ in range(50)]
        config = CorpusSamplingConfig(target_size=100)
        result = sample_corpus(pool, config)
        # Should return everything available from the pool
        assert len(result) <= len(pool)

    def test_all_document_types_represented(self) -> None:
        pool = _generate_diverse_pool(3000)
        result = sample_corpus(pool)
        doc_types = {d.document_type for d in result}
        assert "article" in doc_types
        assert "filing" in doc_types
        assert "transcript" in doc_types
        assert "press_release" in doc_types
        assert "macro_event" in doc_types

    def test_all_difficulty_levels_represented(self) -> None:
        pool = _generate_diverse_pool(3000)
        result = sample_corpus(pool)
        difficulties = {d.difficulty for d in result}
        assert Difficulty.EASY in difficulties
        assert Difficulty.MEDIUM in difficulties
        assert Difficulty.HARD in difficulties


# ---------------------------------------------------------------------------
# Tests — validate_corpus_coverage
# ---------------------------------------------------------------------------


class TestValidateCorpusCoverage:
    def test_valid_corpus_passes(self) -> None:
        pool = _generate_diverse_pool(3000)
        config = CorpusSamplingConfig(target_size=1000)
        corpus = sample_corpus(pool, config)
        report = validate_corpus_coverage(corpus, config)
        # The report should have reasonable coverage
        assert report.total_documents >= 1000
        assert isinstance(report, CoverageReport)

    def test_empty_corpus_fails(self) -> None:
        config = CorpusSamplingConfig(target_size=100)
        report = validate_corpus_coverage([], config)
        assert not report.is_valid
        assert not report.meets_target_size

    def test_reports_dimension_gaps(self) -> None:
        # Create a corpus missing some document types
        docs = [_make_doc(document_type="article") for _ in range(100)]
        config = CorpusSamplingConfig(target_size=100)
        report = validate_corpus_coverage(docs, config)
        # Should report gaps for missing document types
        assert "document_type" in report.dimension_gaps
        assert "filing" in report.dimension_gaps["document_type"]

    def test_reports_diversity_gaps(self) -> None:
        # Create a corpus without diversity tags
        docs = [_make_doc() for _ in range(100)]
        config = CorpusSamplingConfig(target_size=100)
        report = validate_corpus_coverage(docs, config)
        assert report.diversity_gaps  # Should have gaps


# ---------------------------------------------------------------------------
# Tests — StratificationDimensions
# ---------------------------------------------------------------------------


class TestStratificationDimensions:
    def test_default_dimensions_cover_all_types(self) -> None:
        dims = StratificationDimensions()
        assert "article" in dims.document_type
        assert "filing" in dims.document_type
        assert "short" in dims.length_bucket
        assert "medium" in dims.length_bucket
        assert "long" in dims.length_bucket

    def test_all_event_classes_have_minimums(self) -> None:
        dims = StratificationDimensions()
        assert len(dims.event_class) == 13  # All EventClass values


# ---------------------------------------------------------------------------
# Tests — DiversityRequirements
# ---------------------------------------------------------------------------


class TestDiversityRequirements:
    def test_default_requirements(self) -> None:
        req = DiversityRequirements()
        assert req.duplicate_story >= 1
        assert req.long_filing >= 1
        assert req.contradictory_reports >= 1

    def test_as_tag_minimums(self) -> None:
        req = DiversityRequirements()
        tag_mins = req.as_tag_minimums()
        assert DiversityTag.DUPLICATE_STORY in tag_mins
        assert DiversityTag.MACRO_EVENT in tag_mins
        assert all(v > 0 for v in tag_mins.values())
