from __future__ import annotations

import json
import math
import threading
import uuid
from abc import ABC, abstractmethod
from collections import Counter
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, fields, is_dataclass
from numbers import Integral, Real
from typing import Any, ClassVar

from models import BatchDiagnosticRecord, RunContext
from storage import LocalSessionStore, utc_now


def _serialize_dataclass(instance: Any) -> dict[str, Any]:
    payload: dict[str, Any] = {}
    for item in fields(instance):
        payload[item.name] = _serialize_value(getattr(instance, item.name))
    return payload


def _serialize_value(value: Any) -> Any:
    if is_dataclass(value):
        return _serialize_dataclass(value)
    if isinstance(value, tuple):
        return [_serialize_value(item) for item in value]
    if isinstance(value, list):
        return [_serialize_value(item) for item in value]
    if isinstance(value, dict):
        return {key: _serialize_value(item) for key, item in value.items()}
    return value


def _validate_positive_int(name: str, value: Any) -> int:
    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer.")
    normalized_value = int(value)
    if normalized_value <= 0:
        raise ValueError(f"{name} must be a positive integer.")
    return normalized_value


def _validate_non_negative_int(name: str, value: Any) -> int:
    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer.")
    normalized_value = int(value)
    if normalized_value < 0:
        raise ValueError(f"{name} must not be negative.")
    return normalized_value


def _validate_positive_float(name: str, value: Any) -> float:
    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real number.")
    normalized_value = float(value)
    if normalized_value <= 0.0:
        raise ValueError(f"{name} must be greater than zero.")
    return normalized_value


def _validate_probability(name: str, value: Any) -> float:
    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real number.")
    normalized_value = float(value)
    if normalized_value <= 0.0 or normalized_value >= 1.0:
        raise ValueError(f"{name} must be between 0.0 and 1.0.")
    return normalized_value


def _normalize_optional_key(name: str, value: str | None) -> str | None:
    if value is None:
        return None
    candidate = value.strip()
    if not candidate:
        raise ValueError(f"{name} must not be empty when provided.")
    return candidate


def _normalize_index(name: str, value: int | None) -> int | None:
    if value is None:
        return None
    return _validate_non_negative_int(name, value)


def _normalize_loss(value: float) -> float:
    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError("loss must be a numeric scalar.")
    normalized_value = float(value)
    if not math.isfinite(normalized_value):
        raise ValueError("loss must be a finite real number.")
    return normalized_value


def _normalize_metadata(metadata: Mapping[str, Any] | None) -> dict[str, Any] | None:
    if metadata is None:
        return None
    normalized_payload = dict(metadata)
    json.dumps(normalized_payload, sort_keys=True)
    return normalized_payload


def _flatten_values(value: Any) -> tuple[Any, ...]:
    if value is None:
        return ()
    if isinstance(value, Mapping):
        return (value,)
    if isinstance(value, (str, bytes)):
        return (value,)
    tolist = getattr(value, "tolist", None)
    if callable(tolist):
        try:
            return _flatten_values(tolist())
        except Exception:
            pass
    if isinstance(value, Sequence):
        flattened: list[Any] = []
        for item in value:
            flattened.extend(_flatten_values(item))
        return tuple(flattened)
    if isinstance(value, set):
        flattened = []
        for item in value:
            flattened.extend(_flatten_values(item))
        return tuple(flattened)
    return (value,)


def _normalize_label(value: Any) -> str:
    candidate = str(value).strip()
    if not candidate:
        raise ValueError("labels must not contain empty values.")
    return candidate


def _normalize_single_sample_id(value: Any) -> str:
    candidate = str(value).strip()
    if not candidate:
        raise ValueError("sample IDs must not be empty.")
    return candidate


def _normalize_sample_ids(sample_ids: Any) -> tuple[str, ...]:
    if sample_ids is None:
        return ()
    normalized_values = []
    for raw_value in _flatten_values(sample_ids):
        if raw_value is None:
            raise ValueError("sample IDs must not contain None.")
        normalized_values.append(_normalize_single_sample_id(raw_value))
    return tuple(normalized_values)


