"""Create and validate data splits."""

from __future__ import annotations

import hashlib
import json
import random
from pathlib import Path
from typing import Sequence, Tuple

import pandas as pd


def random_split_indices(
    n_rows: int, seed: int
) -> Tuple[Sequence[int], Sequence[int], Sequence[int]]:
    """Return a seeded 70/15/15 row-random split."""

    indices = list(range(n_rows))
    random.Random(seed).shuffle(indices)
    n_train = int(0.70 * n_rows)
    n_validation = int(0.15 * n_rows)
    return (
        indices[:n_train],
        indices[n_train : n_train + n_validation],
        indices[n_train + n_validation :],
    )


def split_indices_sha256(
    train: Sequence[int], validation: Sequence[int], test: Sequence[int]
) -> str:
    """Hash the ordered split indices."""

    payload = json.dumps(
        {
            "train": [int(index) for index in train],
            "validation": [int(index) for index in validation],
            "test": [int(index) for index in test],
        },
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    return hashlib.sha256(payload).hexdigest()


def load_split_indices(
    path: str | Path, n_rows: int
) -> Tuple[Sequence[int], Sequence[int], Sequence[int]]:
    """Return train, validation and test row indices from ``path``."""

    table = pd.read_csv(path)
    required = {"row_index", "split"}
    missing = required.difference(table.columns)
    if missing:
        raise ValueError(f"split file is missing columns: {sorted(missing)}")

    numeric_indices = pd.to_numeric(table["row_index"], errors="raise")
    if bool((numeric_indices % 1 != 0).any()):
        raise ValueError("split row_index values must be integers")
    indices = numeric_indices.astype(int)
    if len(table) != n_rows or indices.nunique() != n_rows:
        raise ValueError("split file must contain each dataset row exactly once")
    if set(indices.tolist()) != set(range(n_rows)):
        raise ValueError("split row_index values do not match the dataset")

    labels = table["split"].astype(str)
    allowed = {"train", "validation", "test"}
    unknown = set(labels).difference(allowed)
    if unknown:
        raise ValueError(f"unknown split labels: {sorted(unknown)}")
    if "group_id" in table.columns:
        groups = table["group_id"]
        if bool(groups.isna().any()) or bool(groups.astype(str).str.strip().eq("").any()):
            raise ValueError("group_id values must be nonblank")
        leaked = table.assign(group_id=groups.astype(str), split=labels).groupby(
            "group_id"
        )["split"].nunique()
        if bool((leaked > 1).any()):
            raise ValueError("one or more group_id values cross data partitions")

    parts = {
        name: indices.loc[labels.eq(name)].tolist()
        for name in ("train", "validation", "test")
    }
    if any(not part for part in parts.values()):
        raise ValueError("train, validation and test partitions must be non-empty")
    return parts["train"], parts["validation"], parts["test"]
