"""Split one approved plan across parallel implementers. Section 16 of the design.

Each subtask gets its own worktree, they run concurrently, and the reviewer sees the
combined diff. The state machine is untouched: this sits beside it, driving the same
adapter contract, so a single-implementer task behaves exactly as before.
"""

from __future__ import annotations

import subprocess
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from uuid import uuid4

from .adapters.base import AgentAdapter, AgentRequest, AgentResult
from .domain import ImplementationArtifact, PlanArtifact, PlanStep


@dataclass(frozen=True)
class Subtask:
    subtask_id: str
    steps: tuple[PlanStep, ...]
    worktree: Path
    #: Commit the worktree started from. Everything after it is this subtask's work,
    #: whether the orchestrator committed it or the provider committed it itself.
    base: str | None = None

    def goal(self, parent_goal: str) -> str:
        listed = "; ".join(f"{step.id}: {step.action}" for step in self.steps)
        return f"{parent_goal} — only these steps: {listed}"


#: The orchestrator commits as itself, so a host with no global Git identity still works.
IDENTITY = ("-c", "user.email=orchestrator@local", "-c", "user.name=dual-agent")


class MergeConflict(RuntimeError):
    pass


@dataclass
class ParallelResult:
    artifacts: list[ImplementationArtifact] = field(default_factory=list)
    failures: list[AgentResult] = field(default_factory=list)
    merged_files: list[str] = field(default_factory=list)

    @property
    def ok(self) -> bool:
        return not self.failures

    def merged(self) -> ImplementationArtifact:
        """One report for the reviewer, with each file and test named once."""
        files, tests, notes = [], [], []
        for artifact in self.artifacts:
            files += artifact.files_changed
            tests += artifact.tests_run
            if artifact.notes:
                notes.append(artifact.notes)
        return ImplementationArtifact(
            files_changed=sorted(dict.fromkeys(files)),
            tests_run=sorted(dict.fromkeys(tests)),
            notes=" | ".join(notes),
        )


def partition(plan: PlanArtifact, workers: int) -> list[tuple[PlanStep, ...]]:
    """Group plan steps into at most `workers` buckets, preserving order.

    Callers gate on `unsafe_to_parallelise` first, so by here the steps are known to be
    independent and round-robin is enough.
    """
    if workers < 2 or len(plan.steps) < 2:
        return [tuple(plan.steps)] if plan.steps else []
    buckets: list[list[PlanStep]] = [[] for _ in range(min(workers, len(plan.steps)))]
    for index, step in enumerate(plan.steps):
        buckets[index % len(buckets)].append(step)
    return [tuple(bucket) for bucket in buckets if bucket]


def unsafe_to_parallelise(plan: PlanArtifact) -> str | None:
    """Why this plan must run in one worktree, or None when splitting it is safe.

    Spending several real provider calls only to hit a cherry-pick conflict is the worst
    outcome, so the checks are deliberately conservative: anything undeclared counts
    against parallelising.
    """
    if len(plan.steps) < 2:
        return "the plan has fewer than two steps"
    if any(step.depends_on for step in plan.steps):
        return "some steps declare dependencies on others"
    declared = [set(step.files) for step in plan.steps]
    if not all(declared):
        if len(set(plan.files_expected)) < 2:
            return "the plan expects fewer than two distinct files"
        return "not every step declares the files it touches"
    for index, files in enumerate(declared):
        for other in declared[index + 1:]:
            overlap = files & other
            if overlap:
                return f"steps share files: {', '.join(sorted(overlap))}"
    return None


def touches_same_file(plan: PlanArtifact) -> bool:
    """Kept for callers that only want the file-overlap question."""
    return unsafe_to_parallelise(plan) is not None


