"""Corpus sampling framework for the Gold Corpus.

Implements stratified sampling across document type, event class, length,
source, company count, and difficulty dimensions. Ensures diversity requirements
including duplicates, long filings, transcripts, contradictory reports, macro events,
and opposing multi-company effects.
"""

from __future__ import annotations

import random
from collections import defaultdict
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Stratification dimensions
# ---------------------------------------------------------------------------


class LengthBucket(str, Enum):
    """Document length classification."""

    SHORT = "short"  # < 2000 chars
    MEDIUM = "medium"  # 2000-8000 chars
    LONG = "long"  # > 8000 chars


class SourceType(str, Enum):
    """Document source type."""

    NEWS = "news"
    FILING = "filing"
    TRANSCRIPT = "transcript"
    PRESS_RELEASE = "press_release"
    MACRO_EVENT = "macro_event"


class CompanyCountBucket(str, Enum):
    """How many companies the document mentions."""

    SINGLE = "single"  # 1 company
    MULTI = "multi"  # 2-3 companies
    MANY = "many"  # 4+ companies


class Difficulty(str, Enum):
    """Annotation difficulty level."""

    EASY = "easy"
    MEDIUM = "medium"
    HARD = "hard"


class DiversityTag(str, Enum):
    """Tags for diversity requirements that must be represented."""

    DUPLICATE_STORY = "duplicate_story"
    LONG_FILING = "long_filing"
    TRANSCRIPT = "transcript"
    CONTRADICTORY_REPORTS = "contradictory_reports"
    MACRO_EVENT = "macro_event"
    OPPOSING_MULTI_COMPANY_EFFECTS = "opposing_multi_company_effects"


# ---------------------------------------------------------------------------
# Document metadata model
# ---------------------------------------------------------------------------


class DocumentMetadata(BaseModel):
    """Metadata for a candidate document in the sampling pool.

    This model represents the document-level attributes needed for stratified
    sampling. The actual document content is not included — only identifiers
    and classification metadata.
    """

    document_id: str = Field(description="Unique document identifier.")
    document_type: str = Field(description="article, filing, transcript, press_release, macro_event")
    event_class: str | None = Field(
        default=None, description="Primary event class if classified."
    )
    length_bucket: LengthBucket = Field(description="short/medium/long classification.")
    source_type: SourceType = Field(description="Source type classification.")
    company_count_bucket: CompanyCountBucket = Field(
        description="single/multi/many company mentions."
    )
    difficulty: Difficulty = Field(description="Annotation difficulty: easy/medium/hard.")
    diversity_tags: list[DiversityTag] = Field(
        default_factory=list,
        description="Diversity tags this document satisfies.",
    )
    extra_metadata: dict[str, Any] = Field(
        default_factory=dict,
        description="Additional metadata for filtering or reporting.",
    )


# ---------------------------------------------------------------------------
# Stratification dimensions config
# ---------------------------------------------------------------------------


class StratificationDimensions(BaseModel):
    """Configuration for stratification dimensions and their minimum counts.

    Each dimension specifies the minimum number of documents required per category
    within that dimension.
    """

    document_type: dict[str, int] = Field(
        default_factory=lambda: {
            "article": 300,
            "filing": 200,
            "transcript": 150,
            "press_release": 200,
            "macro_event": 150,
        },
        description="Minimum documents per document type.",
    )
    event_class: dict[str, int] = Field(
        default_factory=lambda: {
            "earnings_beat": 80,
            "earnings_miss": 80,
            "guidance_raise": 60,
            "guidance_cut": 60,
            "ma_announcement": 60,
            "legal_regulatory": 60,
            "product_launch": 60,
            "supply_chain": 50,
            "rating_change": 50,
            "management_change": 50,
            "macro_event": 80,
            "dividend_change": 40,
            "buyback": 40,
        },
        description="Minimum documents per event class.",
    )
    length_bucket: dict[str, int] = Field(
        default_factory=lambda: {
            "short": 250,
            "medium": 400,
            "long": 350,
        },
        description="Minimum documents per length bucket.",
    )
    source_type: dict[str, int] = Field(
        default_factory=lambda: {
            "news": 300,
            "filing": 200,
            "transcript": 150,
            "press_release": 200,
            "macro_event": 150,
        },
        description="Minimum documents per source type.",
    )
    company_count_bucket: dict[str, int] = Field(
        default_factory=lambda: {
            "single": 400,
            "multi": 350,
            "many": 250,
        },
        description="Minimum documents per company count bucket.",
    )
    difficulty: dict[str, int] = Field(
        default_factory=lambda: {
            "easy": 300,
            "medium": 400,
            "hard": 300,
        },
        description="Minimum documents per difficulty level.",
    )


