from __future__ import annotations

from collections import Counter
from dataclasses import dataclass, field, fields, is_dataclass
from numbers import Real
from typing import Any, Iterable, Literal

from models import AlertRecord, LayerSummaryRecord, MetricRecord, PhaseRecord, RunContext
from storage import LocalSessionStore, utc_now

COMPARISON_SCHEMA_VERSION = 1
ComparisonMode = Literal["pairwise", "baseline"]
MetricTargetDirection = Literal["at_least", "at_most"]

_ALERT_SEVERITY_ORDER = {
    "info": 0,
    "warning": 1,
    "error": 2,
    "critical": 3,
}
_ALERT_SEVERITIES = ("info", "warning", "error", "critical")
_ALERT_CATEGORIES = ("optimization", "representation", "generalization", "runtime")


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_real(name: str, value: Real) -> float:
    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real number.")
    return float(value)


def _validate_metric_key(metric_key: str) -> str:
    candidate = metric_key.strip()
    if not candidate:
        raise ValueError("metric_key must not be empty.")
    return candidate


def _ordered_union(primary_values: Iterable[Any], reference_values: Iterable[Any]) -> tuple[Any, ...]:
    ordered_values: list[Any] = []
    seen: set[Any] = set()
    for value in tuple(primary_values) + tuple(reference_values):
        if value in seen:
            continue
        seen.add(value)
        ordered_values.append(value)
    return tuple(ordered_values)


@dataclass(frozen=True, slots=True)
class MetricTargetSpec:
    """Target rule used to compute time-to-target comparisons for a metric."""

    metric_key: str
    target_value: float
    direction: MetricTargetDirection

    def __post_init__(self) -> None:
        object.__setattr__(self, "metric_key", _validate_metric_key(self.metric_key))
        object.__setattr__(self, "target_value", _validate_real("target_value", self.target_value))
        if self.direction not in {"at_least", "at_most"}:
            raise ValueError("direction must be either 'at_least' or 'at_most'.")

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


@dataclass(frozen=True, slots=True)
class ComparisonEngineConfig:
    """Optional comparison settings for metric targets and milestone diffs."""

    metric_targets: tuple[MetricTargetSpec, ...] = ()

    def __post_init__(self) -> None:
        normalized_targets: list[MetricTargetSpec] = []
        seen_metric_keys: set[str] = set()
        for target in self.metric_targets:
            if not isinstance(target, MetricTargetSpec):
                raise TypeError("metric_targets must contain MetricTargetSpec instances.")
            if target.metric_key in seen_metric_keys:
                raise ValueError(
                    f"Duplicate metric target '{target.metric_key}' detected in comparison config."
                )
            seen_metric_keys.add(target.metric_key)
            normalized_targets.append(target)
        object.__setattr__(self, "metric_targets", tuple(normalized_targets))

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


@dataclass(frozen=True, slots=True)
class NumericDiff:
    """Paired numeric values plus their delta relative to the reference run."""

    primary: float | int | None
    reference: float | int | None
    delta: float | int | None
    percent_delta: float | None = None

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


@dataclass(frozen=True, slots=True)
class ComparisonRunTarget:
    """Stable metadata describing one side of a run comparison."""

    run_id: str
    project_name: str
    status: str
    tags: tuple[str, ...] = ()
    is_baseline: bool = False

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


@dataclass(frozen=True, slots=True)
class MetricSeriesSummary:
    """Aggregated summary for one metric stream within a single run."""

    sample_count: int
    step_start: int | None
    step_end: int | None
    first_value: float | None
    latest_value: float | None
    min_value: float | None
    max_value: float | None
    mean_value: float | None
    net_change: float | None
    target: MetricTargetSpec | None = None
    time_to_target_step: int | None = None

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


@dataclass(frozen=True, slots=True)
class MetricComparison:
    """Machine-readable metric diff between a primary and reference run."""

    metric_key: str
    primary: MetricSeriesSummary
    reference: MetricSeriesSummary
    latest_value_diff: NumericDiff
    mean_value_diff: NumericDiff
    min_value_diff: NumericDiff
    max_value_diff: NumericDiff
    net_change_diff: NumericDiff
    sample_count_diff: NumericDiff
    time_to_target_step_diff: NumericDiff | None = None

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


