from __future__ import annotations

import hashlib
import subprocess
import threading
import time
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Protocol
from uuid import uuid4

from .adapters.base import AgentAdapter, AgentRequest, AgentResult
from .checkpoint import Checkpoint
from .conversation import Conversation, render
from .domain import (
    AgentTuning,
    Agents,
    ErrorCategory,
    FixArtifact,
    Limits,
    PlanArtifact,
    ReviewArtifact,
    TaskRecord,
    TaskSpec,
    TaskState,
    StageTestReport,
    is_terminal,
)
from .infra.process import is_running, terminate_group
from .locking import TaskBusy, task_lock
from .memory import ProjectMemory
from .parallel import MergeConflict, ParallelImplementer, unsafe_to_parallelise
from .persistence import TaskStore
from .retry import RetryPolicy
from .state_machine import InvalidTransition, next_state


class RepositoryNotAllowed(ValueError):
    pass


class WorkspaceProvider(Protocol):
    def create(self, task_id: str, repo_path: Path) -> tuple[Path, str]: ...


TestRunner = Callable[[TaskRecord, Path], "StageTestReport | bool"]

#: How many times an operator may send a plan back before the task stops on its own.
MAX_OPERATOR_REPLANS = 2

#: Stages that write to the workspace, and so may be put behind an operator approval.
WRITING_STAGES = {TaskState.IMPLEMENT, TaskState.FIX}

#: Which adapter owns each agent stage. Codex plans and reviews; Claude reviews the plan, implements and fixes.
PLANNER_STAGES = {TaskState.PLAN, TaskState.REVIEW, TaskState.FINAL_VERIFY}

#: Artifact type each agent stage must return before the state machine will move.
STAGE_SCHEMA = {
    TaskState.PLAN: PlanArtifact,
    TaskState.PLAN_REVIEW: ReviewArtifact,
    TaskState.REVIEW: ReviewArtifact,
    TaskState.FINAL_VERIFY: ReviewArtifact,
    TaskState.FIX: FixArtifact,
}


def _category(status: str) -> ErrorCategory:
    try:
        return ErrorCategory(status)
    except ValueError:
        return ErrorCategory.AGENT_FAILURE


