"""Sequential, resumable execution of an immutable paired trial plan."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Callable, Mapping

from .io import (
    AtomicTrialWriter,
    TrialPayload,
    atomic_write_json,
    record_trial_failure,
    utc_now,
)
from .manifest import ensure_batch
from .plan import validate_trial_plan
from .validate import validate_trial_directory


TrialExecutor = Callable[[Mapping[str, Any]], TrialPayload | Mapping[str, Any]]


def _normalize_payload(value: TrialPayload | Mapping[str, Any]) -> TrialPayload:
    if isinstance(value, TrialPayload):
        return value
    if not isinstance(value, Mapping) or "samples" not in value:
        raise TypeError("executor must return TrialPayload or a mapping with samples")
    return TrialPayload(
        samples=value["samples"],
        events=value.get("events", ()),
        metadata=value.get("metadata", {}),
    )


def run_trial_plan(
    plan: Mapping[str, Any],
    batch_dir: Path,
    executor: TrialExecutor | None = None,
    *,
    resume: bool = True,
    dry_run: bool = False,
    continue_on_error: bool = True,
    max_trials: int | None = None,
    source_root: Path | None = None,
) -> dict[str, Any]:
    """Run trials in plan order and atomically commit successful artifacts."""
    validate_trial_plan(plan)
    selected = list(plan["trials"])
    if max_trials is not None:
        if max_trials < 0:
            raise ValueError("max_trials must be non-negative")
        selected = selected[:max_trials]
    if dry_run:
        return {
            "dry_run": True,
            "study_id": plan["study_id"],
            "split": plan["split"],
            "pair_count": len({trial["pair_id"] for trial in selected}),
            "trial_count": len(selected),
            "plan_hash": plan["plan_hash"],
        }
    if executor is None:
        raise ValueError("executor is required unless dry_run=True")

    batch_dir = Path(batch_dir)
    ensure_batch(batch_dir, plan, source_root=source_root)
    writer = AtomicTrialWriter(batch_dir)
    completed = 0
    skipped = 0
    failed = 0
    failures: list[dict[str, str]] = []
    for trial in selected:
        destination = writer.trial_dir(trial["trial_id"])
        if destination.exists():
            validation_errors = validate_trial_directory(destination, trial)
            if resume and not validation_errors:
                skipped += 1
                continue
            if validation_errors:
                raise RuntimeError(
                    f"existing trial {trial['trial_id']} is invalid: "
                    + "; ".join(validation_errors)
                )
            raise FileExistsError(f"trial already completed: {trial['trial_id']}")
        try:
            payload = _normalize_payload(executor(trial))
            writer.write(trial, payload)
            completed += 1
        except Exception as exc:
            failed += 1
            record_trial_failure(batch_dir, trial, exc)
            failures.append(
                {
                    "trial_id": trial["trial_id"],
                    "error_type": type(exc).__name__,
                    "error_message": str(exc),
                }
            )
            if not continue_on_error:
                raise
    summary = {
        "kind": "run_summary",
        "schema_version": plan["schema_version"],
        "plan_hash": plan["plan_hash"],
        "recorded_utc": utc_now(),
        "selected_trials": len(selected),
        "completed_trials": completed,
        "skipped_trials": skipped,
        "failed_trials": failed,
        "failures": failures,
    }
    atomic_write_json(batch_dir / "run_summary.json", summary)
    return summary