# ---------------------------------------------------------------------------
# Diversity requirements
# ---------------------------------------------------------------------------


class DiversityRequirements(BaseModel):
    """Minimum counts for diversity tags that must be present in the corpus."""

    duplicate_story: int = Field(default=30, ge=1)
    long_filing: int = Field(default=50, ge=1)
    transcript: int = Field(default=50, ge=1)
    contradictory_reports: int = Field(default=30, ge=1)
    macro_event: int = Field(default=50, ge=1)
    opposing_multi_company_effects: int = Field(default=30, ge=1)

    def as_tag_minimums(self) -> dict[DiversityTag, int]:
        """Return a mapping of DiversityTag to minimum count."""
        return {
            DiversityTag.DUPLICATE_STORY: self.duplicate_story,
            DiversityTag.LONG_FILING: self.long_filing,
            DiversityTag.TRANSCRIPT: self.transcript,
            DiversityTag.CONTRADICTORY_REPORTS: self.contradictory_reports,
            DiversityTag.MACRO_EVENT: self.macro_event,
            DiversityTag.OPPOSING_MULTI_COMPANY_EFFECTS: self.opposing_multi_company_effects,
        }


# ---------------------------------------------------------------------------
# Sampling configuration
# ---------------------------------------------------------------------------


class CorpusSamplingConfig(BaseModel):
    """Complete sampling configuration for Gold Corpus construction."""

    target_size: int = Field(default=1000, ge=100, description="Target corpus size.")
    stratification: StratificationDimensions = Field(
        default_factory=StratificationDimensions
    )
    diversity: DiversityRequirements = Field(default_factory=DiversityRequirements)
    random_seed: int = Field(default=42, description="Random seed for reproducibility.")
    allow_oversampling: bool = Field(
        default=True,
        description="Allow sampling more than target_size to meet stratification minimums.",
    )


# ---------------------------------------------------------------------------
# Sampling logic
# ---------------------------------------------------------------------------


def _get_stratum_value(doc: DocumentMetadata, dimension: str) -> str | None:
    """Extract the stratum value for a document along a given dimension."""
    if dimension == "document_type":
        return doc.document_type
    elif dimension == "event_class":
        return doc.event_class
    elif dimension == "length_bucket":
        return doc.length_bucket.value
    elif dimension == "source_type":
        return doc.source_type.value
    elif dimension == "company_count_bucket":
        return doc.company_count_bucket.value
    elif dimension == "difficulty":
        return doc.difficulty.value
    return None


def sample_corpus(
    pool: list[DocumentMetadata],
    config: CorpusSamplingConfig | None = None,
) -> list[DocumentMetadata]:
    """Sample a stratified corpus from a pool of document metadata.

    The algorithm:
    1. First, ensure diversity requirements are met by selecting documents
       that carry required diversity tags.
    2. Then, fill stratification minimums dimension by dimension.
    3. Finally, if under target_size, add remaining documents proportionally.

    Args:
        pool: Available documents to sample from.
        config: Sampling configuration. Uses defaults if None.

    Returns:
        Selected documents forming the Gold Corpus sample.

    Raises:
        ValueError: If the pool cannot satisfy minimum requirements.
    """
    if config is None:
        config = CorpusSamplingConfig()

    rng = random.Random(config.random_seed)

    selected_ids: set[str] = set()
    selected: list[DocumentMetadata] = []

    def _add(doc: DocumentMetadata) -> bool:
        if doc.document_id not in selected_ids:
            selected_ids.add(doc.document_id)
            selected.append(doc)
            return True
        return False

    # Step 1: Satisfy diversity requirements
    tag_minimums = config.diversity.as_tag_minimums()
    tag_counts: dict[DiversityTag, int] = defaultdict(int)

    for tag, minimum in tag_minimums.items():
        candidates = [d for d in pool if tag in d.diversity_tags and d.document_id not in selected_ids]
        rng.shuffle(candidates)
        for doc in candidates[:minimum]:
            _add(doc)
            for t in doc.diversity_tags:
                tag_counts[t] += 1

    # Step 2: Fill stratification minimums
    dimensions = {
        "document_type": config.stratification.document_type,
        "event_class": config.stratification.event_class,
        "length_bucket": config.stratification.length_bucket,
        "source_type": config.stratification.source_type,
        "company_count_bucket": config.stratification.company_count_bucket,
        "difficulty": config.stratification.difficulty,
    }

    for dim_name, minimums in dimensions.items():
        # Count already selected for this dimension
        current_counts: dict[str, int] = defaultdict(int)
        for doc in selected:
            val = _get_stratum_value(doc, dim_name)
            if val is not None:
                current_counts[val] += 1

        for category, minimum in minimums.items():
            deficit = minimum - current_counts.get(category, 0)
            if deficit <= 0:
                continue

            candidates = [
                d
                for d in pool
                if d.document_id not in selected_ids
                and _get_stratum_value(d, dim_name) == category
            ]
            rng.shuffle(candidates)
            for doc in candidates[:deficit]:
                _add(doc)

    # Step 3: Fill up to target size if needed
    if len(selected) < config.target_size:
        remaining = [d for d in pool if d.document_id not in selected_ids]
        rng.shuffle(remaining)
        for doc in remaining[: config.target_size - len(selected)]:
            _add(doc)

    return selected


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------