@dataclass(frozen=True, slots=True)
class AlertRuleDiff:
    """Per-rule alert count and severity differences between runs."""

    rule_id: str
    primary_category: str | None
    reference_category: str | None
    primary_count: int
    reference_count: int
    count_diff: int
    primary_max_severity: str | None
    reference_max_severity: str | None
    severity_changed: bool

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


@dataclass(frozen=True, slots=True)
class AlertDiff:
    """Aggregate alert-profile differences between two runs."""

    primary_total_count: int
    reference_total_count: int
    total_count_diff: int
    primary_severity_counts: dict[str, int] = field(default_factory=dict)
    reference_severity_counts: dict[str, int] = field(default_factory=dict)
    severity_count_diffs: dict[str, int] = field(default_factory=dict)
    primary_category_counts: dict[str, int] = field(default_factory=dict)
    reference_category_counts: dict[str, int] = field(default_factory=dict)
    category_count_diffs: dict[str, int] = field(default_factory=dict)
    rule_diffs: tuple[AlertRuleDiff, ...] = ()

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


@dataclass(frozen=True, slots=True)
class PhaseAggregate:
    """Aggregated phase duration summary for a phase label within a single run."""

    display_name: str | None
    segment_count: int
    total_duration_steps: int
    first_step: int | None
    last_step: int | None

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


@dataclass(frozen=True, slots=True)
class PhaseComparison:
    """Phase duration and presence differences between two runs."""

    phase_name: str
    primary: PhaseAggregate
    reference: PhaseAggregate
    segment_count_diff: int
    total_duration_diff: int

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


@dataclass(frozen=True, slots=True)
class LayerSummaryAggregate:
    """Aggregated layer-summary statistics for one layer probe stream."""

    sample_count: int
    step_start: int | None
    step_end: int | None
    average_mean: float | None
    average_std: float | None
    average_norm: float | None
    average_variance: float | None
    average_dead_unit_ratio: float | None
    average_saturation_ratio: float | None

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


@dataclass(frozen=True, slots=True)
class LayerSummaryComparison:
    """Layer-summary differences keyed by layer, summary type, and hook direction."""

    module_name: str
    module_type: str
    summary_type: str
    hook_direction: str
    primary: LayerSummaryAggregate
    reference: LayerSummaryAggregate
    sample_count_diff: NumericDiff
    mean_diff: NumericDiff
    std_diff: NumericDiff
    norm_diff: NumericDiff
    variance_diff: NumericDiff
    dead_unit_ratio_diff: NumericDiff
    saturation_ratio_diff: NumericDiff

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


@dataclass(frozen=True, slots=True)
class RunComparisonResult:
    """Structured API result for pairwise or baseline run comparison."""

    schema_version: int
    generated_at: str
    mode: ComparisonMode
    primary: ComparisonRunTarget
    reference: ComparisonRunTarget
    metric_diffs: tuple[MetricComparison, ...]
    alert_diff: AlertDiff
    phase_diffs: tuple[PhaseComparison, ...]
    layer_summary_diffs: tuple[LayerSummaryComparison, ...]

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