def _normalize_class_counts(labels: Any) -> dict[str, int]:
    if labels is None:
        return {}
    if isinstance(labels, Mapping):
        normalized_counts: dict[str, int] = {}
        for raw_label, raw_count in labels.items():
            label = _normalize_label(raw_label)
            count = _validate_positive_int(f"class count for '{label}'", raw_count)
            normalized_counts[label] = normalized_counts.get(label, 0) + count
        return normalized_counts

    counts: Counter[str] = Counter()
    for raw_label in _flatten_values(labels):
        if raw_label is None:
            raise ValueError("labels must not contain None.")
        counts[_normalize_label(raw_label)] += 1
    return dict(counts)


def _normalize_sequence_lengths(sequence_lengths: Any) -> tuple[int, ...]:
    if sequence_lengths is None:
        return ()
    normalized_values = []
    for raw_value in _flatten_values(sequence_lengths):
        normalized_values.append(_validate_non_negative_int("sequence_lengths", raw_value))
    return tuple(normalized_values)


def _infer_batch_size(
    batch_size: int | None,
    *,
    sample_ids: tuple[str, ...],
    class_counts: dict[str, int],
    sequence_lengths: tuple[int, ...],
) -> int | None:
    normalized_batch_size = batch_size
    class_count_total = sum(class_counts.values()) if class_counts else None
    derived_sizes = [
        len(sample_ids) if sample_ids else None,
        class_count_total,
        len(sequence_lengths) if sequence_lengths else None,
    ]
    if normalized_batch_size is not None:
        for derived_size in derived_sizes:
            if derived_size is not None and derived_size > normalized_batch_size:
                raise ValueError(
                    "batch_size must be at least as large as the provided sample IDs, class counts, and sequence lengths."
                )
        return normalized_batch_size

    observed_sizes = [size for size in derived_sizes if size is not None]
    if not observed_sizes:
        return None
    return max(observed_sizes)


def _merge_metadata(
    adapter_metadata: dict[str, Any] | None,
    explicit_metadata: dict[str, Any] | None,
) -> dict[str, Any] | None:
    if adapter_metadata is None and explicit_metadata is None:
        return None
    merged = dict(adapter_metadata or {})
    merged.update(explicit_metadata or {})
    return merged


def _mean(values: Sequence[float]) -> float:
    if not values:
        return 0.0
    return sum(values) / len(values)


def _percentile(values: Sequence[float], quantile: float) -> float:
    if not values:
        raise ValueError("Cannot compute percentiles for an empty sequence.")
    if len(values) == 1:
        return float(values[0])
    position = (len(values) - 1) * quantile
    lower_index = int(math.floor(position))
    upper_index = int(math.ceil(position))
    lower_value = float(values[lower_index])
    upper_value = float(values[upper_index])
    if lower_index == upper_index:
        return lower_value
    interpolation = position - lower_index
    return lower_value + ((upper_value - lower_value) * interpolation)


def _standard_deviation(values: Sequence[float]) -> float:
    if len(values) < 2:
        return 0.0
    mean_value = _mean(values)
    variance = sum((value - mean_value) ** 2 for value in values) / len(values)
    return math.sqrt(max(variance, 0.0))


def _z_score(value: float, *, mean_value: float, std_value: float) -> float | None:
    if std_value <= 0.0:
        return None
    return (value - mean_value) / std_value


