"""Dataset split management for the Gold Corpus.

Implements train/calibration/holdout/agreement splits with:
- Configurable ratios (default: 60/15/20/5)
- Frozen holdout that cannot be used for prompt or model tuning
- Hard-case subset selection for double annotation
- Immutable manifest generation with SHA-256 hashes
"""

from __future__ import annotations

import hashlib
import json
import random
from datetime import datetime, timezone
from enum import Enum

from pydantic import BaseModel, Field

from services.intelligence_pipeline_v3.gold_corpus.sampler import (
    Difficulty,
    DocumentMetadata,
)

# ---------------------------------------------------------------------------
# Split enum and config
# ---------------------------------------------------------------------------


class CorpusSplit(str, Enum):
    """Corpus split identifiers."""

    TRAIN = "train"
    CALIBRATION = "calibration"
    HOLDOUT = "holdout"
    ANNOTATOR_AGREEMENT = "annotator_agreement"


class SplitConfig(BaseModel):
    """Configuration for corpus split ratios.

    Ratios must sum to 1.0. The holdout split is frozen and restricted
    from any use in prompt engineering or model tuning.
    """

    train_ratio: float = Field(default=0.60, ge=0.0, le=1.0)
    calibration_ratio: float = Field(default=0.15, ge=0.0, le=1.0)
    holdout_ratio: float = Field(default=0.20, ge=0.0, le=1.0)
    agreement_ratio: float = Field(default=0.05, ge=0.0, le=1.0)
    random_seed: int = Field(default=42)
    hard_case_priority_for_agreement: bool = Field(
        default=True,
        description="Prioritize hard cases for the annotator agreement subset.",
    )

    def validate_ratios(self) -> bool:
        """Check that split ratios sum to 1.0 (within tolerance)."""
        total = (
            self.train_ratio
            + self.calibration_ratio
            + self.holdout_ratio
            + self.agreement_ratio
        )
        return abs(total - 1.0) < 0.001


# ---------------------------------------------------------------------------
# Split manifest
# ---------------------------------------------------------------------------


class SplitManifest(BaseModel):
    """Immutable manifest recording which documents belong to which split.

    The holdout manifest includes SHA-256 hashes of document IDs to prevent
    accidental use in tuning workflows.
    """

    split: CorpusSplit
    document_ids: list[str]
    document_id_hashes: list[str] = Field(
        default_factory=list,
        description="SHA-256 hashes of document IDs for integrity verification.",
    )
    frozen: bool = Field(default=False, description="Whether this split is frozen (holdout).")
    frozen_at: datetime | None = Field(default=None)
    restricted_uses: list[str] = Field(
        default_factory=list,
        description="Uses this split is restricted from (e.g., prompt_tuning, model_training).",
    )
    total_count: int = Field(default=0)

    def verify_integrity(self) -> bool:
        """Verify that document_id_hashes match the document_ids."""
        if len(self.document_ids) != len(self.document_id_hashes):
            return False
        for doc_id, expected_hash in zip(self.document_ids, self.document_id_hashes):
            computed = hashlib.sha256(doc_id.encode()).hexdigest()
            if computed != expected_hash:
                return False
        return True


# ---------------------------------------------------------------------------
# Hard-case selection
# ---------------------------------------------------------------------------


def select_hard_cases(
    corpus: list[DocumentMetadata],
    max_count: int | None = None,
) -> list[DocumentMetadata]:
    """Select hard-case documents suitable for double annotation.

    Hard cases are documents with difficulty='hard' or those with
    multiple ambiguity-inducing characteristics (multi-company,
    long filings, contradictory content).

    Args:
        corpus: Full corpus to select from.
        max_count: Maximum number of hard cases to return.

    Returns:
        List of hard-case documents.
    """
    hard_cases = [
        doc
        for doc in corpus
        if doc.difficulty == Difficulty.HARD
    ]

    if max_count is not None and len(hard_cases) > max_count:
        hard_cases = hard_cases[:max_count]

    return hard_cases


# ---------------------------------------------------------------------------
# Split creation
# ---------------------------------------------------------------------------


