from __future__ import annotations

import random
from dataclasses import dataclass
from typing import Iterable, Iterator, Sequence

from dovla_cil.data.datasets import CILDataset
from dovla_cil.data.schema import CILRecord
from dovla_cil.data.sharding import group_records


class BatchIndices(list[int]):
    """Batch index list with same-state ranking metadata."""

    def __init__(
        self,
        indices: Iterable[int],
        *,
        group_ids: Iterable[str] = (),
        pair_indices: Iterable[tuple[int, int]] = (),
    ) -> None:
        super().__init__(int(index) for index in indices)
        self.group_ids = list(group_ids)
        self.pair_indices = [tuple(pair) for pair in pair_indices]


class GroupAwareBatchSampler:
    """Yield batches that preserve CIL same-state group structure.

    Modes:
    - `full_group`: include complete groups, up to `batch_groups` groups per batch.
    - `pairs`: sample positive/negative same-group pairs for ranking losses.
    - `mixed`: sample `records_per_group` records per group while preferring reward-diverse sets.
    """

    def __init__(
        self,
        dataset: CILDataset,
        *,
        mode: str = "full_group",
        batch_groups: int = 1,
        records_per_group: int | None = None,
        pair_count_per_group: int = 1,
        shuffle: bool = False,
        seed: int = 17,
        reward_margin: float = 1e-6,
        group_ids: Sequence[str] | None = None,
    ) -> None:
        if mode not in {"full_group", "pairs", "mixed"}:
            raise ValueError("mode must be 'full_group', 'pairs', or 'mixed'")
        if batch_groups <= 0:
            raise ValueError("batch_groups must be positive")
        if records_per_group is not None and records_per_group <= 0:
            raise ValueError("records_per_group must be positive when provided")
        if pair_count_per_group <= 0:
            raise ValueError("pair_count_per_group must be positive")
        if reward_margin < 0:
            raise ValueError("reward_margin must be non-negative")
        self.dataset = dataset
        self.mode = mode
        self.batch_groups = int(batch_groups)
        self.records_per_group = records_per_group
        self.pair_count_per_group = int(pair_count_per_group)
        self.shuffle = bool(shuffle)
        self.seed = int(seed)
        self.reward_margin = float(reward_margin)
        unknown = [group_id for group_id in (group_ids or []) if group_id not in dataset.group_ids]
        if unknown:
            raise KeyError(f"Unknown CIL groups for sampler: {unknown[:3]}")
        self.group_ids = list(group_ids) if group_ids is not None else list(dataset.group_ids)

    def __iter__(self) -> Iterator[BatchIndices]:
        rng = random.Random(self.seed)
        group_ids = list(self.group_ids)
        if self.shuffle:
            rng.shuffle(group_ids)
        for offset in range(0, len(group_ids), self.batch_groups):
            selected_group_ids = group_ids[offset : offset + self.batch_groups]
            if self.mode == "full_group":
                yield self._full_group_batch(selected_group_ids)
            elif self.mode == "pairs":
                yield self._pair_batch(selected_group_ids, rng)
            else:
                yield self._mixed_batch(selected_group_ids, rng)

    def __len__(self) -> int:
        if not self.group_ids:
            return 0
        return (len(self.group_ids) + self.batch_groups - 1) // self.batch_groups

    def _full_group_batch(self, group_ids: Sequence[str]) -> BatchIndices:
        indices: list[int] = []
        for group_id in group_ids:
            indices.extend(self.dataset.group_indices(group_id))
        pair_indices = _local_pair_indices(
            [self.dataset[index] for index in indices], reward_margin=self.reward_margin
        )
        return BatchIndices(indices, group_ids=group_ids, pair_indices=pair_indices)

    def _pair_batch(self, group_ids: Sequence[str], rng: random.Random) -> BatchIndices:
        indices: list[int] = []
        pair_indices: list[tuple[int, int]] = []
        for group_id in group_ids:
            group_indices = self.dataset.group_indices(group_id)
            records = [self.dataset[index] for index in group_indices]
            global_pairs = _sample_reward_ordered_pairs(
                group_indices,
                records,
                pair_count=self.pair_count_per_group,
                reward_margin=self.reward_margin,
                rng=rng,
            )
            local_by_global: dict[int, int] = {}
            for better, worse in global_pairs:
                if better not in local_by_global:
                    local_by_global[better] = len(indices)
                    indices.append(better)
                if worse not in local_by_global:
                    local_by_global[worse] = len(indices)
                    indices.append(worse)
                pair_indices.append((local_by_global[better], local_by_global[worse]))
        return BatchIndices(indices, group_ids=group_ids, pair_indices=pair_indices)

    def _mixed_batch(self, group_ids: Sequence[str], rng: random.Random) -> BatchIndices:
        indices: list[int] = []
        for group_id in group_ids:
            group_indices = self.dataset.group_indices(group_id)
            records = [self.dataset[index] for index in group_indices]
            selected = _sample_records_from_group(
                group_indices,
                records,
                count=self.records_per_group or len(group_indices),
                reward_margin=self.reward_margin,
                rng=rng,
            )
            indices.extend(selected)
        pair_indices = _local_pair_indices(
            [self.dataset[index] for index in indices], reward_margin=self.reward_margin
        )
        return BatchIndices(indices, group_ids=group_ids, pair_indices=pair_indices)


