"""Expand (task x variant x replicate) into a shuffled queue and run it
under a concurrency semaphore.

Shuffled so a mid-run abort or --max-spend stop doesn't systematically
starve one arm — a partial run must still be a roughly balanced sample of
the design, which is what makes it analyzable (docs/03-architecture.md).
"""

from __future__ import annotations

import random
from collections.abc import Callable, Iterable
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

from ablate.agents.base import AgentAdapter
from ablate.config.variant import ConfigVariant
from ablate.record import TrialRecord
from ablate.runner.trial import TrialJob, run_trial
from ablate.task import Task


def build_jobs(
    tasks: Iterable[Task], variants: Iterable[ConfigVariant], trials: int, seed: int | None = None
) -> list[TrialJob]:
    jobs = [
        TrialJob(task=t, variant=v, replicate=r)
        for t in tasks
        for v in variants
        for r in range(1, trials + 1)
    ]
    random.Random(seed).shuffle(jobs)
    return jobs


#: Consecutive errored trials, from the very start of a run, that mean the
#: problem is the setup rather than the agent's performance on a task. Set
#: above the concurrency default so a single unlucky batch cannot trip it, and
#: low enough that a misconfigured run costs pennies instead of the whole
#: budget. Only counted from the start: a run that produced real results and
#: then hits a rate limit is a partial design, which is analyzable and handled
#: by `should_stop`.
FAILFAST_THRESHOLD = 6


def run_jobs(
    jobs: list[TrialJob],
    repo: Path,
    worktrees_root: Path,
    adapter: AgentAdapter,
    run_id: str,
    concurrency: int = 4,
    max_turns: int = 40,
    keep_worktrees: bool = False,
    model: str | None = None,
    on_result: Callable[[TrialRecord], None] | None = None,
    should_stop: Callable[[], bool] | None = None,
    on_abort: Callable[[str], None] | None = None,
) -> tuple[list[TrialRecord], bool]:
    """Runs `jobs` under a `concurrency`-wide thread pool.

    `on_result` fires as each trial completes, in completion order — used to
    stream progress and append to the log incrementally rather than buffering
    the whole run in memory.

    `should_stop` is polled after each completion; when it returns True the
    queue is drained and no further trials start. This is how `--max-spend`
    is enforced against real costs mid-run rather than only against the
    pre-run estimate. Trials already in flight are allowed to finish — killing
    an agent mid-edit would leak a worktree and waste money already spent.

    `on_abort` fires when the run is stopped because nothing is succeeding —
    see FAILFAST_THRESHOLD. An expired key, an exhausted quota or a model id
    the provider does not serve fails every trial identically; without this the
    queue ran to completion and the failure arrived as a report rather than as
    an error.

    Returns the records collected and whether the run was stopped early.
    """
    records: list[TrialRecord] = []
    stopped = False
    consecutive_errors = 0
    any_success = False

    with ThreadPoolExecutor(max_workers=max(1, concurrency)) as pool:
        futures = {
            pool.submit(
                run_trial,
                job=job,
                repo=repo,
                worktrees_root=worktrees_root,
                adapter=adapter,
                run_id=run_id,
                max_turns=max_turns,
                keep_worktrees=keep_worktrees,
                model=model,
            ): job
            for job in jobs
        }
        consumed = set()
        for future in as_completed(futures):
            consumed.add(future)
            record = future.result()
            records.append(record)
            if on_result:
                on_result(record)

            if record.error is None:
                any_success = True
                consecutive_errors = 0
            else:
                consecutive_errors += 1

            # Only before anything has worked. Once one trial has succeeded the
            # setup is proven, and later failures are data about the tasks.
            if not any_success and consecutive_errors >= FAILFAST_THRESHOLD:
                stopped = True
                if on_abort:
                    on_abort(_abort_reason(records))
                for pending in futures:
                    pending.cancel()
                break

            if should_stop is not None and should_stop():
                stopped = True
                for pending in futures:
                    pending.cancel()
                break

        # Trials already in flight when the stop fired still finish, and their
        # results are kept: they have been paid for, and because the queue was
        # shuffled they extend a roughly balanced partial design rather than
        # skewing one arm.
        for future in futures:
            if future in consumed or future.cancelled():
                continue
            try:
                record = future.result()
            except Exception:  # noqa: BLE001 — a lost straggler must not mask the stop
                continue
            records.append(record)
            if on_result:
                on_result(record)

    return records, stopped


def _abort_reason(records: list[TrialRecord]) -> str:
    """The agent's own error, quoted back.

    The useful diagnostic is almost always in the first error verbatim
    ("Not logged in", "credit balance too low", "model not found"). Quoting it
    beats any message this layer could compose, because this layer does not
    know which of those it is.
    """
    first = next((r.error for r in records if r.error), "")
    detail = " ".join(first.split())[:300]
    return (
        f"{len(records)} trials ran and every one failed. Stopping rather than "
        f"spending the rest of the budget on the same error.\n  First error: {detail}"
    )