@dataclass(frozen=True, slots=True)
class BatchMetadata:
    batch_size: int | None = None
    sample_ids: tuple[str, ...] = ()
    class_counts: dict[str, int] = field(default_factory=dict)
    sequence_lengths: tuple[int, ...] = ()
    metadata: dict[str, Any] | None = None

    def __post_init__(self) -> None:
        normalized_sample_ids = _normalize_sample_ids(self.sample_ids)
        normalized_class_counts = _normalize_class_counts(self.class_counts)
        normalized_sequence_lengths = _normalize_sequence_lengths(self.sequence_lengths)
        normalized_batch_size = (
            _validate_positive_int("batch_size", self.batch_size)
            if self.batch_size is not None
            else None
        )
        inferred_batch_size = _infer_batch_size(
            normalized_batch_size,
            sample_ids=normalized_sample_ids,
            class_counts=normalized_class_counts,
            sequence_lengths=normalized_sequence_lengths,
        )
        object.__setattr__(self, "batch_size", inferred_batch_size)
        object.__setattr__(self, "sample_ids", normalized_sample_ids)
        object.__setattr__(self, "class_counts", normalized_class_counts)
        object.__setattr__(self, "sequence_lengths", normalized_sequence_lengths)
        object.__setattr__(self, "metadata", _normalize_metadata(self.metadata))

    def to_dict(self) -> dict[str, Any]:
        return _serialize_dataclass(self)


class BatchMetadataAdapter(ABC):
    """Interface for extracting lightweight batch metadata without persisting raw samples."""

    @abstractmethod
    def extract_batch_metadata(self, batch: Any) -> BatchMetadata:
        """Return normalized metadata for one training batch."""


@dataclass(frozen=True, slots=True)
class MappingBatchMetadataAdapter(BatchMetadataAdapter):
    """Adapter for dictionary-like batches with configurable keys."""

    label_key: str | None = None
    sample_id_key: str | None = None
    sequence_length_key: str | None = None
    batch_size_key: str | None = None
    metadata_keys: tuple[str, ...] = ()

    def __post_init__(self) -> None:
        object.__setattr__(self, "label_key", _normalize_optional_key("label_key", self.label_key))
        object.__setattr__(
            self,
            "sample_id_key",
            _normalize_optional_key("sample_id_key", self.sample_id_key),
        )
        object.__setattr__(
            self,
            "sequence_length_key",
            _normalize_optional_key("sequence_length_key", self.sequence_length_key),
        )
        object.__setattr__(
            self,
            "batch_size_key",
            _normalize_optional_key("batch_size_key", self.batch_size_key),
        )
        normalized_metadata_keys: list[str] = []
        for raw_key in self.metadata_keys:
            normalized_key = _normalize_optional_key("metadata_keys", raw_key)
            if normalized_key is None:
                continue
            if normalized_key not in normalized_metadata_keys:
                normalized_metadata_keys.append(normalized_key)
        object.__setattr__(self, "metadata_keys", tuple(normalized_metadata_keys))

    def extract_batch_metadata(self, batch: Any) -> BatchMetadata:
        if not isinstance(batch, Mapping):
            raise TypeError("MappingBatchMetadataAdapter expects a mapping-like batch.")

        metadata_payload = {key: batch[key] for key in self.metadata_keys if key in batch}
        return BatchMetadata(
            batch_size=(
                _validate_positive_int("batch_size", batch[self.batch_size_key])
                if self.batch_size_key is not None and self.batch_size_key in batch
                else None
            ),
            sample_ids=(
                batch[self.sample_id_key]
                if self.sample_id_key is not None and self.sample_id_key in batch
                else ()
            ),
            class_counts=(
                _normalize_class_counts(batch[self.label_key])
                if self.label_key is not None and self.label_key in batch
                else {}
            ),
            sequence_lengths=(
                batch[self.sequence_length_key]
                if self.sequence_length_key is not None and self.sequence_length_key in batch
                else ()
            ),
            metadata=metadata_payload or None,
        )


