from __future__ import annotations

import hashlib
import random
from collections import defaultdict
from datetime import datetime
from typing import Any, Callable

from pydantic import BaseModel, ConfigDict, Field


class SamplingManifestV1(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    schema_version: str = "1.0"
    sample_id: str = Field(min_length=1)
    strategy: str = Field(min_length=1)
    population_hash: str = Field(min_length=1)
    random_seed: str = Field(min_length=1)
    selected_record_ids: tuple[str, ...]
    strata_counts: dict[str, int]


def population_hash(record_ids: list[str]) -> str:
    return hashlib.sha256("\n".join(sorted(record_ids)).encode("utf-8")).hexdigest()


def reproducible_random_sample(record_ids: list[str], size: int, seed: str) -> SamplingManifestV1:
    unique = sorted(set(record_ids))
    if size < 1 or size > len(unique):
        raise ValueError("Sample size must be within the unique population size.")
    selected = sorted(random.Random(seed).sample(unique, size))
    return SamplingManifestV1(
        sample_id=hashlib.sha256(f"random:{seed}:{population_hash(unique)}:{size}".encode("utf-8")).hexdigest(),
        strategy="random",
        population_hash=population_hash(unique),
        random_seed=seed,
        selected_record_ids=tuple(selected),
        strata_counts={"all": len(selected)},
    )


def stratified_sample(
    records: list[dict[str, Any]],
    size: int,
    seed: str,
    stratum: Callable[[dict[str, Any]], str],
) -> SamplingManifestV1:
    groups: dict[str, list[str]] = defaultdict(list)
    for record in records:
        groups[stratum(record)].append(record["record_id"])
    if size < len(groups):
        raise ValueError("Sample size must include at least one record per stratum.")
    population = sorted({record_id for values in groups.values() for record_id in values})
    if size > len(population):
        raise ValueError("Sample size exceeds population.")
    rng = random.Random(seed)
    selected: list[str] = []
    counts: dict[str, int] = {}
    remaining = size
    strata = sorted(groups)
    for index, key in enumerate(strata):
        available = sorted(set(groups[key]))
        target = max(1, round(size * len(available) / len(population)))
        target = min(target, len(available), remaining - (len(strata) - index - 1))
        picked = rng.sample(available, target)
        selected.extend(picked)
        counts[key] = len(picked)
        remaining -= len(picked)
    if remaining:
        candidates = sorted(set(population).difference(selected))
        selected.extend(rng.sample(candidates, remaining))
    return SamplingManifestV1(
        sample_id=hashlib.sha256(f"stratified:{seed}:{population_hash(population)}:{size}".encode("utf-8")).hexdigest(),
        strategy="stratified",
        population_hash=population_hash(population),
        random_seed=seed,
        selected_record_ids=tuple(sorted(selected)),
        strata_counts=counts,
    )


def capacity_report(review_events: list[dict[str, Any]]) -> dict[str, Any]:
    finalized = [event for event in review_events if event.get("outcome") not in {"skip", None}]
    hours = sum(float(event.get("duration_seconds", 0)) for event in finalized) / 3600
    disagreements = sum(bool(event.get("disagreed")) for event in finalized)
    return {
        "finalized_reviews": len(finalized),
        "review_hours": round(hours, 2),
        "reviews_per_hour": round(len(finalized) / hours, 2) if hours else None,
        "disagreement_rate": round(disagreements / len(finalized), 4) if finalized else None,
        "generated_at": datetime.utcnow().isoformat() + "Z",
    }