class Orchestrator:
    def __init__(
        self,
        store: TaskStore,
        planner: AgentAdapter,
        implementer: AgentAdapter,
        workspace: WorkspaceProvider,
        test_runner: TestRunner,
        retry_policy: RetryPolicy | None = None,
        require_approval: bool = False,
        memory: ProjectMemory | None = None,
        allowed_repo_roots: list[str] | None = None,
        implement_workers: int = 1,
        lock_timeout_s: float = 0.0,
        limits: Limits | None = None,
        discussion_rounds: int = 0,
        adapter_factory: Callable[[str, AgentTuning], AgentAdapter] | None = None,
        agents: Agents | None = None,
    ) -> None:
        self.store = store
        self.planner = planner
        self.implementer = implementer
        #: Admin-maintained model/effort catalog, checked before any tuning is trusted.
        self.agents_config = agents or Agents()
        self.workspace = workspace
        self.test_runner = test_runner
        self.retry_policy = retry_policy or RetryPolicy()
        self.require_approval = require_approval
        self.allowed_repo_roots = [Path(root) for root in (allowed_repo_roots or [])]
        #: More than one splits an approved plan across parallel worktrees.
        self.implement_workers = max(1, implement_workers)
        #: How long advance waits for another process to release the task before refusing.
        self.lock_timeout_s = lock_timeout_s
        #: Configured limits. Without this every task silently used the schema defaults.
        self.limits = limits or Limits()
        #: Planner/reviewer exchanges allowed over a rejected plan. 0 keeps them apart.
        self.discussion_rounds = max(0, discussion_rounds)
        #: Builds a per-task adapter when the operator chose a different model or effort.
        self.adapter_factory = adapter_factory
        self.memory = memory or ProjectMemory(store.root)
        #: One advance per task at a time. Cancel deliberately stays outside this lock
        #: so an operator can always stop a stage that is already running.
        self._locks: dict[str, threading.Lock] = defaultdict(threading.Lock)
        self._locks_guard = threading.Lock()
        #: Parallel subtasks report their pids concurrently; without this the
        #: read-modify-write loses one and cancel then misses that process.
        self._pid_guard = threading.Lock()

    # ---------------------------------------------------------------- lifecycle

    def _check_repo(self, repo_path: Path) -> Path:
        """Confine tasks to configured roots so a UI caller cannot reach any repo on the host."""
        resolved = Path(repo_path).resolve(strict=False)
        if not self.allowed_repo_roots:
            return resolved
        for root in self.allowed_repo_roots:
            base = root.resolve(strict=False)
            if resolved == base or base in resolved.parents:
                return resolved
        allowed = ", ".join(str(root) for root in self.allowed_repo_roots)
        raise RepositoryNotAllowed(f"{resolved} is outside the allowed roots: {allowed}")

    def _validate_tuning(self, tuning: dict[str, AgentTuning]) -> None:
        """Check every chosen model and effort against the provider's catalog.

        Called before anything the tuning could reach: creating a task's worktree,
        building a per-task adapter, or resuming a discussion.
        """
        for role in ("planner", "implementer"):
            chosen = tuning.get(role)
            if chosen:
                self.agents_config.validate_tuning(role, chosen)

    def create_task(self, spec: TaskSpec) -> TaskRecord:
        spec = spec.model_copy(update={"repo_path": self._check_repo(spec.repo_path)})
        self._validate_tuning(spec.tuning)
        task_id = uuid4().hex
        worktree_path, baseline = self.workspace.create(task_id, spec.repo_path)
        record = TaskRecord.new(
            task_id,
            spec.repo_path,
            spec.goal,
            constraints=spec.constraints,
            worktree_path=worktree_path,
            baseline_commit=baseline,
            dirty_files=self._dirty_files(spec.repo_path),
            require_approval=self.require_approval,
            limits=self.limits,
            tuning=spec.tuning,
        )
        self.store.create(record)
        return record

    def get_task(self, task_id: str) -> TaskRecord:
        return self.store.load(task_id)

    def conversation(self, task_id: str) -> Conversation:
        return Conversation(self.store.task_dir(task_id))

    def say(self, task_id: str, text: str) -> TaskRecord:
        """Record what the operator wants and act on it where it changes the next step.

        Speaking while a plan is waiting for approval is how you reject it: the task goes
        back to planning carrying what you said. At any other point the message is kept
        and handed to whichever agent runs next.
        """
        with self._lock_for(task_id), task_lock(self.store.task_dir(task_id), self.lock_timeout_s):
            task = self.store.load(task_id)
            if not text.strip():
                raise ValueError("an empty message carries nothing")
            self.conversation(task_id).add("operator", text, stage=task.state.value)
            self.store.append_event(task_id, {"event": "operator_message", "state": task.state.value})

            if self.awaiting_approval(task) and task.state is TaskState.IMPLEMENT:
                task = self._save(task, plan_rounds=task.plan_rounds + 1, discussion_rounds=0)
                if task.plan_rounds > MAX_OPERATOR_REPLANS:
                    return self._needs_human(
                        task, f"the plan was sent back {task.plan_rounds} times", ErrorCategory.STALLED
                    )
                return self._transition(task, "replan_requested")
            return task

    def cancel(self, task_id: str) -> TaskRecord:
        """Stop the task and every agent it has running.

        Holding the pid guard for the whole operation closes the start/cancel race: a
        subtask that is still launching either registers before this runs and is killed
        here, or finds the task already terminal and kills itself in `_record_pid`.
        """
        with self._pid_guard:
            task = self.store.load(task_id)
            killed = []
            for pid in self._recorded_pids(task):
                if is_running(pid) and terminate_group(pid):
                    killed.append(pid)
                    self.store.append_event(
                        task.task_id, {"event": "run_terminated", "pid": pid, "signalled": True}
                    )
            reason = "cancelled by operator" + (
                f" ({len(killed)} agent(s) terminated)" if killed else ""
            )
            # Persisted inside the lock so any launching subtask observes the terminal state.
            return self._needs_human(task, reason, ErrorCategory.OPERATOR)

    def _record_pid(self, task_id: str, pid: int) -> None:
        with self._pid_guard:
            record = self.store.load(task_id)
            if is_terminal(record.state):
                # Cancelled while this provider was starting. Kill it now rather than
                # registering a process nobody will come back for.
                terminate_group(pid)
                self.store.append_event(
                    task_id, {"event": "run_terminated", "pid": pid, "signalled": True, "late_start": True}
                )
                return
            pids = [existing for existing in record.active_run_pids if existing != pid] + [pid]
            self.store.save(
                record.model_copy(update={"active_run_pid": pid, "active_run_pids": pids})
            )

    @staticmethod
    def _recorded_pids(task: TaskRecord) -> list[int]:
        """Every pid the task ever registered. Parallel stages start several at once."""
        known = [*task.active_run_pids, *([task.active_run_pid] if task.active_run_pid else [])]
        return list(dict.fromkeys(known))

    # Mapping: last_successful_stage value → state to resume from after recovery.
    # NOTE: FINAL_VERIFY maps back to FINAL_VERIFY (not DONE) because
    # last_successful_stage is written when the stage artifact is saved, before the
    # verdict is applied.  A failed FINAL_VERIFY (CHANGES_REQUIRED) that exhausts
    # fix cycles enters NEEDS_HUMAN with last_successful_stage=FINAL_VERIFY; resuming
    # must re-run the verification rather than skip it and mark the task complete.
    _RESUME_STATE: dict[str | None, TaskState] = {
        None: TaskState.INIT,
        TaskState.INIT.value: TaskState.INIT,
        TaskState.PLAN.value: TaskState.PLAN_REVIEW,
        TaskState.PLAN_REVIEW.value: TaskState.IMPLEMENT,
        TaskState.IMPLEMENT.value: TaskState.TEST,
        TaskState.TEST.value: TaskState.REVIEW,
        TaskState.REVIEW.value: TaskState.FINAL_VERIFY,
        TaskState.FIX.value: TaskState.TEST,
        TaskState.FINAL_VERIFY.value: TaskState.FINAL_VERIFY,
    }

    def recover(self, task_id: str) -> TaskRecord:
        """Resume a task, including NEEDS_HUMAN tasks where recovery is possible.

        Operator-cancelled tasks (ErrorCategory.OPERATOR) and DONE tasks are
        not recoverable and raise InvalidTransition / return as-is.
        """
        task = self.store.load(task_id)
        if task.state is TaskState.DONE:
            return task
        if task.state is TaskState.NEEDS_HUMAN:
            if task.error_category is ErrorCategory.OPERATOR:
                raise InvalidTransition(
                    "task was cancelled by the operator; create a new task to retry"
                )
            # Fall through to workspace checks; recovery will clear terminal state.
        if task.active_run_id:
            alive = [pid for pid in self._recorded_pids(task) if is_running(pid)]
            if alive:
                # Work is still in flight. Moving to a terminal state here would clear the
                # run record and orphan those processes: they would keep editing the
                # worktree with nothing left able to name, let alone stop, them. Leave the
                # task exactly as it is so `cancel` can still reach every pid.
                self.store.append_event(
                    task.task_id,
                    {"event": "recovery_declined", "run_id": task.active_run_id, "alive": alive},
                )
                return self._save(
                    task,
                    error=(
                        f"run {task.active_run_id} is still executing as "
                        f"pid {', '.join(str(pid) for pid in alive)}; cancel it before recovering"
                    ),
                    error_category=ErrorCategory.AGENT_FAILURE,
                )
            self.store.append_event(
                task.task_id,
                {"event": "interrupted", "run_id": task.active_run_id, "after": task.last_successful_stage},
            )
            return self._needs_human(
                task,
                f"run interrupted after {task.last_successful_stage or 'INIT'}; workspace needs review",
                ErrorCategory.AGENT_FAILURE,
            )
        worktree = task.worktree_path
        if worktree is not None and not Path(worktree).is_dir():
            return self._needs_human(task, f"worktree missing: {worktree}", ErrorCategory.ISOLATION)
        if task.baseline_commit and worktree and not self._commit_exists(Path(worktree), task.baseline_commit):
            return self._needs_human(task, "baseline commit is unreachable", ErrorCategory.ISOLATION)
        if task.last_checkpoint and worktree and not Checkpoint(self.store.task_dir(task_id)).verify(
            task.last_checkpoint, Path(worktree)
        ):
            return self._needs_human(task, "workspace changed since the last checkpoint", ErrorCategory.ISOLATION)
        if task.state is TaskState.NEEDS_HUMAN:
            resume_state = self._RESUME_STATE.get(task.last_successful_stage, TaskState.INIT)
            self.store.append_event(
                task.task_id,
                {"event": "resumed", "from": TaskState.NEEDS_HUMAN.value, "to": resume_state.value},
            )
            return self._save(
                task,
                state=resume_state,
                error=None,
                error_category=None,
                active_run_id=None,
                active_run_pid=None,
                active_run_pids=[],
            )
        self.store.append_event(task.task_id, {"event": "resumed", "state": task.state.value})
        return task

    def recover_all(self) -> list[TaskRecord]:
        """Sweep every non-terminal task after a restart. Section 10.1 of the design."""
        root = self.store.root / "tasks"
        if not root.is_dir():
            return []
        recovered = []
        for entry in sorted(root.iterdir()):
            if not entry.is_dir():
                continue
            try:
                task = self.store.load(entry.name)
            except Exception:  # noqa: BLE001 - a corrupt task must not stop the sweep
                continue
            if not is_terminal(task.state):
                recovered.append(self.recover(task.task_id))
        return recovered

    def _stop_interrupted_run(self, task: TaskRecord) -> TaskRecord:
        """Terminate recorded provider processes before restart recovery continues."""
        alive = [pid for pid in self._recorded_pids(task) if is_running(pid)]
        for pid in alive:
            if not terminate_group(pid):
                raise InvalidTransition(f"provider pid {pid} is still running and could not be terminated")
            self.store.append_event(task.task_id, {"event": "run_terminated", "pid": pid, "signalled": True, "restart": True})
        deadline = time.monotonic() + 5.0
        remaining = [pid for pid in alive if is_running(pid)]
        while remaining and time.monotonic() < deadline:
            time.sleep(0.05)
            remaining = [pid for pid in remaining if is_running(pid)]
        if remaining:
            raise InvalidTransition(f"provider pid(s) {', '.join(str(pid) for pid in remaining)} are still running after termination")
        if task.active_run_id or task.active_run_pid or task.active_run_pids:
            task = self._save(task, active_run_id=None, active_run_pid=None, active_run_pids=[])
        return task

    def mark_waiting_after_restart(self, task_id: str, resume_state: TaskState) -> TaskRecord:
        task = self.store.load(task_id)
        task = self._stop_interrupted_run(task)
        task = self._save(task, resume_state=resume_state, state=TaskState.WAITING_CONTINUE)
        self.store.append_event(task_id, {"event": "waiting_continue", "resume_state": resume_state.value})
        return task

    def resume(self, task_id: str) -> TaskRecord:
        task = self.store.load(task_id)
        if task.state is not TaskState.WAITING_CONTINUE:
            raise InvalidTransition(f"task {task_id} is in {task.state.value}, not WAITING_CONTINUE")
        if task.resume_state is None:
            raise InvalidTransition(f"task {task_id} has no resume_state to restore")
        alive = [pid for pid in self._recorded_pids(task) if is_running(pid)]
        if alive:
            raise InvalidTransition(f"provider pid(s) {', '.join(str(pid) for pid in alive)} are still running")
        restored = task.resume_state
        task = self._save(task, state=restored, resume_state=None)
        self.store.append_event(task_id, {"event": "resumed_from_waiting", "state": restored.value})
        return task

    def startup_recovery_scan(self) -> list[TaskRecord]:
        """Mark every in-progress task waiting after a process restart, once only."""
        root = self.store.root / "tasks"
        if not root.is_dir():
            return []
        results = []
        for entry in sorted(root.iterdir()):
            if not entry.is_dir():
                continue
            try:
                task = self.store.load(entry.name)
            except Exception:  # noqa: BLE001
                continue
            if is_terminal(task.state) or task.state is TaskState.WAITING_CONTINUE:
                results.append(task)
                continue
            results.append(self.mark_waiting_after_restart(task.task_id, task.state))
        return results

    def run_to_completion(self, task_id: str, max_steps: int = 64) -> TaskRecord:
        task = self.store.load(task_id)
        for _ in range(max_steps):
            if is_terminal(task.state) or self.awaiting_approval(task) or task.state is TaskState.WAITING_CONTINUE:
                return task
            task = self.advance(task_id)
        return self._needs_human(task, "step budget exhausted", ErrorCategory.STALLED)

    # ----------------------------------------------------------------- advance

    def _lock_for(self, task_id: str) -> threading.Lock:
        with self._locks_guard:
            return self._locks[task_id]

    def advance(self, task_id: str) -> TaskRecord:
        # Two locks, two scopes: the in-process one serialises threads cheaply, the file
        # one covers a second Uvicorn worker or a CLI run against the same runtime.
        with self._lock_for(task_id), task_lock(self.store.task_dir(task_id), self.lock_timeout_s):
            return self._advance(task_id)

    def _advance(self, task_id: str) -> TaskRecord:
        task = self.store.load(task_id)
        if is_terminal(task.state):
            return task
        if task.state is TaskState.WAITING_CONTINUE:
            return task
        if task.state is TaskState.INIT:
            return self._transition(task, "planning_started")
        if task.state is TaskState.TEST:
            return self._run_tests(task)
        return self._agent_stage(task)

    # ------------------------------------------------------------- agent stages

    def approve(self, task_id: str) -> TaskRecord:
        """Record the operator's go-ahead for the stage the task is waiting on."""
        with self._lock_for(task_id), task_lock(self.store.task_dir(task_id), self.lock_timeout_s):
            task = self.store.load(task_id)
            if not self.awaiting_approval(task):
                return task
            task = self._save(task, approved_stages=[*task.approved_stages, task.state.value])
            self.store.append_event(task.task_id, {"event": "approved", "stage": task.state.value})
            return task

    @staticmethod
    def awaiting_approval(task: TaskRecord) -> bool:
        return (
            task.require_approval
            and task.state in WRITING_STAGES
            and task.state.value not in task.approved_stages
        )

    def _parallel_plan(self, task: TaskRecord) -> PlanArtifact | None:
        """The approved plan, when it has enough independent steps to be worth splitting."""
        if self.implement_workers < 2:
            return None
        path = self.store.task_dir(task.task_id) / "plan.json"
        if not path.is_file():
            return None
        try:
            plan = PlanArtifact.model_validate_json(path.read_text(encoding="utf-8"))
        except ValueError:
            return None
        reason = unsafe_to_parallelise(plan)
        if reason:
            self.store.append_event(
                task.task_id, {"event": "parallel_declined", "reason": reason}
            )
            return None
        return plan

    def _implement_in_parallel(self, task: TaskRecord, plan: PlanArtifact) -> TaskRecord:
        """Split the plan across worktrees, then merge the work back before review."""
        destination = Path(task.worktree_path or task.repo_path)
        runner = ParallelImplementer(
            self._adapter_for(task, "implementer"), self.workspace, self.implement_workers,
            patch_dir=self.store.task_dir(task.task_id) / "patches",
        )
        run_id = uuid4().hex
        task = self._save(task, active_run_id=run_id, active_run_pids=[])
        self.store.append_event(
            task.task_id, {"event": "parallel_started", "workers": self.implement_workers}
        )
        try:
            outcome = runner.run(
                plan, task.task_id, Path(task.repo_path), task.goal,
                self._handoff(task, TaskState.IMPLEMENT), destination=destination,
                on_start=lambda pid: self._record_pid(task.task_id, pid),
            )
        except MergeConflict as error:
            task = self._save(self.store.load(task.task_id), active_run_id=None, active_run_pids=[])
            return self._needs_human(task, f"parallel merge failed: {error}", ErrorCategory.ISOLATION)
        task = self.store.load(task.task_id)
        task = self._save(task, active_run_id=None, active_run_pid=None, active_run_pids=[])
        if is_terminal(task.state) or task.state is not TaskState.IMPLEMENT:
            self.store.append_event(
                task.task_id,
                {"event": "stage_abandoned", "stage": "IMPLEMENT", "state": task.state.value},
            )
            return task
        if not outcome.ok:
            failure = outcome.failures[0]
            return self._fail_stage(task, failure, TaskState.IMPLEMENT)
        artifact = outcome.merged()
        path = self.store.write_artifact(task.task_id, "impl_report.json", artifact)
        self.store.append_event(
            task.task_id, {"event": "parallel_merged", "subtasks": outcome.merged_files}
        )
        task = self._save(
            task,
            stage_hashes={
                **task.stage_hashes,
                "IMPLEMENT": hashlib.sha256(path.read_bytes()).hexdigest(),
            },
            last_successful_stage="IMPLEMENT",
        )
        task = self._checkpoint(task, TaskState.IMPLEMENT)
        return self._transition(task, "implementation_finished")

    def _adapter_for(self, task: TaskRecord, role: str):
        """The adapter for this role, honouring a model or effort the task chose."""
        default = self.planner if role == "planner" else self.implementer
        wanted = task.tuning.get(role)
        if not wanted or not (wanted.model or wanted.effort) or self.adapter_factory is None:
            return default
        try:
            return self.adapter_factory(role, wanted)
        except Exception as error:  # noqa: BLE001 - a bad choice must not lose the task
            self.store.append_event(
                task.task_id, {"event": "tuning_ignored", "role": role, "reason": str(error)}
            )
            return default

    def _agent_stage(self, task: TaskRecord) -> TaskRecord:
        stage = task.state
        if stage is TaskState.IMPLEMENT and not self.awaiting_approval(task):
            plan = self._parallel_plan(task)
            if plan is not None:
                return self._implement_in_parallel(task, plan)
        if self.awaiting_approval(task):
            return task
        adapter = self._adapter_for(task, "planner" if stage in PLANNER_STAGES else "implementer")
        run_id = uuid4().hex
        task = self._save(task, active_run_id=run_id)
        talk = self.conversation(task.task_id)
        request = AgentRequest(
            run_id,
            stage.value,
            Path(task.worktree_path or task.repo_path),
            task.goal,
            self._handoff(task, stage),
            task.session_ids.get(stage.value),
            task.baseline_commit,
            tuple(task.constraints),
            talk.guidance(),
        )
        result = adapter.run(request, on_start=lambda pid: self._record_pid(task.task_id, pid))
        task = self.store.load(task.task_id)
        task = self._save(task, active_run_id=None, active_run_pid=None)
        if is_terminal(task.state) or task.state is not stage:
            # Cancelled, or moved on, while the agent was running. Record the orphaned
            # result and leave the decision that was already made standing.
            self.store.append_event(
                task.task_id,
                {"event": "stage_abandoned", "stage": stage.value, "status": result.status,
                 "state": task.state.value},
            )
            return task
        if result.session_id:
            task = self._save(task, session_ids={**task.session_ids, stage.value: result.session_id})
        if result.usage:
            task = self._save(task, usage=_accumulate(task.usage, result.usage))
        expected = STAGE_SCHEMA.get(stage)
        if result.status != "SUCCESS" or result.artifact is None or (expected and not isinstance(result.artifact, expected)):
            return self._fail_stage(task, result, stage)
        problem = self._contract_violation(task, stage, result.artifact)
        if problem:
            return self._fail_stage(task, AgentResult("PROTOCOL", error=problem), stage)
        talk.mark_delivered()
        name = self._artifact_name(task, stage)
        path = self.store.write_artifact(task.task_id, name, result.artifact)
        talk.add(
            "planner" if stage in PLANNER_STAGES else "implementer",
            render(result.artifact, stage.value),
            stage=stage.value,
            artifact=name,
        )
        task = self._save(
            task,
            stage_hashes={**task.stage_hashes, stage.value: hashlib.sha256(path.read_bytes()).hexdigest()},
            attempts={name: count for name, count in task.attempts.items() if name != stage.value},
            approved_stages=[name for name in task.approved_stages if name != stage.value],
            last_successful_stage=stage.value,
        )
        task = self._checkpoint(task, stage)
        if stage is TaskState.PLAN:
            self.store.write_text(task.task_id, "plan.md", _render_plan(task, result.artifact))
            self.memory.learn_from_plan(result.artifact, task.task_id)
        if stage in {TaskState.REVIEW, TaskState.FINAL_VERIFY} and result.artifact.verdict != "PASS":
            self.memory.learn_from_review(result.artifact, task.task_id)
        return self._apply_stage_outcome(task, stage, result.artifact)

    def _apply_stage_outcome(self, task: TaskRecord, stage: TaskState, artifact) -> TaskRecord:
        if stage is TaskState.PLAN:
            return self._transition(task, "plan_drafted")
        if stage is TaskState.IMPLEMENT:
            return self._transition(task, "implementation_finished")
        if stage is TaskState.FIX:
            return self._transition(self._save(task, cycle=task.cycle + 1), "fix_finished")
        if stage is TaskState.PLAN_REVIEW:
            if artifact.verdict == "PASS":
                return self._transition(self._save(task, discussion_rounds=0), "plan_approved")
            if task.discussion_rounds < self.discussion_rounds:
                return self._discuss_plan(task, artifact)
            task = self._save(task, plan_rounds=task.plan_rounds + 1, discussion_rounds=0)
            if task.plan_rounds > MAX_OPERATOR_REPLANS:
                return self._needs_human(task, "plan rejected twice", ErrorCategory.STALLED)
            return self._transition(task, "plan_rejected")
        if stage in {TaskState.REVIEW, TaskState.FINAL_VERIFY}:
            advisory = artifact.advisory_issue_ids()
            if advisory:
                self.store.append_event(task.task_id, {"event": "advisory_issues", "stage": stage.value, "issues": advisory})
                task = self._save(task, advisory_issues=sorted(set(task.advisory_issues + advisory)))
        if stage is TaskState.FINAL_VERIFY:
            if artifact.verdict == "PASS" or not artifact.blocking_issue_ids():
                task = self._transition(task, "verified")
                self._learn_from_completed_plan(task)
                self._write_summary(task)
                return task
            return self._enter_fix(task, artifact, stage)
        if artifact.verdict == "PASS" or not artifact.blocking_issue_ids():
            return self._transition(self._save(task, last_issues={**task.last_issues, stage.value: []}), "review_passed")
        return self._enter_fix(task, artifact, stage)

    def _learn_from_completed_plan(self, task: TaskRecord) -> None:
        """Only a finished task promotes its assumptions from guesses to facts."""
        path = self.store.task_dir(task.task_id) / "plan.json"
        if not path.is_file():
            return
        try:
            plan = PlanArtifact.model_validate_json(path.read_text(encoding="utf-8"))
        except ValueError:
            return
        self.memory.learn_from_completion(plan, task.task_id)

    def _discuss_plan(self, task: TaskRecord, review: ReviewArtifact) -> TaskRecord:
        """One exchange between the planner and the reviewer over a rejected plan.

        The planner is handed the reviewer's specific objections and answers them with a
        revised plan; the reviewer then judges again. Still one validated artifact per
        turn, never free-form chat, so the state machine and every guard stay intact.
        An exchange that changes nothing is stalling and ends the discussion.
        """
        raised = review.issue_ids()
        if raised and raised == task.last_issues.get(TaskState.PLAN_REVIEW.value):
            self.store.append_event(
                task.task_id, {"event": "discussion_stalled", "issues": raised}
            )
            task = self._save(task, plan_rounds=task.plan_rounds + 1, discussion_rounds=0)
            if task.plan_rounds > MAX_OPERATOR_REPLANS:
                return self._needs_human(task, "plan rejected twice", ErrorCategory.STALLED)
            return self._transition(task, "plan_rejected")

        task = self._save(
            task,
            discussion_rounds=task.discussion_rounds + 1,
            last_issues={**task.last_issues, TaskState.PLAN_REVIEW.value: raised},
        )
        self.store.append_event(
            task.task_id,
            {"event": "discussion_round", "round": task.discussion_rounds, "issues": raised},
        )
        run_id = uuid4().hex
        talk = self.conversation(task.task_id)
        request = AgentRequest(
            run_id,
            TaskState.PLAN.value,
            Path(task.worktree_path or task.repo_path),
            task.goal,
            self._discussion_handoff(task, review),
            task.session_ids.get(TaskState.PLAN.value),
            task.baseline_commit,
            tuple(task.constraints),
            talk.guidance(),
        )
        task = self._save(task, active_run_id=run_id)
        result = self._adapter_for(task, "planner").run(
            request, on_start=lambda pid: self._record_pid(task.task_id, pid)
        )
        task = self._save(self.store.load(task.task_id), active_run_id=None, active_run_pid=None)
        if result.status != "SUCCESS" or not isinstance(result.artifact, PlanArtifact):
            return self._fail_stage(task, result, TaskState.PLAN)
        talk.mark_delivered()
        path = self.store.write_artifact(task.task_id, "plan.json", result.artifact)
        talk.add("planner", render(result.artifact, TaskState.PLAN.value), stage="PLAN", artifact="plan.json")
        self.store.write_text(task.task_id, "plan.md", _render_plan(task, result.artifact))
        # Stay in PLAN_REVIEW: the revised plan is re-judged by the next advance, which
        # keeps the exchange inside one state instead of inventing a transition for it.
        return self._save(
            task, stage_hashes={**task.stage_hashes, "PLAN": hashlib.sha256(path.read_bytes()).hexdigest()}
        )

    def _discussion_handoff(self, task: TaskRecord, review: ReviewArtifact) -> str:
        """The reviewer's objections, as data the planner must answer point by point."""
        issues = chr(10).join(
            f"- {issue.issue_id} [{issue.severity}] {issue.problem}"
            + (f" file: {issue.file}" if issue.file else "")
            + (f" location: {issue.location}" if issue.location else "")
            + (f" expected: {issue.expected_fix}" if issue.expected_fix else "")
            + (f" verification: {issue.verification}" if issue.verification else "")
            for issue in review.issues
        )
        directory = self.store.task_dir(task.task_id)
        plan = (directory / "plan.json").read_text(encoding="utf-8") if (directory / "plan.json").is_file() else ""
        return (
            "The reviewer rejected this plan. Revise it so every objection below is"
            + " answered, keeping whatever was already sound."
            + chr(10) * 2 + "OBJECTIONS:" + chr(10) + issues
            + (chr(10) * 2 + "CURRENT PLAN:" + chr(10) + plan if plan else "")
        )

    def _enter_fix(self, task: TaskRecord, review: ReviewArtifact, stage: TaskState) -> TaskRecord:
        """Gate every entry into FIX so a disagreeing reviewer cannot spin forever."""
        raised = review.blocking_issue_ids()
        if raised and raised == task.last_issues.get(stage.value):
            kind = "blockers" if review.blockers() else "issues"
            return self._needs_human(
                task, f"{stage.value} repeated identical {kind}: {', '.join(raised)}", ErrorCategory.STALLED
            )
        task = self._save(
            task, last_issues={**task.last_issues, stage.value: raised}, open_issues=raised
        )
        if task.cycle >= task.limits.max_fix_cycles:
            return self._needs_human(task, f"max_fix_cycles={task.limits.max_fix_cycles} reached", ErrorCategory.MAX_CYCLES)
        return self._transition(task, "changes_required")

    def _contract_violation(self, task: TaskRecord, stage: TaskState, artifact) -> str | None:
        """A fix must answer every raised issue id by name, or the loop cannot terminate."""
        if stage is not TaskState.FIX or not task.open_issues:
            return None
        answered = {item.issue_id for item in [*artifact.resolved, *artifact.not_resolved]}
        missing = sorted(set(task.open_issues) - answered)
        if missing:
            return f"fix report does not address issues: {', '.join(missing)}"
        return None

    def _fail_stage(self, task: TaskRecord, result: AgentResult, stage: TaskState) -> TaskRecord:
        category = _category(result.status)
        reason = result.error or f"{stage.value} failed"
        spent = task.attempts.get(stage.value, 0)
        if not self.retry_policy.should_retry(category, spent):
            return self._needs_human(task, reason, category)
        delay = self.retry_policy.wait(category, spent + 1)
        task = self._save(
            task,
            attempts={**task.attempts, stage.value: spent + 1},
            error=reason,
            error_category=category,
        )
        self.store.append_event(
            task.task_id,
            {
                "event": "stage_retry",
                "stage": stage.value,
                "category": category.value,
                "attempt": spent + 1,
                "delay_s": delay,
                "reason": reason,
            },
        )
        return task

    # ------------------------------------------------------------------- tests

    def _run_tests(self, task: TaskRecord) -> TaskRecord:
        outcome = self.test_runner(task, Path(task.worktree_path or task.repo_path))
        report = outcome if isinstance(outcome, StageTestReport) else StageTestReport(
            command="<injected>",
            exit_code=0 if outcome else 1,
            duration_ms=0,
            stdout_path="",
            stderr_path="",
            passed=bool(outcome),
        )
        regression = report.failed is not None and task.last_failed is not None and report.failed > task.last_failed
        report = report.model_copy(update={"regression": regression})
        self.store.write_artifact(task.task_id, f"test_report-{task.cycle + 1:02d}.json", report)
        task = self._save(task, last_failed=report.failed, regressions=task.regressions + (1 if regression else 0))
        if regression and task.regressions > 1:
            return self._needs_human(task, "tests regressed on consecutive rounds", ErrorCategory.STALLED)
        if report.passed:
            return self._transition(task, "tests_passed")
        if task.cycle >= task.limits.max_fix_cycles:
            return self._needs_human(task, f"max_fix_cycles={task.limits.max_fix_cycles} reached", ErrorCategory.MAX_CYCLES)
        return self._transition(task, "tests_failed")

    # ------------------------------------------------------------------- utils

    def _handoff(self, task: TaskRecord, stage: TaskState) -> str:
        names = {
            TaskState.PLAN_REVIEW: ["plan.json"],
            TaskState.IMPLEMENT: ["plan.json"],
            TaskState.REVIEW: ["plan.json", "impl_report.json", f"test_report-{task.cycle + 1:02d}.json"],
            TaskState.FIX: ["plan.json", f"review-{task.cycle + 1:02d}.json", f"test_report-{task.cycle + 1:02d}.json"],
            TaskState.FINAL_VERIFY: [
                "plan.json",
                "impl_report.json",
                f"test_report-{task.cycle + 1:02d}.json",
                f"review-{task.cycle + 1:02d}.json",
            ],
        }.get(stage, [])
        directory = self.store.task_dir(task.task_id)
        artifacts = "\n\n".join(
            (directory / name).read_text(encoding="utf-8") for name in names if (directory / name).is_file()
        )
        if stage not in {TaskState.REVIEW, TaskState.FINAL_VERIFY}:
            return artifacts
        rules = (
            "REVIEW ISSUE CLASSIFICATION:\n"
            "- CODE_BLOCKER: production code defect requiring a code change.\n"
            "- TEST_BLOCKER: failing or missing automated acceptance required by the plan.\n"
            "- MANUAL_ACCEPTANCE: physical-device or human-only verification.\n"
            "- DELIVERY_HYGIENE: repository or handoff cleanliness that does not break runtime behavior.\n"
            "- FOLLOW_UP: improvement outside the accepted task scope.\n"
            "Only CODE_BLOCKER and TEST_BLOCKER block completion. Do not reopen a resolved issue without new failing evidence.\n"
            "ORCHESTRATOR TEST EVIDENCE IS AUTHORITATIVE:\n"
            "- When the supplied test_report has passed=true and exit_code=0, accept that execution as complete evidence.\n"
            "- Do not create a TEST_BLOCKER merely because your own sandbox cannot rerun tests.\n"
            "- A TEST_BLOCKER after a passing report requires concrete new evidence of a failing or missing acceptance test."
        )
        return rules + ("\n\n" + artifacts if artifacts else "")

    @staticmethod
    def _artifact_name(task: TaskRecord, stage: TaskState) -> str:
        return {
            TaskState.PLAN: "plan.json",
            TaskState.PLAN_REVIEW: f"plan_review-{task.plan_rounds + 1:02d}.json",
            TaskState.IMPLEMENT: "impl_report.json",
            TaskState.REVIEW: f"review-{task.cycle + 1:02d}.json",
            TaskState.FIX: f"fix_report-{task.cycle + 1:02d}.json",
            TaskState.FINAL_VERIFY: "final_review.json",
        }[stage]

    def _write_summary(self, task: TaskRecord) -> None:
        lines = [
            f"# Task {task.task_id}",
            f"- goal: {task.goal}",
            f"- final state: {task.state.value}",
            f"- baseline: {task.baseline_commit}",
            f"- worktree: {task.worktree_path}",
            f"- fix cycles: {task.cycle}",
            f"- usage: {task.usage or 'not reported by the providers'}",
        ]
        if task.error:
            category = task.error_category.value if task.error_category else "?"
            lines.append(f"- error: [{category}] {task.error}")
        (self.store.task_dir(task.task_id) / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8")

    def _checkpoint(self, task: TaskRecord, stage: TaskState) -> TaskRecord:
        worktree = Path(task.worktree_path or task.repo_path)
        record = Checkpoint(self.store.task_dir(task.task_id)).capture(
            stage.value, task.cycle, worktree, task.baseline_commit
        )
        return self._save(task, last_checkpoint=record) if record else task

    def _save(self, task: TaskRecord, **updates) -> TaskRecord:
        updated = task.model_copy(update={**updates, "updated_at": datetime.now(timezone.utc)})
        self.store.save(updated)
        return updated

    def _transition(self, task: TaskRecord, event: str) -> TaskRecord:
        updated = self._save(task, state=next_state(task.state, event), error=None, error_category=None)
        self.store.append_event(task.task_id, {"event": event, "state": updated.state.value})
        return updated

    def _needs_human(self, task: TaskRecord, reason: str, category: ErrorCategory) -> TaskRecord:
        state = TaskState.NEEDS_HUMAN if is_terminal(task.state) or task.state is TaskState.WAITING_CONTINUE else next_state(task.state, "needs_human")
        checkpoint = task.last_checkpoint
        if (
            task.state not in {TaskState.DONE, TaskState.NEEDS_HUMAN}
            and category is not ErrorCategory.OPERATOR
            and task.active_run_id is None
        ):
            worktree = Path(task.worktree_path or task.repo_path)
            if worktree.is_dir():
                captured = Checkpoint(self.store.task_dir(task.task_id)).capture(
                    TaskState.NEEDS_HUMAN.value, task.cycle, worktree, task.baseline_commit
                )
                if captured is not None:
                    checkpoint = captured
        updated = self._save(
            task,
            state=state,
            error=reason,
            error_category=category,
            active_run_id=None,
            last_checkpoint=checkpoint,
        )
        self.store.append_event(
            task.task_id, {"event": "needs_human", "reason": reason, "category": category.value}
        )
        self.store.write_json(
            task.task_id,
            "failure.json",
            {
                "schema_version": 1,
                "task_id": updated.task_id,
                "stage": task.state.value,
                "category": category.value,
                "reason": reason,
                "cycle": updated.cycle,
                "last_successful_stage": updated.last_successful_stage,
                "worktree_path": str(updated.worktree_path) if updated.worktree_path else None,
                "at": updated.updated_at.isoformat(),
            },
        )
        self._write_summary(updated)
        return updated

    @staticmethod
    def _dirty_files(repo_path: Path) -> list[str]:
        result = subprocess.run(
            ["git", "-C", str(repo_path), "status", "--porcelain"], capture_output=True, text=True
        )
        return [line[3:] for line in result.stdout.splitlines() if line] if result.returncode == 0 else []

    @staticmethod
    def _commit_exists(worktree: Path, commit: str) -> bool:
        return (
            subprocess.run(
                ["git", "-C", str(worktree), "cat-file", "-e", f"{commit}^{{commit}}"], capture_output=True
            ).returncode
            == 0
        )


def _accumulate(total: dict[str, float], reported: dict) -> dict[str, float]:
    """Sum whatever the provider reported. Missing figures stay missing, never zero-filled."""
    merged = dict(total)
    for key, value in reported.items():
        if isinstance(value, (int, float)):
            merged[key] = round(merged.get(key, 0) + value, 6)
    return merged


def _render_plan(task: TaskRecord, plan: PlanArtifact) -> str:
    """Human-readable companion to plan.json. Never parsed, never used to advance state."""
    lines = [f"# 计划 · {plan.goal}", "", f"任务 `{task.task_id}`，基线 `{task.baseline_commit}`。", ""]
    for heading, items in (
        ("验收标准", plan.acceptance_criteria),
        ("假设", plan.assumptions),
        ("约束", plan.constraints),
        ("预期涉及文件", plan.files_expected),
    ):
        if items:
            lines += [f"## {heading}", ""] + [f"- {item}" for item in items] + [""]
    if plan.risks:
        lines += ["## 风险", ""]
        for item in plan.risks:
            mitigation = f" 缓解：{item.mitigation}" if item.mitigation else ""
            lines.append(f"- {item.risk}{mitigation}")
        lines.append("")
    if plan.steps:
        lines += ["## 步骤", ""]
        for step in plan.steps:
            verification = f" 验证：{step.verification}" if step.verification else ""
            lines.append(f"{step.id}. {step.action}{verification}")
        lines.append("")
    return chr(10).join(lines)