class ComparisonEngine:
    """Compare one run against another run or an explicit baseline run."""

    def __init__(
        self,
        session: Any,
        *,
        config: ComparisonEngineConfig | None = None,
    ) -> None:
        self._context, self._store = self._resolve_session_dependencies(session)
        self._config = config or ComparisonEngineConfig()
        self._metric_target_by_key = {
            target.metric_key: target for target in self._config.metric_targets
        }

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

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

    def compare(self, other: Any, *, baseline: bool = False) -> RunComparisonResult:
        """Compare the engine's run against another run and return structured diffs."""

        reference_context, reference_store = self._resolve_comparison_target(
            other,
            default_storage_dir=self._context.storage_dir,
        )

        primary_metrics = self._store.read_scalar_metrics(self.run_id)
        reference_metrics = reference_store.read_scalar_metrics(reference_context.run_id)

        primary_alerts = self._store.read_alerts(self.run_id)
        reference_alerts = reference_store.read_alerts(reference_context.run_id)

        primary_phases = self._store.read_phases(self.run_id)
        reference_phases = reference_store.read_phases(reference_context.run_id)

        primary_layer_summaries = self._store.read_layer_summaries(self.run_id)
        reference_layer_summaries = reference_store.read_layer_summaries(reference_context.run_id)

        return RunComparisonResult(
            schema_version=COMPARISON_SCHEMA_VERSION,
            generated_at=utc_now(),
            mode="baseline" if baseline else "pairwise",
            primary=self._build_run_target(
                self._context,
                is_baseline="baseline" in self._context.run.tags,
            ),
            reference=self._build_run_target(
                reference_context,
                is_baseline=baseline or "baseline" in reference_context.run.tags,
            ),
            metric_diffs=self._build_metric_diffs(primary_metrics, reference_metrics),
            alert_diff=self._build_alert_diff(primary_alerts, reference_alerts),
            phase_diffs=self._build_phase_diffs(primary_phases, reference_phases),
            layer_summary_diffs=self._build_layer_summary_diffs(
                primary_layer_summaries,
                reference_layer_summaries,
            ),
        )

    def compare_to_baseline(self, baseline_run: Any) -> RunComparisonResult:
        """Compare the engine's run against a baseline run."""

        return self.compare(baseline_run, baseline=True)

    @classmethod
    def compare_runs(
        cls,
        primary: Any,
        reference: Any,
        *,
        baseline: bool = False,
        config: ComparisonEngineConfig | None = None,
    ) -> RunComparisonResult:
        """Convenience helper for one-shot pairwise or baseline comparison."""

        return cls(primary, config=config).compare(reference, baseline=baseline)

    def _build_metric_diffs(
        self,
        primary_metrics: tuple[MetricRecord, ...],
        reference_metrics: tuple[MetricRecord, ...],
    ) -> tuple[MetricComparison, ...]:
        primary_by_key = self._group_metrics(primary_metrics)
        reference_by_key = self._group_metrics(reference_metrics)
        metric_keys = _ordered_union(primary_by_key.keys(), reference_by_key.keys())

        comparisons: list[MetricComparison] = []
        for metric_key in metric_keys:
            target_spec = self._metric_target_by_key.get(metric_key)
            primary_summary = self._build_metric_summary(
                primary_by_key.get(metric_key, ()),
                target_spec=target_spec,
            )
            reference_summary = self._build_metric_summary(
                reference_by_key.get(metric_key, ()),
                target_spec=target_spec,
            )
            comparisons.append(
                MetricComparison(
                    metric_key=metric_key,
                    primary=primary_summary,
                    reference=reference_summary,
                    latest_value_diff=self._build_numeric_diff(
                        primary_summary.latest_value,
                        reference_summary.latest_value,
                    ),
                    mean_value_diff=self._build_numeric_diff(
                        primary_summary.mean_value,
                        reference_summary.mean_value,
                    ),
                    min_value_diff=self._build_numeric_diff(
                        primary_summary.min_value,
                        reference_summary.min_value,
                    ),
                    max_value_diff=self._build_numeric_diff(
                        primary_summary.max_value,
                        reference_summary.max_value,
                    ),
                    net_change_diff=self._build_numeric_diff(
                        primary_summary.net_change,
                        reference_summary.net_change,
                    ),
                    sample_count_diff=self._build_numeric_diff(
                        primary_summary.sample_count,
                        reference_summary.sample_count,
                    ),
                    time_to_target_step_diff=(
                        self._build_numeric_diff(
                            primary_summary.time_to_target_step,
                            reference_summary.time_to_target_step,
                        )
                        if target_spec is not None
                        else None
                    ),
                )
            )
        return tuple(comparisons)

    def _build_alert_diff(
        self,
        primary_alerts: tuple[AlertRecord, ...],
        reference_alerts: tuple[AlertRecord, ...],
    ) -> AlertDiff:
        primary_severity_counts = self._build_fixed_count_map(
            Counter(alert.severity for alert in primary_alerts),
            _ALERT_SEVERITIES,
        )
        reference_severity_counts = self._build_fixed_count_map(
            Counter(alert.severity for alert in reference_alerts),
            _ALERT_SEVERITIES,
        )

        category_order = _ordered_union(
            _ALERT_CATEGORIES,
            _ordered_union(
                (alert.category for alert in primary_alerts),
                (alert.category for alert in reference_alerts),
            ),
        )
        primary_category_counts = self._build_fixed_count_map(
            Counter(alert.category for alert in primary_alerts),
            category_order,
        )
        reference_category_counts = self._build_fixed_count_map(
            Counter(alert.category for alert in reference_alerts),
            category_order,
        )

        primary_by_rule = self._group_alerts(primary_alerts)
        reference_by_rule = self._group_alerts(reference_alerts)
        rule_ids = _ordered_union(primary_by_rule.keys(), reference_by_rule.keys())

        rule_diffs = tuple(
            AlertRuleDiff(
                rule_id=rule_id,
                primary_category=(
                    primary_by_rule[rule_id][0].category if rule_id in primary_by_rule else None
                ),
                reference_category=(
                    reference_by_rule[rule_id][0].category if rule_id in reference_by_rule else None
                ),
                primary_count=len(primary_by_rule.get(rule_id, ())),
                reference_count=len(reference_by_rule.get(rule_id, ())),
                count_diff=(
                    len(primary_by_rule.get(rule_id, ()))
                    - len(reference_by_rule.get(rule_id, ()))
                ),
                primary_max_severity=self._max_alert_severity(primary_by_rule.get(rule_id, ())),
                reference_max_severity=self._max_alert_severity(reference_by_rule.get(rule_id, ())),
                severity_changed=(
                    self._max_alert_severity(primary_by_rule.get(rule_id, ()))
                    != self._max_alert_severity(reference_by_rule.get(rule_id, ()))
                ),
            )
            for rule_id in rule_ids
        )

        return AlertDiff(
            primary_total_count=len(primary_alerts),
            reference_total_count=len(reference_alerts),
            total_count_diff=len(primary_alerts) - len(reference_alerts),
            primary_severity_counts=primary_severity_counts,
            reference_severity_counts=reference_severity_counts,
            severity_count_diffs=self._diff_count_maps(
                primary_severity_counts,
                reference_severity_counts,
            ),
            primary_category_counts=primary_category_counts,
            reference_category_counts=reference_category_counts,
            category_count_diffs=self._diff_count_maps(
                primary_category_counts,
                reference_category_counts,
            ),
            rule_diffs=rule_diffs,
        )

    def _build_phase_diffs(
        self,
        primary_phases: tuple[PhaseRecord, ...],
        reference_phases: tuple[PhaseRecord, ...],
    ) -> tuple[PhaseComparison, ...]:
        primary_by_name = self._group_phases(primary_phases)
        reference_by_name = self._group_phases(reference_phases)
        phase_names = _ordered_union(primary_by_name.keys(), reference_by_name.keys())

        comparisons: list[PhaseComparison] = []
        for phase_name in phase_names:
            primary_summary = self._build_phase_aggregate(primary_by_name.get(phase_name, ()))
            reference_summary = self._build_phase_aggregate(reference_by_name.get(phase_name, ()))
            comparisons.append(
                PhaseComparison(
                    phase_name=phase_name,
                    primary=primary_summary,
                    reference=reference_summary,
                    segment_count_diff=(
                        primary_summary.segment_count - reference_summary.segment_count
                    ),
                    total_duration_diff=(
                        primary_summary.total_duration_steps - reference_summary.total_duration_steps
                    ),
                )
            )
        return tuple(comparisons)

    def _build_layer_summary_diffs(
        self,
        primary_layer_summaries: tuple[LayerSummaryRecord, ...],
        reference_layer_summaries: tuple[LayerSummaryRecord, ...],
    ) -> tuple[LayerSummaryComparison, ...]:
        primary_by_key = self._group_layer_summaries(primary_layer_summaries)
        reference_by_key = self._group_layer_summaries(reference_layer_summaries)
        ordered_keys = _ordered_union(primary_by_key.keys(), reference_by_key.keys())

        comparisons: list[LayerSummaryComparison] = []
        for layer_key in ordered_keys:
            module_name, module_type, summary_type, hook_direction = layer_key
            primary_summary = self._build_layer_summary_aggregate(primary_by_key.get(layer_key, ()))
            reference_summary = self._build_layer_summary_aggregate(
                reference_by_key.get(layer_key, ())
            )
            comparisons.append(
                LayerSummaryComparison(
                    module_name=module_name,
                    module_type=module_type,
                    summary_type=summary_type,
                    hook_direction=hook_direction,
                    primary=primary_summary,
                    reference=reference_summary,
                    sample_count_diff=self._build_numeric_diff(
                        primary_summary.sample_count,
                        reference_summary.sample_count,
                    ),
                    mean_diff=self._build_numeric_diff(
                        primary_summary.average_mean,
                        reference_summary.average_mean,
                    ),
                    std_diff=self._build_numeric_diff(
                        primary_summary.average_std,
                        reference_summary.average_std,
                    ),
                    norm_diff=self._build_numeric_diff(
                        primary_summary.average_norm,
                        reference_summary.average_norm,
                    ),
                    variance_diff=self._build_numeric_diff(
                        primary_summary.average_variance,
                        reference_summary.average_variance,
                    ),
                    dead_unit_ratio_diff=self._build_numeric_diff(
                        primary_summary.average_dead_unit_ratio,
                        reference_summary.average_dead_unit_ratio,
                    ),
                    saturation_ratio_diff=self._build_numeric_diff(
                        primary_summary.average_saturation_ratio,
                        reference_summary.average_saturation_ratio,
                    ),
                )
            )
        return tuple(comparisons)

    @staticmethod
    def _group_metrics(metrics: tuple[MetricRecord, ...]) -> dict[str, tuple[MetricRecord, ...]]:
        grouped: dict[str, list[MetricRecord]] = {}
        for metric in metrics:
            grouped.setdefault(metric.metric_key, []).append(metric)
        return {metric_key: tuple(records) for metric_key, records in grouped.items()}

    @staticmethod
    def _group_alerts(alerts: tuple[AlertRecord, ...]) -> dict[str, tuple[AlertRecord, ...]]:
        grouped: dict[str, list[AlertRecord]] = {}
        for alert in alerts:
            grouped.setdefault(alert.rule_id, []).append(alert)
        return {rule_id: tuple(records) for rule_id, records in grouped.items()}

    @staticmethod
    def _group_phases(phases: tuple[PhaseRecord, ...]) -> dict[str, tuple[PhaseRecord, ...]]:
        grouped: dict[str, list[PhaseRecord]] = {}
        for phase in phases:
            grouped.setdefault(phase.phase_name, []).append(phase)
        return {phase_name: tuple(records) for phase_name, records in grouped.items()}

    @staticmethod
    def _group_layer_summaries(
        layer_summaries: tuple[LayerSummaryRecord, ...],
    ) -> dict[tuple[str, str, str, str], tuple[LayerSummaryRecord, ...]]:
        grouped: dict[tuple[str, str, str, str], list[LayerSummaryRecord]] = {}
        for summary in layer_summaries:
            key = (
                summary.module_name,
                summary.module_type,
                summary.summary_type,
                summary.hook_direction,
            )
            grouped.setdefault(key, []).append(summary)
        return {layer_key: tuple(records) for layer_key, records in grouped.items()}

    def _build_metric_summary(
        self,
        records: tuple[MetricRecord, ...],
        *,
        target_spec: MetricTargetSpec | None,
    ) -> MetricSeriesSummary:
        if not records:
            return MetricSeriesSummary(
                sample_count=0,
                step_start=None,
                step_end=None,
                first_value=None,
                latest_value=None,
                min_value=None,
                max_value=None,
                mean_value=None,
                net_change=None,
                target=target_spec,
                time_to_target_step=None,
            )

        values = [record.value for record in records]
        steps = [record.step for record in records if record.step is not None]
        first_value = records[0].value
        latest_value = records[-1].value

        return MetricSeriesSummary(
            sample_count=len(records),
            step_start=min(steps) if steps else None,
            step_end=max(steps) if steps else None,
            first_value=first_value,
            latest_value=latest_value,
            min_value=min(values),
            max_value=max(values),
            mean_value=sum(values) / len(values),
            net_change=latest_value - first_value,
            target=target_spec,
            time_to_target_step=self._find_time_to_target_step(records, target_spec=target_spec),
        )

    @staticmethod
    def _build_phase_aggregate(records: tuple[PhaseRecord, ...]) -> PhaseAggregate:
        if not records:
            return PhaseAggregate(
                display_name=None,
                segment_count=0,
                total_duration_steps=0,
                first_step=None,
                last_step=None,
            )

        durations = [record.step_end - record.step_start + 1 for record in records]
        return PhaseAggregate(
            display_name=records[0].display_name,
            segment_count=len(records),
            total_duration_steps=sum(durations),
            first_step=min(record.step_start for record in records),
            last_step=max(record.step_end for record in records),
        )

    @staticmethod
    def _build_layer_summary_aggregate(
        records: tuple[LayerSummaryRecord, ...],
    ) -> LayerSummaryAggregate:
        if not records:
            return LayerSummaryAggregate(
                sample_count=0,
                step_start=None,
                step_end=None,
                average_mean=None,
                average_std=None,
                average_norm=None,
                average_variance=None,
                average_dead_unit_ratio=None,
                average_saturation_ratio=None,
            )

        return LayerSummaryAggregate(
            sample_count=len(records),
            step_start=min(record.step for record in records),
            step_end=max(record.step for record in records),
            average_mean=ComparisonEngine._mean_optional(record.mean for record in records),
            average_std=ComparisonEngine._mean_optional(record.std for record in records),
            average_norm=ComparisonEngine._mean_optional(record.norm for record in records),
            average_variance=ComparisonEngine._mean_optional(record.variance for record in records),
            average_dead_unit_ratio=ComparisonEngine._mean_optional(
                record.dead_unit_ratio for record in records
            ),
            average_saturation_ratio=ComparisonEngine._mean_optional(
                record.saturation_ratio for record in records
            ),
        )

    @staticmethod
    def _build_numeric_diff(
        primary_value: float | int | None,
        reference_value: float | int | None,
    ) -> NumericDiff:
        if primary_value is None or reference_value is None:
            return NumericDiff(
                primary=primary_value,
                reference=reference_value,
                delta=None,
                percent_delta=None,
            )

        delta = primary_value - reference_value
        percent_delta = None
        if reference_value != 0:
            percent_delta = (delta / abs(reference_value)) * 100.0
        return NumericDiff(
            primary=primary_value,
            reference=reference_value,
            delta=delta,
            percent_delta=percent_delta,
        )

    @staticmethod
    def _mean_optional(values: Iterable[float | None]) -> float | None:
        materialized = [float(value) for value in values if value is not None]
        if not materialized:
            return None
        return sum(materialized) / len(materialized)

    @staticmethod
    def _find_time_to_target_step(
        records: tuple[MetricRecord, ...],
        *,
        target_spec: MetricTargetSpec | None,
    ) -> int | None:
        if target_spec is None:
            return None
        for record in records:
            if record.step is None:
                continue
            if target_spec.direction == "at_least" and record.value >= target_spec.target_value:
                return record.step
            if target_spec.direction == "at_most" and record.value <= target_spec.target_value:
                return record.step
        return None

    @staticmethod
    def _max_alert_severity(alerts: tuple[AlertRecord, ...]) -> str | None:
        if not alerts:
            return None
        return max(alerts, key=lambda alert: _ALERT_SEVERITY_ORDER[alert.severity]).severity

    @staticmethod
    def _build_fixed_count_map(
        counts: Counter[str],
        ordered_keys: Iterable[str],
    ) -> dict[str, int]:
        return {key: int(counts.get(key, 0)) for key in ordered_keys}

    @staticmethod
    def _diff_count_maps(
        primary_counts: dict[str, int],
        reference_counts: dict[str, int],
    ) -> dict[str, int]:
        return {
            key: int(primary_counts.get(key, 0) - reference_counts.get(key, 0))
            for key in _ordered_union(primary_counts.keys(), reference_counts.keys())
        }

    @staticmethod
    def _build_run_target(
        context: RunContext,
        *,
        is_baseline: bool = False,
    ) -> ComparisonRunTarget:
        return ComparisonRunTarget(
            run_id=context.run_id,
            project_name=context.project_name,
            status=context.status,
            tags=context.run.tags,
            is_baseline=is_baseline,
        )

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

    @staticmethod
    def _resolve_comparison_target(
        target: Any,
        *,
        default_storage_dir: str,
    ) -> tuple[RunContext, LocalSessionStore]:
        if isinstance(target, str):
            store = LocalSessionStore(default_storage_dir)
            return store.load_run(target), store
        return ComparisonEngine._resolve_session_dependencies(target)


__all__ = [
    "COMPARISON_SCHEMA_VERSION",
    "AlertDiff",
    "AlertRuleDiff",
    "ComparisonEngine",
    "ComparisonEngineConfig",
    "ComparisonMode",
    "ComparisonRunTarget",
    "LayerSummaryAggregate",
    "LayerSummaryComparison",
    "MetricComparison",
    "MetricSeriesSummary",
    "MetricTargetDirection",
    "MetricTargetSpec",
    "NumericDiff",
    "PhaseAggregate",
    "PhaseComparison",
    "RunComparisonResult",
]