def create_splits(
    corpus: list[DocumentMetadata],
    config: SplitConfig | None = None,
) -> dict[CorpusSplit, SplitManifest]:
    """Create stratified dataset splits from the corpus.

    The agreement subset prioritizes hard cases when configured. The holdout
    split is marked as frozen and restricted from prompt/model tuning.

    Args:
        corpus: The complete Gold Corpus sample.
        config: Split configuration. Uses defaults if None.

    Returns:
        Dictionary mapping each CorpusSplit to its SplitManifest.

    Raises:
        ValueError: If config ratios don't sum to 1.0 or corpus is empty.
    """
    if config is None:
        config = SplitConfig()

    if not config.validate_ratios():
        raise ValueError(
            f"Split ratios must sum to 1.0, got "
            f"{config.train_ratio + config.calibration_ratio + config.holdout_ratio + config.agreement_ratio:.3f}"
        )

    if not corpus:
        raise ValueError("Cannot create splits from an empty corpus.")

    rng = random.Random(config.random_seed)

    # Separate hard cases for agreement subset priority
    agreement_docs: list[DocumentMetadata] = []
    remaining_docs: list[DocumentMetadata] = list(corpus)

    agreement_count = max(1, int(len(corpus) * config.agreement_ratio))

    if config.hard_case_priority_for_agreement:
        hard_cases = select_hard_cases(remaining_docs)
        rng.shuffle(hard_cases)
        agreement_docs = hard_cases[:agreement_count]
        remaining_ids = {d.document_id for d in agreement_docs}
        remaining_docs = [d for d in remaining_docs if d.document_id not in remaining_ids]

        # Fill remaining agreement slots if hard cases weren't enough
        if len(agreement_docs) < agreement_count:
            rng.shuffle(remaining_docs)
            extra_needed = agreement_count - len(agreement_docs)
            agreement_docs.extend(remaining_docs[:extra_needed])
            remaining_docs = remaining_docs[extra_needed:]
    else:
        rng.shuffle(remaining_docs)
        agreement_docs = remaining_docs[:agreement_count]
        remaining_docs = remaining_docs[agreement_count:]

    # Distribute remaining docs into train/calibration/holdout
    rng.shuffle(remaining_docs)
    remaining_total = len(remaining_docs)

    # Calculate proportional sizes for remaining splits (exclude agreement ratio)
    remaining_ratio = config.train_ratio + config.calibration_ratio + config.holdout_ratio
    train_count = int(remaining_total * (config.train_ratio / remaining_ratio))
    calibration_count = int(remaining_total * (config.calibration_ratio / remaining_ratio))
    # Holdout gets the remainder to avoid rounding losses

    train_docs = remaining_docs[:train_count]
    calibration_docs = remaining_docs[train_count : train_count + calibration_count]
    holdout_docs = remaining_docs[train_count + calibration_count :]

    # Build manifests
    def _build_manifest(
        split: CorpusSplit,
        docs: list[DocumentMetadata],
        frozen: bool = False,
    ) -> SplitManifest:
        doc_ids = [d.document_id for d in docs]
        doc_hashes = [hashlib.sha256(did.encode()).hexdigest() for did in doc_ids]
        restricted = (
            ["prompt_tuning", "model_training", "hyperparameter_search"]
            if frozen
            else []
        )
        return SplitManifest(
            split=split,
            document_ids=doc_ids,
            document_id_hashes=doc_hashes,
            frozen=frozen,
            frozen_at=datetime.now(timezone.utc) if frozen else None,
            restricted_uses=restricted,
            total_count=len(doc_ids),
        )

    return {
        CorpusSplit.TRAIN: _build_manifest(CorpusSplit.TRAIN, train_docs),
        CorpusSplit.CALIBRATION: _build_manifest(CorpusSplit.CALIBRATION, calibration_docs),
        CorpusSplit.HOLDOUT: _build_manifest(CorpusSplit.HOLDOUT, holdout_docs, frozen=True),
        CorpusSplit.ANNOTATOR_AGREEMENT: _build_manifest(
            CorpusSplit.ANNOTATOR_AGREEMENT, agreement_docs
        ),
    }


# ---------------------------------------------------------------------------
# Holdout freezing
# ---------------------------------------------------------------------------


def freeze_holdout(manifest: SplitManifest) -> str:
    """Create an immutable holdout manifest as a JSON string with SHA-256 hashes.

    The frozen manifest serves as a contract: documents in the holdout split
    MUST NOT be used for prompt engineering, model fine-tuning, or
    hyperparameter optimization.

    Args:
        manifest: The holdout split manifest.

    Returns:
        JSON string of the frozen manifest with integrity hashes.

    Raises:
        ValueError: If the manifest is not the holdout split.
    """
    if manifest.split != CorpusSplit.HOLDOUT:
        raise ValueError(
            f"Can only freeze holdout manifests, got split={manifest.split.value}"
        )

    # Ensure hashes are computed
    if not manifest.document_id_hashes:
        manifest.document_id_hashes = [
            hashlib.sha256(did.encode()).hexdigest()
            for did in manifest.document_ids
        ]

    # Mark as frozen
    manifest.frozen = True
    manifest.frozen_at = datetime.now(timezone.utc)
    manifest.restricted_uses = [
        "prompt_tuning",
        "model_training",
        "hyperparameter_search",
    ]

    # Create the immutable JSON document
    frozen_doc = {
        "split": manifest.split.value,
        "frozen": True,
        "frozen_at": manifest.frozen_at.isoformat(),
        "restricted_uses": manifest.restricted_uses,
        "total_count": manifest.total_count,
        "document_ids": manifest.document_ids,
        "document_id_hashes": manifest.document_id_hashes,
        "manifest_checksum": "",
    }

    # Compute manifest-level checksum (excluding the checksum field itself)
    content_for_hash = json.dumps(
        {k: v for k, v in frozen_doc.items() if k != "manifest_checksum"},
        sort_keys=True,
    )
    frozen_doc["manifest_checksum"] = hashlib.sha256(
        content_for_hash.encode()
    ).hexdigest()

    return json.dumps(frozen_doc, indent=2, default=str)
