from __future__ import annotations

import json
import math
import re
import threading
import uuid
from collections.abc import Mapping
from dataclasses import dataclass, field
from numbers import Integral, Real
from typing import Any, ClassVar

from hook_engine import HookScheduleConfig
from models import ParameterDistributionRecord, ParameterSummaryRecord, RunContext
from storage import LocalSessionStore, utc_now

try:
    import torch
    import torch.nn as nn
except Exception:  # pragma: no cover - exercised through runtime guard.
    torch = None
    nn = None

_PARAMETER_ROLES = {"weight", "bias", "other"}
_DISTRIBUTION_TYPES = {"value", "gradient", "update"}


def _normalize_string_tuple(values: tuple[str, ...], *, label: str) -> tuple[str, ...]:
    normalized: list[str] = []
    for raw_value in values:
        if not isinstance(raw_value, str):
            raise TypeError(f"{label} values must be strings.")
        value = raw_value.strip()
        if not value:
            raise ValueError(f"{label} values must not be empty.")
        if value not in normalized:
            normalized.append(value)
    return tuple(normalized)


@dataclass(frozen=True, slots=True)
class ParameterDiagnosticsConfig:
    """Sampling configuration for model parameter summaries and histograms."""

    schedule: HookScheduleConfig = field(default_factory=HookScheduleConfig)
    batch_size: int = 100
    capture_summaries: bool = True
    capture_distributions: bool = False
    distribution_types: tuple[str, ...] = ("value",)
    histogram_bins: int = 32
    near_zero_epsilon: float = 1e-8
    max_histogram_values: int | None = 100_000
    include_parameter_patterns: tuple[str, ...] = ()
    exclude_parameter_patterns: tuple[str, ...] = ()
    include_parameter_roles: tuple[str, ...] = ("weight", "bias", "other")

    def __post_init__(self) -> None:
        if not isinstance(self.schedule, HookScheduleConfig):
            raise TypeError("schedule must be a HookScheduleConfig instance.")
        if isinstance(self.batch_size, bool) or not isinstance(self.batch_size, Integral):
            raise TypeError("batch_size must be an integer.")
        if int(self.batch_size) <= 0:
            raise ValueError("batch_size must be positive.")
        object.__setattr__(self, "batch_size", int(self.batch_size))
        if not isinstance(self.capture_summaries, bool):
            raise TypeError("capture_summaries must be a boolean.")
        if not isinstance(self.capture_distributions, bool):
            raise TypeError("capture_distributions must be a boolean.")
        if not self.capture_summaries and not self.capture_distributions:
            raise ValueError("At least one of capture_summaries or capture_distributions must be enabled.")

        distribution_types = _normalize_string_tuple(
            self.distribution_types,
            label="distribution_types",
        )
        unknown_types = sorted(set(distribution_types) - _DISTRIBUTION_TYPES)
        if unknown_types:
            raise ValueError(f"Unsupported distribution_types values: {unknown_types}.")
        object.__setattr__(self, "distribution_types", distribution_types)

        if isinstance(self.histogram_bins, bool) or not isinstance(self.histogram_bins, Integral):
            raise TypeError("histogram_bins must be an integer.")
        if int(self.histogram_bins) <= 0:
            raise ValueError("histogram_bins must be positive.")
        object.__setattr__(self, "histogram_bins", int(self.histogram_bins))

        if isinstance(self.near_zero_epsilon, bool) or not isinstance(self.near_zero_epsilon, Real):
            raise TypeError("near_zero_epsilon must be a real number.")
        if float(self.near_zero_epsilon) < 0.0:
            raise ValueError("near_zero_epsilon must not be negative.")
        object.__setattr__(self, "near_zero_epsilon", float(self.near_zero_epsilon))

        if self.max_histogram_values is not None:
            if isinstance(self.max_histogram_values, bool) or not isinstance(
                self.max_histogram_values,
                Integral,
            ):
                raise TypeError("max_histogram_values must be an integer or None.")
            if int(self.max_histogram_values) <= 0:
                raise ValueError("max_histogram_values must be positive when provided.")
            object.__setattr__(self, "max_histogram_values", int(self.max_histogram_values))

        include_patterns = _normalize_string_tuple(
            self.include_parameter_patterns,
            label="include_parameter_patterns",
        )
        exclude_patterns = _normalize_string_tuple(
            self.exclude_parameter_patterns,
            label="exclude_parameter_patterns",
        )
        for pattern in include_patterns + exclude_patterns:
            re.compile(pattern)
        object.__setattr__(self, "include_parameter_patterns", include_patterns)
        object.__setattr__(self, "exclude_parameter_patterns", exclude_patterns)

        include_roles = _normalize_string_tuple(
            self.include_parameter_roles,
            label="include_parameter_roles",
        )
        unknown_roles = sorted(set(include_roles) - _PARAMETER_ROLES)
        if unknown_roles:
            raise ValueError(f"Unsupported include_parameter_roles values: {unknown_roles}.")
        object.__setattr__(self, "include_parameter_roles", include_roles)