class CoverageReport(BaseModel):
    """Report on corpus coverage of required categories."""

    total_documents: int
    meets_target_size: bool
    dimension_coverage: dict[str, dict[str, int]] = Field(
        default_factory=dict,
        description="Actual counts per dimension per category.",
    )
    dimension_gaps: dict[str, dict[str, int]] = Field(
        default_factory=dict,
        description="Deficit per dimension per category (0 means satisfied).",
    )
    diversity_coverage: dict[str, int] = Field(
        default_factory=dict,
        description="Actual counts per diversity tag.",
    )
    diversity_gaps: dict[str, int] = Field(
        default_factory=dict,
        description="Deficit per diversity tag.",
    )
    is_valid: bool = Field(description="Whether all requirements are met.")


def validate_corpus_coverage(
    corpus: list[DocumentMetadata],
    config: CorpusSamplingConfig | None = None,
) -> CoverageReport:
    """Check that a corpus sample meets all stratification and diversity requirements.

    Args:
        corpus: The sampled corpus.
        config: Sampling configuration to validate against.

    Returns:
        CoverageReport with detailed coverage information.
    """
    if config is None:
        config = CorpusSamplingConfig()

    # Count dimensions
    dimension_counts: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    for doc in corpus:
        dimension_counts["document_type"][doc.document_type] += 1
        if doc.event_class:
            dimension_counts["event_class"][doc.event_class] += 1
        dimension_counts["length_bucket"][doc.length_bucket.value] += 1
        dimension_counts["source_type"][doc.source_type.value] += 1
        dimension_counts["company_count_bucket"][doc.company_count_bucket.value] += 1
        dimension_counts["difficulty"][doc.difficulty.value] += 1

    # Check dimension gaps
    dimensions = {
        "document_type": config.stratification.document_type,
        "event_class": config.stratification.event_class,
        "length_bucket": config.stratification.length_bucket,
        "source_type": config.stratification.source_type,
        "company_count_bucket": config.stratification.company_count_bucket,
        "difficulty": config.stratification.difficulty,
    }

    dimension_gaps: dict[str, dict[str, int]] = {}
    all_satisfied = True

    for dim_name, minimums in dimensions.items():
        gaps: dict[str, int] = {}
        for category, minimum in minimums.items():
            actual = dimension_counts[dim_name].get(category, 0)
            deficit = max(0, minimum - actual)
            if deficit > 0:
                gaps[category] = deficit
                all_satisfied = False
        if gaps:
            dimension_gaps[dim_name] = gaps

    # Check diversity
    tag_minimums = config.diversity.as_tag_minimums()
    diversity_counts: dict[str, int] = defaultdict(int)
    for doc in corpus:
        for tag in doc.diversity_tags:
            diversity_counts[tag.value] += 1

    diversity_gaps: dict[str, int] = {}
    for tag, minimum in tag_minimums.items():
        actual = diversity_counts.get(tag.value, 0)
        deficit = max(0, minimum - actual)
        if deficit > 0:
            diversity_gaps[tag.value] = deficit
            all_satisfied = False

    meets_target = len(corpus) >= config.target_size

    return CoverageReport(
        total_documents=len(corpus),
        meets_target_size=meets_target,
        dimension_coverage=dict(dimension_counts),
        dimension_gaps=dimension_gaps,
        diversity_coverage=dict(diversity_counts),
        diversity_gaps=diversity_gaps,
        is_valid=all_satisfied and meets_target,
    )