class ParallelImplementer:
    def __init__(self, adapter: AgentAdapter, workspace, workers: int = 2,
                 patch_dir: Path | None = None) -> None:
        self.adapter = adapter
        self.workspace = workspace
        self.workers = workers
        #: Where each subtask's diff is kept before its worktree is removed.
        self.patch_dir = Path(patch_dir) if patch_dir else None

    def cleanup(self, subtasks: list[Subtask], repo_path: Path) -> None:
        """Save each subtask's diff, then remove the worktree.

        Without this every parallel task leaves behind one worktree per worker, plus its
        registration in the repository. The diff is kept first so a failed run is still
        diagnosable after the directory is gone.
        """
        for subtask in subtasks:
            if self.patch_dir is not None:
                diff = self._full_diff(subtask)
                if diff.strip():
                    self.patch_dir.mkdir(parents=True, exist_ok=True)
                    (self.patch_dir / f"subtask-{subtask.subtask_id}.diff").write_text(
                        diff, encoding="utf-8", newline=chr(10)
                    )
            remove = getattr(self.workspace, "remove", None)
            if remove is not None:
                remove(repo_path, subtask.worktree)

    @staticmethod
    def _full_diff(subtask: Subtask) -> str:
        """Everything this subtask produced, against the commit its worktree started from.

        Anchoring on the base rather than HEAD~1 is what makes the saved patch complete:
        a provider that committed three times would otherwise have the first two thrown
        away with the worktree, and a failed run is exactly when that evidence matters.
        """
        # Intent-to-add so brand new files appear in the diff. A plain `git diff` omits
        # untracked content entirely, and a file the agent just created is precisely what
        # a failed run needs to preserve. The worktree is about to be removed, so touching
        # its index costs nothing.
        _git(subtask.worktree, "add", "-A", "-N")
        anchor = (subtask.base or "").strip()
        if anchor:
            against_base = _git(subtask.worktree, "diff", "--binary", anchor)
            if against_base is not None:
                return against_base
        # No usable base: keep at least the uncommitted work rather than nothing.
        return _git(subtask.worktree, "diff", "--binary", "HEAD") or ""

    def plan_subtasks(self, plan: PlanArtifact, task_id: str, repo_path: Path) -> list[Subtask]:
        groups = partition(plan, self.workers)
        subtasks = []
        for index, steps in enumerate(groups):
            subtask_id = f"{task_id}-p{index}"
            worktree, base = self.workspace.create(subtask_id, repo_path)
            subtasks.append(Subtask(subtask_id, steps, Path(worktree), base))
        return subtasks

    def collect(self, subtasks: list[Subtask], destination: Path) -> list[str]:
        """Commit each subtask worktree and cherry-pick it into the task worktree.

        Section 7.3 names cherry-pick as the import path, and it is also the only one that
        is safe for binary files and for checkouts whose line endings differ from the
        diff text. Conflicts abort rather than resolving themselves: two implementers that
        edited the same lines need a human or a re-plan, not a guess.
        """
        before = _git(destination, "rev-parse", "HEAD")
        if before is None:
            raise MergeConflict(f"{destination} is not a usable worktree")
        before = before.strip()

        # The rollback below is a hard reset, which would also destroy anything already
        # uncommitted here. Refusing up front is the only way to keep that promise: the
        # caller commits or stashes first, and nothing of theirs can be lost by a conflict.
        dirty = _git(destination, "status", "--porcelain")
        if dirty is None:
            raise MergeConflict(f"cannot read the state of {destination}")
        if dirty.strip():
            raise MergeConflict(
                f"{destination} has uncommitted changes; commit or stash them before merging "
                f"parallel work, because a failed merge is rolled back with a hard reset"
            )

        applied: list[str] = []
        try:
            for subtask in subtasks:
                commits = self._pending_commits(subtask)
                if not commits:
                    continue
                for commit in commits:
                    result = subprocess.run(
                        ["git", "-C", str(destination), *IDENTITY, "cherry-pick", "--allow-empty", commit],
                        capture_output=True, text=True, errors="replace",
                    )
                    if result.returncode:
                        subprocess.run(
                            ["git", "-C", str(destination), "cherry-pick", "--abort"],
                            capture_output=True, check=False,
                        )
                        raise MergeConflict(
                            f"{subtask.subtask_id} does not merge cleanly: "
                            f"{(result.stderr or result.stdout).strip()[:200]}"
                        )
                applied.append(subtask.subtask_id)
        except MergeConflict:
            # All or nothing. A half-merged worktree is worse than none: the reviewer would
            # be handed one subtask's work as if it were the whole change.
            subprocess.run(
                ["git", "-C", str(destination), "reset", "--hard", before],
                capture_output=True, check=False,
            )
            raise
        return applied

    @staticmethod
    def _pending_commits(subtask: Subtask) -> list[str]:
        """Every commit this subtask produced, oldest first.

        Providers are allowed to commit their own work, and several often do. Looking only
        at the dirty worktree would silently drop everything an agent had already
        committed, so the range from the worktree's base to its HEAD is the real answer.
        """
        status = _git(subtask.worktree, "status", "--porcelain")
        if status is None:
            raise MergeConflict(f"{subtask.subtask_id} is not a usable worktree")
        if status.strip():
            if _git(subtask.worktree, "add", "-A") is None:
                raise MergeConflict(f"cannot stage {subtask.subtask_id}")
            if _git(subtask.worktree, *IDENTITY, "commit", "-qm", f"subtask {subtask.subtask_id}") is None:
                raise MergeConflict(f"cannot commit {subtask.subtask_id}")

        head = _git(subtask.worktree, "rev-parse", "HEAD")
        if head is None:
            raise MergeConflict(f"cannot read HEAD of {subtask.subtask_id}")
        head = head.strip()
        if not subtask.base:
            # No recorded base: fall back to the single commit we just made, if any.
            return [head] if status.strip() else []
        if head == subtask.base.strip():
            return []
        listed = _git(subtask.worktree, "rev-list", "--reverse", f"{subtask.base.strip()}..{head}")
        if listed is None:
            raise MergeConflict(f"cannot list the commits of {subtask.subtask_id}")
        return [line.strip() for line in listed.splitlines() if line.strip()]

    def run(self, plan: PlanArtifact, task_id: str, repo_path: Path, goal: str, handoff: str,
            destination: Path | None = None, on_start=None) -> ParallelResult:
        subtasks = self.plan_subtasks(plan, task_id, repo_path)
        outcome = ParallelResult()
        if not subtasks:
            return outcome

        def execute(subtask: Subtask) -> AgentResult:
            return self.adapter.run(
                AgentRequest(uuid4().hex, "IMPLEMENT", subtask.worktree, subtask.goal(goal), handoff),
                on_start=on_start,
            )

        try:
            with ThreadPoolExecutor(max_workers=len(subtasks)) as pool:
                results = list(pool.map(execute, subtasks))
            for result in results:
                if result.status == "SUCCESS" and isinstance(result.artifact, ImplementationArtifact):
                    outcome.artifacts.append(result.artifact)
                else:
                    outcome.failures.append(result)
            if outcome.ok and destination is not None:
                outcome.merged_files = self.collect(subtasks, Path(destination))
            return outcome
        finally:
            self.cleanup(subtasks, repo_path)


def _git(worktree: Path, *args: str) -> str | None:
    result = subprocess.run(
        ["git", "-C", str(worktree), *args], capture_output=True, text=True, errors="replace"
    )
    return result.stdout if result.returncode == 0 else None