"""One trial: checkout, swap config, invoke the agent, capture, assert, clean up.

The unit the scheduler fans out over (task x variant x replicate).
"""

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path

from ablate.agents.base import AgentAdapter, TrialSpec
from ablate.assertions.run import run_assertions
from ablate.config.variant import ConfigVariant
from ablate.gitutil import diff_text
from ablate.record import AssertionRecord, TelemetryRecord, TrialRecord
from ablate.runner.worktree import make_worktree
from ablate.task import Task


@dataclass(frozen=True)
class TrialJob:
    """One cell of the (task x variant x replicate) design."""

    task: Task
    variant: ConfigVariant
    replicate: int

    @property
    def trial_id(self) -> str:
        return f"{self.task.id}-{self.variant.id}-r{self.replicate}"




def run_trial(
    job: TrialJob,
    repo: Path,
    worktrees_root: Path,
    adapter: AgentAdapter,
    run_id: str,
    max_turns: int = 40,
    keep_worktrees: bool = False,
    model: str | None = None,
) -> TrialRecord:
    task = job.task
    base_commit = task.base_commit or "HEAD"
    started_at = datetime.now(UTC).isoformat()
    t0 = time.monotonic()

    with make_worktree(
        repo=repo,
        trial_id=job.trial_id,
        base_commit=base_commit,
        worktrees_root=worktrees_root,
        setup=task.setup,
        keep=keep_worktrees,
    ) as wt:
        job.variant.apply(wt.path, task.config_path)

        spec = TrialSpec(
            prompt=task.prompt,
            cwd=wt.path,
            config_variant_text=job.variant.text,
            config_path=task.config_path,
            max_turns=max_turns,
            timeout_s=task.timeout_s,
            model=model,
        )
        telemetry = adapter.invoke(spec)

        # task.hide (files re-applied after the agent runs, so a mined task's
        # answer is never visible during the trial) is unused until ablate
        # mine ships in M5 — no hand-authored M1 task sets it.

        diff = diff_text(wt.path, base_commit)
        diff_sha = "sha256:" + hashlib.sha256(diff.encode()).hexdigest()[:16]

        results, passed = run_assertions(task, wt.path)
        if telemetry.error:
            passed = False

    duration_s = time.monotonic() - t0

    return TrialRecord(
        trial_id=job.trial_id,
        run_id=run_id,
        task_id=task.id,
        variant_id=job.variant.id,
        variant_hash=job.variant.content_hash,
        replicate=job.replicate,
        agent=adapter.name,
        # What ran, not what was asked for, whenever the agent says.
        model=telemetry.model or model,
        requested_model=model,
        started_at=started_at,
        duration_s=duration_s,
        telemetry=TelemetryRecord(
            input_tokens=telemetry.input_tokens,
            output_tokens=telemetry.output_tokens,
            cache_read=telemetry.cache_read,
            cache_write=telemetry.cache_write,
            cost_usd=telemetry.cost_usd,
            turns=telemetry.turns,
            tool_calls=telemetry.tool_calls,
        ),
        assertions=[
            AssertionRecord(
                id=r.id, type=r.type, passed=r.passed, required=r.required,
                detail=r.detail, value=r.value,
            )
            for r in results
        ],
        passed=passed,
        diff_sha=diff_sha,
        session_id=telemetry.session_id,
        error=telemetry.error,
    )