@dataclass(frozen=True)
class GroupBatchSampler:
    """Backward-compatible complete-group sampler over materialized records."""

    records: tuple[CILRecord, ...]
    shuffle: bool = False
    seed: int = 17

    @classmethod
    def from_records(
        cls, records: Iterable[CILRecord], *, shuffle: bool = False, seed: int = 17
    ) -> "GroupBatchSampler":
        return cls(records=tuple(records), shuffle=shuffle, seed=seed)

    def __iter__(self) -> Iterator[list[CILRecord]]:
        groups = list(group_records(self.records).values())
        if self.shuffle:
            rng = random.Random(self.seed)
            rng.shuffle(groups)
        yield from groups


def _sample_reward_ordered_pairs(
    group_indices: list[int],
    records: list[CILRecord],
    *,
    pair_count: int,
    reward_margin: float,
    rng: random.Random,
) -> list[tuple[int, int]]:
    candidates: list[tuple[int, int]] = []
    for left in range(len(records)):
        for right in range(len(records)):
            if left == right:
                continue
            reward_left = _ranking_reward(records[left])
            reward_right = _ranking_reward(records[right])
            if reward_left > reward_right + reward_margin:
                candidates.append((group_indices[left], group_indices[right]))
    rng.shuffle(candidates)
    return candidates[:pair_count]


def _sample_records_from_group(
    group_indices: list[int],
    records: list[CILRecord],
    *,
    count: int,
    reward_margin: float,
    rng: random.Random,
) -> list[int]:
    if count >= len(group_indices):
        return list(group_indices)
    pairs = _sample_reward_ordered_pairs(
        group_indices,
        records,
        pair_count=max(1, count // 2),
        reward_margin=reward_margin,
        rng=rng,
    )
    selected: list[int] = []
    for better, worse in pairs:
        for index in (better, worse):
            if index not in selected:
                selected.append(index)
            if len(selected) >= count:
                return selected
    remaining = [index for index in group_indices if index not in selected]
    rng.shuffle(remaining)
    return selected + remaining[: max(0, count - len(selected))]


def _local_pair_indices(records: list[CILRecord], *, reward_margin: float) -> list[tuple[int, int]]:
    pairs: list[tuple[int, int]] = []
    by_group = group_records(records)
    offset_by_record_id = {record.record_id: index for index, record in enumerate(records)}
    for group in by_group.values():
        for left in range(len(group)):
            for right in range(len(group)):
                if left == right:
                    continue
                if _ranking_reward(group[left]) > _ranking_reward(group[right]) + reward_margin:
                    pairs.append(
                        (
                            offset_by_record_id[group[left].record_id],
                            offset_by_record_id[group[right].record_id],
                        )
                    )
    return pairs


def _ranking_reward(record: CILRecord) -> float:
    return record.reward.score