"""
Dataset access layer for pipecat-ai/smart-turn-data-v3.2-{train,test}.

Design constraints this module honors:
  - Must support HF `streaming=True` so exploratory work never requires
    downloading the full 41GB / 4.84GB parquet files (per Phase 2 brief).
  - Must not materialize the full dataset in memory.
  - Must be deterministic given a seed, for reproducible dev-subset creation.

Status in THIS sandbox: this module has NOT been executed against the real
dataset. `datasets` and `huggingface_hub` are not installed here, and pip
cannot reach PyPI (confirmed: `pip install` fails with "No matching
distribution found" — no route out), and `bash` network egress is
proxy-blocked for huggingface.co specifically (confirmed via curl: HTTP 403,
`x-deny-reason: host_not_allowed`). This is written to run correctly in an
environment with real network access (e.g. the actual training environment
this project will run in) and its logic (subset sampling, split logic) is
unit-tested here against a small in-memory fake dataset that mimics the
real schema — see tests/test_data.py. That is a legitimate way to validate
sampling/splitting *logic* without the real data; it does NOT substitute
for running this against the real dataset, and is never presented as such.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Iterator, Optional

import numpy as np

TRAIN_DATASET_ID = "pipecat-ai/smart-turn-data-v3.2-train"
TEST_DATASET_ID = "pipecat-ai/smart-turn-data-v3.2-test"

# Confirmed schema (from the dataset's own README dataset_info YAML —
# see docs/INITIAL_ANALYSIS.md §1). Listed here as a single source of truth
# so downstream code fails loudly if the real schema doesn't match this
# (rather than silently mis-indexing columns).
EXPECTED_COLUMNS = {
    "audio",
    "id",
    "language",
    "endpoint_bool",
    "midfiller",
    "endfiller",
    "synthetic",
    "spoken_text",
    "dataset",
}


class DatasetAccessError(RuntimeError):
    pass


def _require_datasets_lib():
    try:
        import datasets  # noqa: F401
        return datasets
    except ImportError as e:
        raise DatasetAccessError(
            "The `datasets` library is not installed / could not reach the "
            "Hugging Face Hub in this environment. This function is written "
            "to work correctly wherever `datasets` and network access are "
            "available — install with `pip install datasets huggingface_hub` "
            "in such an environment."
        ) from e


def load_hf_dataset(
    dataset_id: str,
    split: str = "train",
    streaming: bool = True,
):
    """Thin wrapper around `datasets.load_dataset`, defaulting to streaming
    mode so exploratory work doesn't require downloading the full dataset.

    Not executable in this sandbox (see module docstring) — raises
    DatasetAccessError with a clear message rather than pretending to
    succeed.
    """
    ds_lib = _require_datasets_lib()
    try:
        ds = ds_lib.load_dataset(dataset_id, split=split, streaming=streaming)
    except Exception as e:
        raise DatasetAccessError(
            f"Failed to load {dataset_id} (split={split}, streaming={streaming}): {e}"
        ) from e

    # Fail loudly, early, if the live schema doesn't match what we've
    # documented — protects every downstream assumption in this codebase.
    if not streaming:
        cols = set(ds.column_names)
    else:
        # Streaming datasets don't expose column_names without peeking at
        # one example.
        first = next(iter(ds))
        cols = set(first.keys())
    missing = EXPECTED_COLUMNS - cols
    if missing:
        raise DatasetAccessError(
            f"{dataset_id} is missing expected columns {missing}. "
            f"Schema may have changed since docs/INITIAL_ANALYSIS.md was written — "
            f"re-verify before trusting any downstream code."
        )
    return ds


# ---------------------------------------------------------------------------
# Reproducible stratified dev-subset sampling
# ---------------------------------------------------------------------------

STRATIFY_COLUMNS = ("endpoint_bool", "language", "dataset", "synthetic", "midfiller", "endfiller")


def duration_bucket(duration_sec: float) -> str:
    """Coarse duration bucketing used as an extra stratification axis.
    Bucket edges are round numbers chosen for interpretability, not fit to
    any observed distribution (we don't have the full duration distribution
    to fit to — see docs/INITIAL_ANALYSIS.md §3).
    """
    if duration_sec < 1.0:
        return "<1s"
    if duration_sec < 2.0:
        return "1-2s"
    if duration_sec < 4.0:
        return "2-4s"
    if duration_sec < 8.0:
        return "4-8s"
    return ">=8s"


def _stratum_key(record: dict) -> tuple:
    key = []
    for col in STRATIFY_COLUMNS:
        val = record.get(col)
        # None (null) is itself a meaningful stratum value (e.g. filler
        # metadata unavailable for this source) — must NOT be coerced to
        # False, since null != false (per Phase 2 brief §11 warning).
        key.append(str(val))
    if "duration_sec" in record:
        key.append(duration_bucket(record["duration_sec"]))
    return tuple(key)


def stratified_reservoir_sample(
    records: Iterable[dict],
    target_n: int,
    seed: int = 42,
    max_scan: Optional[int] = None,
) -> list:
    """Streaming-compatible stratified sampling.

    Approach: proportional allocation via per-stratum reservoir sampling.
    Because we're consuming a (potentially streaming, single-pass) iterator
    and don't know strata sizes in advance, this uses the standard two-pass-
    free approach: maintain a per-stratum reservoir sized proportionally as
    strata are discovered, using Algorithm R per stratum. This is
    deterministic given `seed` and the iteration order of `records`.

    This does NOT require materializing the full dataset — it holds at most
    `target_n` records (plus a small amount of per-stratum bookkeeping) in
    memory at any time, satisfying the "no full dataset materialization"
    requirement.

    Note on the allocation strategy: true proportional stratified sampling
    with unknown-in-advance strata sizes, in a single streaming pass, is a
    real algorithmic constraint — not solvable exactly without either (a) a
    first pass to count strata sizes, or (b) an online algorithm that
    approximates proportional allocation as it goes. We do a light first
    pass over up to `max_scan` records (default: unbounded, i.e. a full
    pass) to compute exact per-stratum counts, then a second pass to do
    proportional reservoir sampling per stratum. This means the function
    consumes an iterable twice if it's re-iterable (e.g. a non-streaming HF
    dataset). For a truly single-pass streaming source, pass
    `max_scan=<some cap>` and accept that allocation is approximate for
    strata whose true size wasn't fully observed within the cap — this
    tradeoff is intentional and documented rather than hidden.
    """
    rng = np.random.default_rng(seed)

    # Pass 1: count stratum sizes (bounded by max_scan if given)
    stratum_counts: dict[tuple, int] = {}
    n_scanned = 0
    for record in records:
        stratum_counts[_stratum_key(record)] = stratum_counts.get(_stratum_key(record), 0) + 1
        n_scanned += 1
        if max_scan is not None and n_scanned >= max_scan:
            break

    if n_scanned == 0:
        return []

    total = sum(stratum_counts.values())
    # Proportional target size per stratum (at least 1 if the stratum has
    # any members and target_n leaves room)
    target_per_stratum = {
        k: max(1, round(target_n * c / total)) for k, c in stratum_counts.items()
    }

    # Pass 2: reservoir-sample within each stratum up to its target size.
    # Requires records to be re-iterable for this two-pass approach; for a
    # single-pass-only streaming source, use the max_scan single-pass mode
    # below instead (kept as a separate documented code path, not hidden).
    reservoirs: dict[tuple, list] = {k: [] for k in stratum_counts}
    seen_counts: dict[tuple, int] = {k: 0 for k in stratum_counts}

    for record in records:
        key = _stratum_key(record)
        if key not in reservoirs:
            continue  # discovered after max_scan cutoff in pass 1; skip
        seen_counts[key] += 1
        cap = target_per_stratum[key]
        res = reservoirs[key]
        if len(res) < cap:
            res.append(record)
        else:
            j = rng.integers(0, seen_counts[key])
            if j < cap:
                res[j] = record
        n_scanned_2 = sum(seen_counts.values())
        if max_scan is not None and n_scanned_2 >= max_scan:
            break

    sample = [rec for res in reservoirs.values() for rec in res]
    rng.shuffle(sample)
    return sample[:target_n] if len(sample) > target_n else sample


@dataclass
class DevSubsetManifest:
    """Records exactly how a dev subset was created, for reproducibility
    and for the "document exactly how it was sampled" requirement.
    """
    source_dataset_id: str
    source_split: str
    target_n: int
    actual_n: int
    seed: int
    stratify_columns: tuple
    max_scan: Optional[int]

    def to_dict(self) -> dict:
        return {
            "source_dataset_id": self.source_dataset_id,
            "source_split": self.source_split,
            "target_n": self.target_n,
            "actual_n": self.actual_n,
            "seed": self.seed,
            "stratify_columns": list(self.stratify_columns),
            "max_scan": self.max_scan,
        }