class ParameterDiagnostics:
    """Run-scoped writer for parameter summaries and optional histogram distributions."""

    _summary_sequence_lock: ClassVar[threading.RLock] = threading.RLock()
    _distribution_sequence_lock: ClassVar[threading.RLock] = threading.RLock()
    _next_summary_sequence_by_run: ClassVar[dict[tuple[str, str], int]] = {}
    _next_distribution_sequence_by_run: ClassVar[dict[tuple[str, str], int]] = {}

    def __init__(
        self,
        session: Any,
        model: Any,
        optimizer: Any | None = None,
        *,
        config: ParameterDiagnosticsConfig | None = None,
    ) -> None:
        if torch is None or nn is None:
            raise ImportError("ParameterDiagnostics requires PyTorch to be installed.")
        if not isinstance(model, nn.Module):
            raise TypeError("ParameterDiagnostics expects a torch.nn.Module model instance.")

        self._context, self._store = self._resolve_session_dependencies(session)
        self._session = session if callable(getattr(session, "_register_flushable", None)) else None
        self._model = model
        self._optimizer = optimizer
        self._config = config or ParameterDiagnosticsConfig()
        self._include_parameter_patterns = tuple(
            re.compile(pattern) for pattern in self._config.include_parameter_patterns
        )
        self._exclude_parameter_patterns = tuple(
            re.compile(pattern) for pattern in self._config.exclude_parameter_patterns
        )
        self._summary_buffer: list[ParameterSummaryRecord] = []
        self._distribution_buffer: list[ParameterDistributionRecord] = []
        self._previous_parameters: dict[str, Any] = {}
        self._lock = threading.RLock()
        self._closed = False
        if self._session is not None:
            self._session._register_flushable(self)

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

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

    @property
    def buffered_summary_count(self) -> int:
        with self._lock:
            return len(self._summary_buffer)

    @property
    def buffered_distribution_count(self) -> int:
        with self._lock:
            return len(self._distribution_buffer)

    def snapshot(
        self,
        step: int,
        *,
        epoch: int | None = None,
        metadata: Mapping[str, Any] | None = None,
    ) -> tuple[tuple[ParameterSummaryRecord, ...], tuple[ParameterDistributionRecord, ...]]:
        """Capture one scheduled parameter snapshot for the current model state."""

        self._ensure_open()
        normalized_step = self._normalize_index("step", step, required=True)
        normalized_epoch = self._normalize_index("epoch", epoch)
        if not self._config.schedule.should_probe(normalized_step):
            return (), ()

        normalized_metadata = self._normalize_metadata(metadata)
        module_by_name = dict(self._model.named_modules())
        optimizer_name = self._optimizer.__class__.__name__ if self._optimizer is not None else None
        optimizer_group_by_parameter_id = self._optimizer_group_by_parameter_id()
        timestamp = utc_now()
        summaries: list[ParameterSummaryRecord] = []
        distributions: list[ParameterDistributionRecord] = []

        for parameter_key, parameter in self._model.named_parameters():
            parameter_name, module_name, module_type, parameter_role = self._parameter_identity(
                parameter_key,
                module_by_name=module_by_name,
            )
            if not self._matches_filters(parameter_key, parameter_role):
                continue

            parameter_tensor = self._prepare_tensor(parameter)
            if parameter_tensor is None:
                continue
            gradient_tensor = self._prepare_tensor(parameter.grad) if parameter.grad is not None else None
            previous_tensor = self._previous_parameters.get(parameter_key)
            update_tensor = None
            if previous_tensor is not None and tuple(previous_tensor.shape) == tuple(parameter_tensor.shape):
                update_tensor = parameter_tensor - previous_tensor

            summary = None
            if self._config.capture_summaries:
                summary = self._build_summary_record(
                    parameter_key=parameter_key,
                    parameter_name=parameter_name,
                    parameter_role=parameter_role,
                    module_name=module_name,
                    module_type=module_type,
                    parameter=parameter,
                    parameter_tensor=parameter_tensor,
                    gradient_tensor=gradient_tensor,
                    update_tensor=update_tensor,
                    step=normalized_step,
                    epoch=normalized_epoch,
                    timestamp=timestamp,
                    optimizer_name=optimizer_name,
                    optimizer_group_index=optimizer_group_by_parameter_id.get(id(parameter)),
                    metadata=normalized_metadata,
                )
                summaries.append(summary)

            if self._config.capture_distributions:
                distributions.extend(
                    self._build_distribution_records(
                        parameter_key=parameter_key,
                        parameter_name=parameter_name,
                        parameter_role=parameter_role,
                        module_name=module_name,
                        module_type=module_type,
                        parameter_tensor=parameter_tensor,
                        gradient_tensor=gradient_tensor,
                        update_tensor=update_tensor,
                        step=normalized_step,
                        epoch=normalized_epoch,
                        timestamp=timestamp,
                        summary_id=summary.id if summary is not None else None,
                        metadata=normalized_metadata,
                    )
                )

            self._previous_parameters[parameter_key] = parameter_tensor.detach().clone()

        self._enqueue(tuple(summaries), tuple(distributions))
        return tuple(summaries), tuple(distributions)

    def flush(self) -> tuple[tuple[ParameterSummaryRecord, ...], tuple[ParameterDistributionRecord, ...]]:
        """Persist buffered parameter summaries and distributions."""

        self._ensure_open()
        with self._lock:
            summaries = tuple(self._summary_buffer)
            distributions = tuple(self._distribution_buffer)
            self._summary_buffer.clear()
            self._distribution_buffer.clear()

        persisted_summaries: tuple[ParameterSummaryRecord, ...] = ()
        persisted_distributions: tuple[ParameterDistributionRecord, ...] = ()
        try:
            if summaries:
                persisted_summaries = self._store.write_parameter_summaries(
                    self.run_id,
                    summaries,
                )
            if distributions:
                persisted_distributions = self._store.write_parameter_distributions(
                    self.run_id,
                    distributions,
                )
        except Exception:
            with self._lock:
                self._summary_buffer = list(summaries) + self._summary_buffer
                self._distribution_buffer = list(distributions) + self._distribution_buffer
            raise
        return persisted_summaries, persisted_distributions

    def read_summaries(
        self,
        *,
        module_name: str | None = None,
        module_type: str | None = None,
        parameter_name: str | None = None,
        parameter_key: str | None = None,
        parameter_role: 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,
        latest_only: bool | None = None,
        limit: int | None = None,
    ) -> tuple[ParameterSummaryRecord, ...]:
        if not self._closed:
            self.flush()
        return self._store.read_parameter_summaries(
            self.run_id,
            module_name=module_name,
            module_type=module_type,
            parameter_name=parameter_name,
            parameter_key=parameter_key,
            parameter_role=parameter_role,
            step=step,
            step_min=step_min,
            step_max=step_max,
            epoch=epoch,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
            latest_only=latest_only,
            limit=limit,
        )

    def read_distributions(
        self,
        *,
        module_name: str | None = None,
        module_type: str | None = None,
        parameter_name: str | None = None,
        parameter_key: str | None = None,
        parameter_role: str | None = None,
        distribution_type: 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,
        latest_only: bool | None = None,
        limit: int | None = None,
    ) -> tuple[ParameterDistributionRecord, ...]:
        if not self._closed:
            self.flush()
        return self._store.read_parameter_distributions(
            self.run_id,
            module_name=module_name,
            module_type=module_type,
            parameter_name=parameter_name,
            parameter_key=parameter_key,
            parameter_role=parameter_role,
            distribution_type=distribution_type,
            step=step,
            step_min=step_min,
            step_max=step_max,
            epoch=epoch,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
            latest_only=latest_only,
            limit=limit,
        )

    def read(self) -> tuple[tuple[ParameterSummaryRecord, ...], tuple[ParameterDistributionRecord, ...]]:
        """Flush and return all persisted parameter summaries and distributions."""

        return self.read_summaries(), self.read_distributions()

    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) -> ParameterDiagnostics:
        return self

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

    def _build_summary_record(
        self,
        *,
        parameter_key: str,
        parameter_name: str,
        parameter_role: str,
        module_name: str,
        module_type: str | None,
        parameter: Any,
        parameter_tensor: Any,
        gradient_tensor: Any | None,
        update_tensor: Any | None,
        step: int,
        epoch: int | None,
        timestamp: str,
        optimizer_name: str | None,
        optimizer_group_index: int | None,
        metadata: dict[str, Any] | None,
    ) -> ParameterSummaryRecord:
        parameter_stats = self._tensor_stats(parameter_tensor)
        gradient_stats = self._tensor_stats(gradient_tensor) if gradient_tensor is not None else {}
        update_stats = self._tensor_stats(update_tensor) if update_tensor is not None else {}
        parameter_norm = parameter_stats.get("norm")
        update_norm = update_stats.get("norm")
        update_ratio = None
        if parameter_norm is not None and parameter_norm > 0.0 and update_norm is not None:
            update_ratio = min(update_norm / parameter_norm, 1.0)

        return ParameterSummaryRecord(
            id=uuid.uuid4().hex,
            run_id=self.run_id,
            sequence_id=self._allocate_summary_sequence_id(),
            step=step,
            epoch=epoch,
            timestamp=timestamp,
            module_name=module_name,
            module_type=module_type,
            parameter_name=parameter_name,
            parameter_key=parameter_key,
            parameter_role=parameter_role,  # type: ignore[arg-type]
            tensor_shape=tuple(int(dimension) for dimension in parameter.shape),
            dtype=str(parameter.dtype).replace("torch.", ""),
            requires_grad=bool(parameter.requires_grad),
            element_count=int(parameter.numel()),
            mean=parameter_stats.get("mean"),
            std=parameter_stats.get("std"),
            variance=parameter_stats.get("variance"),
            norm=parameter_stats.get("norm"),
            min=parameter_stats.get("min"),
            max=parameter_stats.get("max"),
            abs_mean=parameter_stats.get("abs_mean"),
            abs_max=parameter_stats.get("abs_max"),
            zero_ratio=parameter_stats.get("zero_ratio"),
            near_zero_ratio=parameter_stats.get("near_zero_ratio"),
            finite_ratio=parameter_stats.get("finite_ratio"),
            nan_count=parameter_stats.get("nan_count"),
            inf_count=parameter_stats.get("inf_count"),
            grad_mean=gradient_stats.get("mean"),
            grad_std=gradient_stats.get("std"),
            grad_norm=gradient_stats.get("norm"),
            grad_abs_max=gradient_stats.get("abs_max"),
            update_norm=update_stats.get("norm"),
            update_mean=update_stats.get("mean"),
            update_std=update_stats.get("std"),
            update_to_parameter_norm_ratio=update_ratio,
            optimizer_name=optimizer_name,
            optimizer_group_index=optimizer_group_index,
            metadata=metadata,
        )

    def _build_distribution_records(
        self,
        *,
        parameter_key: str,
        parameter_name: str,
        parameter_role: str,
        module_name: str,
        module_type: str | None,
        parameter_tensor: Any,
        gradient_tensor: Any | None,
        update_tensor: Any | None,
        step: int,
        epoch: int | None,
        timestamp: str,
        summary_id: str | None,
        metadata: dict[str, Any] | None,
    ) -> tuple[ParameterDistributionRecord, ...]:
        source_by_type = {
            "value": parameter_tensor,
            "gradient": gradient_tensor,
            "update": update_tensor,
        }
        records: list[ParameterDistributionRecord] = []
        for distribution_type in self._config.distribution_types:
            source_tensor = source_by_type.get(distribution_type)
            if source_tensor is None:
                continue
            histogram = self._histogram(source_tensor)
            if histogram is None:
                continue
            records.append(
                ParameterDistributionRecord(
                    id=uuid.uuid4().hex,
                    run_id=self.run_id,
                    sequence_id=self._allocate_distribution_sequence_id(),
                    step=step,
                    epoch=epoch,
                    timestamp=timestamp,
                    module_name=module_name,
                    module_type=module_type,
                    parameter_name=parameter_name,
                    parameter_key=parameter_key,
                    parameter_role=parameter_role,  # type: ignore[arg-type]
                    distribution_type=distribution_type,  # type: ignore[arg-type]
                    bin_edges=histogram["bin_edges"],
                    counts=histogram["counts"],
                    density=False,
                    sample_count=histogram["sample_count"],
                    total_count=histogram["total_count"],
                    min=histogram["min"],
                    max=histogram["max"],
                    nan_count=histogram["nan_count"],
                    inf_count=histogram["inf_count"],
                    clipped_low_count=0,
                    clipped_high_count=0,
                    summary_id=summary_id,
                    metadata=metadata,
                )
            )
        return tuple(records)

    def _tensor_stats(self, tensor: Any | None) -> dict[str, Any]:
        if tensor is None:
            return {}
        flattened = tensor.reshape(-1).float()
        element_count = int(flattened.numel())
        if element_count == 0:
            return {
                "finite_ratio": 1.0,
                "nan_count": 0,
                "inf_count": 0,
            }
        finite_mask = torch.isfinite(flattened)
        finite_count = int(finite_mask.sum().item())
        nan_count = int(torch.isnan(flattened).sum().item())
        inf_count = int(torch.isinf(flattened).sum().item())
        stats: dict[str, Any] = {
            "finite_ratio": finite_count / element_count,
            "nan_count": nan_count,
            "inf_count": inf_count,
        }
        if finite_count == 0:
            return stats
        finite_values = flattened[finite_mask]
        abs_values = finite_values.abs()
        stats.update(
            {
                "mean": float(finite_values.mean().item()),
                "std": float(finite_values.std(unbiased=False).item()),
                "variance": float(finite_values.var(unbiased=False).item()),
                "norm": float(finite_values.norm(2).item()),
                "min": float(finite_values.min().item()),
                "max": float(finite_values.max().item()),
                "abs_mean": float(abs_values.mean().item()),
                "abs_max": float(abs_values.max().item()),
                "zero_ratio": int(finite_values.eq(0).sum().item()) / finite_count,
                "near_zero_ratio": int(abs_values.le(self._config.near_zero_epsilon).sum().item()) / finite_count,
            }
        )
        return stats

    def _histogram(self, tensor: Any) -> dict[str, Any] | None:
        flattened = tensor.reshape(-1).float()
        total_count = int(flattened.numel())
        if total_count == 0:
            return None
        nan_count = int(torch.isnan(flattened).sum().item())
        inf_count = int(torch.isinf(flattened).sum().item())
        finite_values = flattened[torch.isfinite(flattened)]
        if int(finite_values.numel()) == 0:
            return None
        if (
            self._config.max_histogram_values is not None
            and int(finite_values.numel()) > self._config.max_histogram_values
        ):
            stride = math.ceil(int(finite_values.numel()) / self._config.max_histogram_values)
            finite_values = finite_values[::stride][: self._config.max_histogram_values]
        observed_min = float(finite_values.min().item())
        observed_max = float(finite_values.max().item())
        if not math.isfinite(observed_min) or not math.isfinite(observed_max):
            return None
        lower = observed_min
        upper = observed_max
        if lower == upper:
            width = max(abs(lower) * 0.01, 1e-6)
            lower -= width
            upper += width
        counts_tensor = torch.histc(
            finite_values,
            bins=self._config.histogram_bins,
            min=lower,
            max=upper,
        )
        bin_width = (upper - lower) / self._config.histogram_bins
        bin_edges = tuple(float(lower + (bin_width * index)) for index in range(self._config.histogram_bins + 1))
        counts = tuple(int(value) for value in counts_tensor.to(dtype=torch.int64).tolist())
        return {
            "bin_edges": bin_edges,
            "counts": counts,
            "sample_count": int(finite_values.numel()),
            "total_count": total_count,
            "min": observed_min,
            "max": observed_max,
            "nan_count": nan_count,
            "inf_count": inf_count,
        }

    def _enqueue(
        self,
        summaries: tuple[ParameterSummaryRecord, ...],
        distributions: tuple[ParameterDistributionRecord, ...],
    ) -> None:
        should_flush = False
        with self._lock:
            self._summary_buffer.extend(summaries)
            self._distribution_buffer.extend(distributions)
            should_flush = (
                len(self._summary_buffer) >= self._config.batch_size
                or len(self._distribution_buffer) >= self._config.batch_size
            )
        if should_flush:
            self.flush()

    def _allocate_summary_sequence_id(self) -> int:
        run_key = (self._context.db_path, self.run_id)
        with self._summary_sequence_lock:
            next_sequence = self._next_summary_sequence_by_run.get(run_key)
            if next_sequence is None:
                next_sequence = self._store.get_next_parameter_summary_sequence(self.run_id)
            self._next_summary_sequence_by_run[run_key] = next_sequence + 1
            return next_sequence

    def _allocate_distribution_sequence_id(self) -> int:
        run_key = (self._context.db_path, self.run_id)
        with self._distribution_sequence_lock:
            next_sequence = self._next_distribution_sequence_by_run.get(run_key)
            if next_sequence is None:
                next_sequence = self._store.get_next_parameter_distribution_sequence(self.run_id)
            self._next_distribution_sequence_by_run[run_key] = next_sequence + 1
            return next_sequence

    def _prepare_tensor(self, value: Any) -> Any | None:
        if torch is None or not isinstance(value, torch.Tensor):
            return None
        tensor = value.detach()
        if tensor.is_sparse:
            tensor = tensor.coalesce().values()
        if tensor.is_complex():
            tensor = tensor.abs()
        if tensor.numel() == 0:
            return None
        if not torch.is_floating_point(tensor) or tensor.dtype != torch.float32:
            tensor = tensor.float()
        return tensor.cpu()

    def _parameter_identity(
        self,
        parameter_key: str,
        *,
        module_by_name: dict[str, Any],
    ) -> tuple[str, str, str | None, str]:
        if "." in parameter_key:
            module_name, parameter_name = parameter_key.rsplit(".", 1)
        else:
            module_name = "<root>"
            parameter_name = parameter_key
        module = module_by_name.get("" if module_name == "<root>" else module_name)
        module_type = module.__class__.__name__ if module is not None else None
        if parameter_name == "weight":
            parameter_role = "weight"
        elif parameter_name == "bias":
            parameter_role = "bias"
        else:
            parameter_role = "other"
        return parameter_name, module_name, module_type, parameter_role

    def _matches_filters(self, parameter_key: str, parameter_role: str) -> bool:
        if parameter_role not in self._config.include_parameter_roles:
            return False
        if self._include_parameter_patterns and not any(
            pattern.search(parameter_key) for pattern in self._include_parameter_patterns
        ):
            return False
        if any(pattern.search(parameter_key) for pattern in self._exclude_parameter_patterns):
            return False
        return True

    def _optimizer_group_by_parameter_id(self) -> dict[int, int]:
        if self._optimizer is None:
            return {}
        groups = getattr(self._optimizer, "param_groups", ())
        mapping: dict[int, int] = {}
        for group_index, group in enumerate(groups):
            for parameter in group.get("params", ()):
                mapping[id(parameter)] = group_index
        return mapping

    @staticmethod
    def _normalize_index(label: str, value: int | None, *, required: bool = False) -> int | None:
        if value is None:
            if required:
                raise ValueError(f"{label} is required.")
            return None
        if isinstance(value, bool) or not isinstance(value, Integral):
            raise TypeError(f"{label} must be an integer.")
        normalized = int(value)
        if normalized < 0:
            raise ValueError(f"{label} must not be negative.")
        return normalized

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

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

    @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(
            "ParameterDiagnostics expects a TrainingSession or RunContext created by the Quark SDK."
        )