@dataclass(frozen=True, slots=True)
class DataDiagnosticsConfig:
    batch_size: int = 100
    adapter: BatchMetadataAdapter | None = None
    outlier_z_score_threshold: float = 2.5
    outlier_iqr_multiplier: float = 1.5
    hard_batch_percentile: float = 0.9
    min_batches_for_outlier_detection: int = 5

    def __post_init__(self) -> None:
        object.__setattr__(self, "batch_size", _validate_positive_int("batch_size", self.batch_size))
        if self.adapter is not None and not isinstance(self.adapter, BatchMetadataAdapter):
            raise TypeError("adapter must implement BatchMetadataAdapter.")
        object.__setattr__(
            self,
            "outlier_z_score_threshold",
            _validate_positive_float(
                "outlier_z_score_threshold",
                self.outlier_z_score_threshold,
            ),
        )
        object.__setattr__(
            self,
            "outlier_iqr_multiplier",
            _validate_positive_float(
                "outlier_iqr_multiplier",
                self.outlier_iqr_multiplier,
            ),
        )
        object.__setattr__(
            self,
            "hard_batch_percentile",
            _validate_probability("hard_batch_percentile", self.hard_batch_percentile),
        )
        object.__setattr__(
            self,
            "min_batches_for_outlier_detection",
            _validate_positive_int(
                "min_batches_for_outlier_detection",
                self.min_batches_for_outlier_detection,
            ),
        )

    def to_dict(self) -> dict[str, Any]:
        payload = _serialize_dataclass(self)
        if self.adapter is not None:
            payload["adapter"] = self.adapter.__class__.__name__
        return payload


@dataclass(frozen=True, slots=True)
class LossDistributionSummary:
    count: int
    min_value: float
    max_value: float
    mean_value: float
    median_value: float
    std_value: float
    upper_fence: float
    hard_batch_threshold: float

    def to_dict(self) -> dict[str, Any]:
        return _serialize_dataclass(self)


@dataclass(frozen=True, slots=True)
class SequenceLengthDistributionSummary:
    observed_batch_count: int
    sample_count: int
    min_value: int
    max_value: int
    mean_value: float
    median_value: float
    upper_fence: float

    def to_dict(self) -> dict[str, Any]:
        return _serialize_dataclass(self)


@dataclass(frozen=True, slots=True)
class ClassExposure:
    label: str
    count: int
    ratio: float
    batch_count: int

    def to_dict(self) -> dict[str, Any]:
        return _serialize_dataclass(self)


@dataclass(frozen=True, slots=True)
class ExposureSummary:
    total_batches: int
    batches_with_labels: int
    batches_with_batch_size: int
    total_examples: int
    total_label_observations: int
    unique_sample_id_count: int
    class_exposures: tuple[ClassExposure, ...] = ()

    def to_dict(self) -> dict[str, Any]:
        return _serialize_dataclass(self)


@dataclass(frozen=True, slots=True)
class HardBatchFinding:
    batch: BatchDiagnosticRecord
    reasons: tuple[str, ...]
    loss_z_score: float | None
    loss_upper_fence: float | None
    hard_batch_threshold: float | None
    mean_sequence_length: float | None = None
    sequence_length_z_score: float | None = None
    sequence_length_upper_fence: float | None = None

    def to_dict(self) -> dict[str, Any]:
        return _serialize_dataclass(self)


@dataclass(frozen=True, slots=True)
class DataDiagnosticsSummary:
    total_batches: int
    step_start: int | None
    step_end: int | None
    loss_summary: LossDistributionSummary | None
    sequence_length_summary: SequenceLengthDistributionSummary | None
    exposure_summary: ExposureSummary
    hard_batch_count: int
    hard_batches: tuple[HardBatchFinding, ...] = ()

    def to_dict(self) -> dict[str, Any]:
        return _serialize_dataclass(self)


@dataclass(frozen=True, slots=True)
class _DistributionStats:
    count: int
    mean_value: float
    std_value: float
    min_value: float
    max_value: float
    median_value: float
    q1_value: float
    q3_value: float
    upper_fence: float
    high_percentile_value: float


