from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from src.configuration import changed_configuration_fields
from src.models import Result


@dataclass(frozen=True)
class MetricDelta:
    baseline: float | None
    treatment: float | None
    absolute: float | None
    percent: float | None


@dataclass(frozen=True)
class Comparison:
    baseline: Result
    treatment: Result
    changed_fields: tuple[str, ...]
    matched_dimensions: tuple[str, ...]
    unmatched_dimensions: tuple[str, ...]
    status: str
    metric_deltas: dict[str, MetricDelta]


def safe_percent_change(baseline: float | int | None, treatment: float | int | None) -> float | None:
    if baseline is None or treatment is None or baseline == 0:
        return None
    return (float(treatment) - float(baseline)) / abs(float(baseline)) * 100.0


def metric_delta(baseline: float | int | None, treatment: float | int | None) -> MetricDelta:
    absolute = None if baseline is None or treatment is None else float(treatment) - float(baseline)
    return MetricDelta(
        baseline=None if baseline is None else float(baseline),
        treatment=None if treatment is None else float(treatment),
        absolute=absolute,
        percent=safe_percent_change(baseline, treatment),
    )


def _error_rate(result: Result) -> float | None:
    tasks = result.metrics.n_tasks
    errors = result.metrics.n_errors
    if tasks is None or errors is None or tasks <= 0:
        return None
    return errors / tasks * 100.0


def _per_success(value: float | int | None, result: Result) -> float | None:
    if value is None or result.metrics.score <= 0:
        return None
    return float(value) / result.metrics.score


def comparison_metric_deltas(baseline: Result, treatment: Result) -> dict[str, MetricDelta]:
    bm = baseline.metrics
    tm = treatment.metrics
    return {
        # Score absolute delta is percentage points; relative delta remains percent.
        "score_percentage_points": metric_delta(bm.score * 100.0, tm.score * 100.0),
        "total_tokens": metric_delta(bm.n_total_tokens, tm.n_total_tokens),
        "cost_usd": metric_delta(bm.cost_usd, tm.cost_usd),
        "total_time_seconds": metric_delta(bm.total_time_seconds, tm.total_time_seconds),
        "agent_time_seconds": metric_delta(bm.agent_time_seconds, tm.agent_time_seconds),
        "execution_error_rate_percentage_points": metric_delta(_error_rate(baseline), _error_rate(treatment)),
        "tokens_per_successful_task": metric_delta(
            _per_success(bm.mean_tokens_per_task, baseline),
            _per_success(tm.mean_tokens_per_task, treatment),
        ),
        "cost_per_successful_task": metric_delta(
            _per_success(bm.mean_cost_usd_per_task, baseline),
            _per_success(tm.mean_cost_usd_per_task, treatment),
        ),
        "time_per_successful_task": metric_delta(
            _per_success(bm.mean_total_time_seconds_per_task, baseline),
            _per_success(tm.mean_total_time_seconds_per_task, treatment),
        ),
    }


def _dimension_values(result: Result) -> dict[str, Any]:
    # Use only metadata represented by the current schema. Future schema additions
    # (benchmark revision, budgets, seed, protocol, harness version) can extend this
    # mapping without changing the comparison API.
    return {
        "benchmark": result.benchmark.name,
        "model": result.model.repo or result.model.name,
        "harness": result.harness.name,
        "environment": result.environment.name,
        "environment_config": result.environment.config,
    }


def compare_runs(baseline: Result, treatment: Result) -> Comparison:
    changed = tuple(changed_configuration_fields(baseline.configuration, treatment.configuration))
    left = _dimension_values(baseline)
    right = _dimension_values(treatment)
    matched = tuple(key for key in left if left[key] is not None and left[key] == right[key])
    unmatched = tuple(key for key in left if left[key] != right[key])

    if baseline.configuration is None or treatment.configuration is None:
        status = "observational"
    elif not changed:
        status = "ambiguous"
    elif not unmatched:
        status = "exact"
    elif {"benchmark", "model"}.issubset(matched):
        status = "partial"
    else:
        status = "unmatched"

    return Comparison(
        baseline=baseline,
        treatment=treatment,
        changed_fields=changed,
        matched_dimensions=matched,
        unmatched_dimensions=unmatched,
        status=status,
        metric_deltas=comparison_metric_deltas(baseline, treatment),
    )