class DataDiagnostics:
    """Batch-level diagnostics logger and analyzer for one training run."""

    _sequence_lock: ClassVar[threading.RLock] = threading.RLock()
    _next_sequence_by_run: ClassVar[dict[tuple[str, str], int]] = {}

    def __init__(
        self,
        session: Any,
        *,
        config: DataDiagnosticsConfig | None = None,
    ) -> None:
        self._context, self._store = self._resolve_session_dependencies(session)
        self._session = session if callable(getattr(session, "_register_flushable", None)) else None
        self._config = config or DataDiagnosticsConfig()
        self._buffer: list[BatchDiagnosticRecord] = []
        self._lock = threading.RLock()
        self._closed = False
        if self._session is not None:
            self._session._register_flushable(self)

    @property
    def config(self) -> DataDiagnosticsConfig:
        return self._config

    @property
    def run_id(self) -> str:
        return self._context.run_id

    @property
    def buffered_count(self) -> int:
        with self._lock:
            return len(self._buffer)

    def log_batch(
        self,
        loss: float,
        *,
        step: int | None = None,
        epoch: int | None = None,
        batch: Any | None = None,
        batch_size: int | None = None,
        labels: Mapping[Any, int] | Sequence[Any] | Any | None = None,
        sample_ids: Sequence[Any] | Any | None = None,
        sequence_lengths: Sequence[int] | int | None = None,
        metadata: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> BatchDiagnosticRecord:
        """Queue one batch diagnostic record for batched persistence."""

        self._ensure_open()
        normalized_loss = _normalize_loss(loss)
        normalized_step = _normalize_index("step", step)
        normalized_epoch = _normalize_index("epoch", epoch)
        if normalized_step is None and normalized_epoch is None:
            raise ValueError("Batch diagnostics must include at least one of step or epoch.")

        adapter_metadata = BatchMetadata()
        if batch is not None:
            if self._config.adapter is None:
                raise ValueError("batch was provided but no batch metadata adapter is configured.")
            adapter_metadata = self._config.adapter.extract_batch_metadata(batch)
            if not isinstance(adapter_metadata, BatchMetadata):
                raise TypeError("Batch metadata adapters must return BatchMetadata instances.")

        normalized_sample_ids = (
            _normalize_sample_ids(sample_ids)
            if sample_ids is not None
            else adapter_metadata.sample_ids
        )
        normalized_class_counts = (
            _normalize_class_counts(labels)
            if labels is not None
            else dict(adapter_metadata.class_counts)
        )
        normalized_sequence_lengths = (
            _normalize_sequence_lengths(sequence_lengths)
            if sequence_lengths is not None
            else adapter_metadata.sequence_lengths
        )
        normalized_batch_size = (
            _validate_positive_int("batch_size", batch_size)
            if batch_size is not None
            else adapter_metadata.batch_size
        )
        normalized_metadata = _merge_metadata(
            adapter_metadata.metadata,
            _normalize_metadata(metadata),
        )

        record = BatchDiagnosticRecord(
            id=uuid.uuid4().hex,
            run_id=self.run_id,
            sequence_id=self._allocate_sequence_id(),
            loss=normalized_loss,
            step=normalized_step,
            epoch=normalized_epoch,
            batch_size=_infer_batch_size(
                normalized_batch_size,
                sample_ids=normalized_sample_ids,
                class_counts=normalized_class_counts,
                sequence_lengths=normalized_sequence_lengths,
            ),
            timestamp=timestamp or utc_now(),
            sample_ids=normalized_sample_ids,
            class_counts=normalized_class_counts,
            sequence_lengths=normalized_sequence_lengths,
            metadata=normalized_metadata,
        )
        self._enqueue(record)
        return record

    def flush(self) -> tuple[BatchDiagnosticRecord, ...]:
        """Persist any buffered batch diagnostics as one storage transaction."""

        self._ensure_open()
        with self._lock:
            if not self._buffer:
                return ()
            batch = tuple(self._buffer)
            self._buffer.clear()

        try:
            return self._store.write_batch_diagnostics(self.run_id, batch)
        except Exception:
            with self._lock:
                self._buffer = list(batch) + self._buffer
            raise

    def read(
        self,
        *,
        sample_id: str | None = None,
        step: int | None = None,
        step_min: int | None = None,
        step_max: int | None = None,
        epoch: int | None = None,
        epoch_min: int | None = None,
        epoch_max: int | None = None,
        limit: int | None = None,
    ) -> tuple[BatchDiagnosticRecord, ...]:
        """Flush pending records and return persisted batch diagnostics for the run."""

        if limit is not None and limit <= 0:
            raise ValueError("limit must be a positive integer when provided.")
        if not self._closed:
            self.flush()

        normalized_sample_id = (
            _normalize_single_sample_id(sample_id)
            if sample_id is not None
            else None
        )
        records = self._store.read_batch_diagnostics(
            self.run_id,
            step=step,
            step_min=step_min,
            step_max=step_max,
            epoch=epoch,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
            limit=None if normalized_sample_id is not None else limit,
        )
        if normalized_sample_id is not None:
            records = tuple(
                record for record in records if normalized_sample_id in record.sample_ids
            )
            if limit is not None:
                records = records[:limit]
        return records

    def summarize_exposure(
        self,
        *,
        step_min: int | None = None,
        step_max: int | None = None,
        epoch_min: int | None = None,
        epoch_max: int | None = None,
    ) -> ExposureSummary:
        records = self.read(
            step_min=step_min,
            step_max=step_max,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
        )
        return self._build_exposure_summary(records)

    def find_hard_batches(
        self,
        *,
        step_min: int | None = None,
        step_max: int | None = None,
        epoch_min: int | None = None,
        epoch_max: int | None = None,
        limit: int | None = 10,
    ) -> tuple[HardBatchFinding, ...]:
        if limit is not None and limit <= 0:
            raise ValueError("limit must be a positive integer when provided.")
        records = self.read(
            step_min=step_min,
            step_max=step_max,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
        )
        return self._build_hard_batch_findings(records, limit=limit)

    def summarize(
        self,
        *,
        step_min: int | None = None,
        step_max: int | None = None,
        epoch_min: int | None = None,
        epoch_max: int | None = None,
        hard_batch_limit: int | None = 10,
    ) -> DataDiagnosticsSummary:
        if hard_batch_limit is not None and hard_batch_limit <= 0:
            raise ValueError("hard_batch_limit must be a positive integer when provided.")

        records = self.read(
            step_min=step_min,
            step_max=step_max,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
        )
        hard_batches = self._build_hard_batch_findings(records, limit=None)
        step_start, step_end = self._step_range(records)
        return DataDiagnosticsSummary(
            total_batches=len(records),
            step_start=step_start,
            step_end=step_end,
            loss_summary=self._build_loss_summary(records),
            sequence_length_summary=self._build_sequence_length_summary(records),
            exposure_summary=self._build_exposure_summary(records),
            hard_batch_count=len(hard_batches),
            hard_batches=(
                hard_batches
                if hard_batch_limit is None
                else hard_batches[:hard_batch_limit]
            ),
        )

    def close(self) -> None:
        if self._closed:
            return
        try:
            self.flush()
        finally:
            self._closed = True
            if self._session is not None:
                self._session._unregister_flushable(self)

    def __enter__(self) -> DataDiagnostics:
        return self

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        self.close()

    def _enqueue(self, record: BatchDiagnosticRecord) -> None:
        should_flush = False
        with self._lock:
            self._buffer.append(record)
            should_flush = len(self._buffer) >= self._config.batch_size

        if should_flush:
            self.flush()

    def _allocate_sequence_id(self) -> int:
        run_key = (self._context.db_path, self.run_id)
        with self._sequence_lock:
            next_sequence = self._next_sequence_by_run.get(run_key)
            if next_sequence is None:
                next_sequence = self._store.get_next_batch_diagnostic_sequence(self.run_id)
            self._next_sequence_by_run[run_key] = next_sequence + 1
            return next_sequence

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("DataDiagnostics is closed.")

    def _build_loss_summary(
        self,
        records: tuple[BatchDiagnosticRecord, ...],
    ) -> LossDistributionSummary | None:
        if not records:
            return None
        stats = self._build_distribution(tuple(record.loss for record in records))
        return LossDistributionSummary(
            count=stats.count,
            min_value=stats.min_value,
            max_value=stats.max_value,
            mean_value=stats.mean_value,
            median_value=stats.median_value,
            std_value=stats.std_value,
            upper_fence=stats.upper_fence,
            hard_batch_threshold=stats.high_percentile_value,
        )

    def _build_sequence_length_summary(
        self,
        records: tuple[BatchDiagnosticRecord, ...],
    ) -> SequenceLengthDistributionSummary | None:
        batch_lengths = [record.sequence_lengths for record in records if record.sequence_lengths]
        all_lengths = [float(length) for lengths in batch_lengths for length in lengths]
        if not all_lengths:
            return None

        stats = self._build_distribution(tuple(all_lengths))
        return SequenceLengthDistributionSummary(
            observed_batch_count=len(batch_lengths),
            sample_count=len(all_lengths),
            min_value=int(stats.min_value),
            max_value=int(stats.max_value),
            mean_value=stats.mean_value,
            median_value=stats.median_value,
            upper_fence=stats.upper_fence,
        )

    @staticmethod
    def _build_exposure_summary(
        records: tuple[BatchDiagnosticRecord, ...],
    ) -> ExposureSummary:
        label_counts: Counter[str] = Counter()
        label_batches: Counter[str] = Counter()
        total_examples = 0
        batches_with_batch_size = 0
        batches_with_labels = 0
        unique_sample_ids: set[str] = set()

        for record in records:
            if record.batch_size is not None:
                total_examples += record.batch_size
                batches_with_batch_size += 1
            if record.sample_ids:
                unique_sample_ids.update(record.sample_ids)
            if not record.class_counts:
                continue
            batches_with_labels += 1
            for label, count in record.class_counts.items():
                label_counts[label] += count
                label_batches[label] += 1

        total_label_observations = sum(label_counts.values())
        exposures = tuple(
            ClassExposure(
                label=label,
                count=count,
                ratio=(count / total_label_observations) if total_label_observations > 0 else 0.0,
                batch_count=int(label_batches[label]),
            )
            for label, count in sorted(
                label_counts.items(),
                key=lambda item: (-item[1], item[0]),
            )
        )
        return ExposureSummary(
            total_batches=len(records),
            batches_with_labels=batches_with_labels,
            batches_with_batch_size=batches_with_batch_size,
            total_examples=total_examples,
            total_label_observations=total_label_observations,
            unique_sample_id_count=len(unique_sample_ids),
            class_exposures=exposures,
        )

    def _build_hard_batch_findings(
        self,
        records: tuple[BatchDiagnosticRecord, ...],
        *,
        limit: int | None,
    ) -> tuple[HardBatchFinding, ...]:
        if not records:
            return ()

        loss_stats = self._build_distribution(tuple(record.loss for record in records))
        batch_length_means = {
            record.sequence_id: _mean([float(length) for length in record.sequence_lengths])
            for record in records
            if record.sequence_lengths
        }
        length_stats = (
            self._build_distribution(tuple(batch_length_means.values()))
            if batch_length_means
            else None
        )
        outlier_detection_active = len(records) >= self._config.min_batches_for_outlier_detection

        findings: list[HardBatchFinding] = []
        for record in records:
            reasons: list[str] = []
            loss_z_score = _z_score(
                record.loss,
                mean_value=loss_stats.mean_value,
                std_value=loss_stats.std_value,
            )
            if outlier_detection_active:
                if loss_z_score is not None and loss_z_score >= self._config.outlier_z_score_threshold:
                    reasons.append("loss_outlier")
                elif record.loss > loss_stats.upper_fence:
                    reasons.append("loss_outlier")

            if record.loss >= loss_stats.high_percentile_value and record.loss > loss_stats.mean_value:
                reasons.append("high_loss")

            mean_sequence_length = batch_length_means.get(record.sequence_id)
            sequence_length_z_score = None
            sequence_length_upper_fence = None
            if mean_sequence_length is not None and length_stats is not None:
                sequence_length_z_score = _z_score(
                    mean_sequence_length,
                    mean_value=length_stats.mean_value,
                    std_value=length_stats.std_value,
                )
                sequence_length_upper_fence = length_stats.upper_fence
                if outlier_detection_active:
                    if (
                        sequence_length_z_score is not None
                        and sequence_length_z_score >= self._config.outlier_z_score_threshold
                    ):
                        reasons.append("long_sequence_batch")
                    elif mean_sequence_length > length_stats.upper_fence:
                        reasons.append("long_sequence_batch")

            if not reasons:
                continue

            findings.append(
                HardBatchFinding(
                    batch=record,
                    reasons=tuple(dict.fromkeys(reasons)),
                    loss_z_score=loss_z_score,
                    loss_upper_fence=loss_stats.upper_fence if outlier_detection_active else None,
                    hard_batch_threshold=loss_stats.high_percentile_value,
                    mean_sequence_length=mean_sequence_length,
                    sequence_length_z_score=sequence_length_z_score,
                    sequence_length_upper_fence=sequence_length_upper_fence,
                )
            )

        findings.sort(
            key=lambda item: (
                "loss_outlier" in item.reasons,
                "high_loss" in item.reasons,
                "long_sequence_batch" in item.reasons,
                item.batch.loss,
                item.loss_z_score if item.loss_z_score is not None else 0.0,
                item.mean_sequence_length if item.mean_sequence_length is not None else 0.0,
            ),
            reverse=True,
        )
        if limit is not None:
            findings = findings[:limit]
        return tuple(findings)

    def _build_distribution(self, values: tuple[float, ...]) -> _DistributionStats:
        ordered_values = tuple(sorted(float(value) for value in values))
        q1_value = _percentile(ordered_values, 0.25)
        q3_value = _percentile(ordered_values, 0.75)
        iqr = q3_value - q1_value
        return _DistributionStats(
            count=len(ordered_values),
            mean_value=_mean(ordered_values),
            std_value=_standard_deviation(ordered_values),
            min_value=float(ordered_values[0]),
            max_value=float(ordered_values[-1]),
            median_value=_percentile(ordered_values, 0.5),
            q1_value=q1_value,
            q3_value=q3_value,
            upper_fence=q3_value + (iqr * self._config.outlier_iqr_multiplier),
            high_percentile_value=_percentile(
                ordered_values,
                self._config.hard_batch_percentile,
            ),
        )

    @staticmethod
    def _step_range(
        records: tuple[BatchDiagnosticRecord, ...],
    ) -> tuple[int | None, int | None]:
        steps = [record.step for record in records if record.step is not None]
        if not steps:
            return None, None
        return min(steps), max(steps)

    @staticmethod
    def _resolve_session_dependencies(session: Any) -> tuple[RunContext, LocalSessionStore]:
        if isinstance(session, RunContext):
            return session, LocalSessionStore(session.storage_dir)

        context = getattr(session, "context", None)
        store = getattr(session, "_store", None)
        if isinstance(context, RunContext) and isinstance(store, LocalSessionStore):
            return context, store

        raise TypeError(
            "DataDiagnostics expects a TrainingSession or RunContext created by the Quark SDK."
        )


__all__ = [
    "BatchMetadata",
    "BatchMetadataAdapter",
    "MappingBatchMetadataAdapter",
    "DataDiagnostics",
    "DataDiagnosticsConfig",
    "DataDiagnosticsSummary",
    "LossDistributionSummary",
    "SequenceLengthDistributionSummary",
    "ClassExposure",
    "ExposureSummary",
    "HardBatchFinding",
]
