"""Implementation stage run orchestration.

This module owns the Stage Run Claim for one ``implementation`` stage run:
recover consumer state, select an available Stage Map entry, provision the
isolated stage worktree, and record the started event.
"""
from __future__ import annotations

import datetime as _dt
import subprocess
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from . import stage_targets
from .design_prep import DesignPrepDecision, DesignPrepError, resolve_design_prep
from .final_report_paths import final_report_data_path
from .stage_reconcile import auto_reconcile_best_effort
from .worker_prompt_policy import IMPLEMENTATION_STAGE_HEADER


class ImplementationStageError(Exception):
    """Implementation stage claim or provisioning failed."""


@dataclass
class StageRunClaim:
    """Claim over one Stage Map stage and its isolated worktree coordinates."""

    stage: int
    worktree_path: str
    worktree_branch: str
    worktree_base_ref: str
    worktree_status: str
    worktree_note: str
    started_head_commit: str
    design_prep_decision: DesignPrepDecision
    concurrent_stages: list[int] = field(default_factory=list)


def _git_out(cwd: str | Path, *args: str) -> str:
    result = subprocess.run(
        ["git", "-C", str(cwd), *args],
        capture_output=True,
        text=True,
    )
    return result.stdout.strip() if result.returncode == 0 else ""


def _as_implementation_stage_error(exc: Exception) -> ImplementationStageError:
    if isinstance(exc, ImplementationStageError):
        return exc
    return ImplementationStageError(str(exc))


def _resolve_stage_design_prep(plan_path: Path, stage: int) -> DesignPrepDecision:
    plan_data_path = final_report_data_path(plan_path)
    if plan_data_path.is_file():
        try:
            return resolve_design_prep(plan_data_path, stage)
        except DesignPrepError as exc:
            # prepare catches ImplementationStageError → PrepareError; an
            # untranslated DesignPrepError leaves okstra_ctl.run as a traceback.
            raise _as_implementation_stage_error(exc) from exc
    return DesignPrepDecision(
        outcome="proceed",
        effective_items=(),
        assumptions_to_inject=(),
        warnings=("legacy-unassessed",),
        request_paths=(),
        reason="legacy plan has no design preparation assessment",
    )


def claim_implementation_stage_run(
    inp: Any,
    ctx_stage_map: list[dict[str, Any]],
    task_group_segment: str,
    task_id_segment: str,
    task_key: str,
    executor_worktree_status: str,
) -> StageRunClaim:
    """Claim a ready implementation stage for one run.

    The returned claim is path-independent with respect to run artifacts:
    callers can compute stage-specific run paths after this function resolves
    the selected stage.
    """
    from .consumers import backfill_done_from_carry
    from . import worktree as _worktree
    from . import worktree_registry as _reg

    plan_run_root = Path(inp.approved_plan_path).resolve().parents[1]
    backfill_done_from_carry(plan_run_root)
    auto_reconcile_best_effort(inp, plan_run_root)
    reserved_stages = _reg.list_active_stage_numbers(
        inp.project_id,
        inp.task_group,
        inp.task_id,
    )
    snapshot = stage_targets.read_stage_lifecycle_snapshot(
        ctx_stage_map,
        plan_run_root,
        reserved_stages=reserved_stages,
    )

    try:
        selected = snapshot.resolve_implementation_stage(inp.stage)
    except stage_targets.StageTargetError as exc:
        raise _as_implementation_stage_error(exc) from exc

    design_prep_decision = _resolve_stage_design_prep(
        Path(inp.approved_plan_path), selected
    )
    if design_prep_decision.outcome == "wait_for_input":
        requests = ", ".join(design_prep_decision.request_paths) or "no request path"
        raise ImplementationStageError(
            f"stage {selected} waits for design input: "
            f"{design_prep_decision.reason}; {requests}"
        )
    if design_prep_decision.outcome == "replan":
        raise ImplementationStageError(
            f"stage {selected} requires implementation-planning rerun: "
            f"{design_prep_decision.reason}"
        )

    concurrent_stages = snapshot.concurrent_stage_numbers(
        selected_stage=selected,
    )

    if executor_worktree_status.startswith("skipped"):
        head = _git_out(inp.project_root, "rev-parse", "HEAD")
        claim = StageRunClaim(
            stage=selected,
            worktree_path="",
            worktree_branch="",
            worktree_base_ref="",
            worktree_status=executor_worktree_status,
            worktree_note="",
            started_head_commit=head,
            design_prep_decision=design_prep_decision,
            concurrent_stages=concurrent_stages,
        )
        _record_stage_run_claim_started(task_key, plan_run_root, claim)
        return claim

    # The anchor and multi-dep candidate base must come from the task-key
    # worktree HEAD (where stage work accumulates), not from inp.project_root
    # (the user's invocation cwd, which may sit on an older/unrelated commit).
    # The task-key worktree is the SSOT in the registry; whole-task
    # final-verification (run.py) resolves its base from the same path.
    task_entry = _reg.lookup(inp.project_id, inp.task_group, inp.task_id)
    task_worktree_path = (task_entry.worktree_path if task_entry else "") or inp.project_root
    head_sha = _git_out(task_worktree_path, "rev-parse", "HEAD")
    if head_sha:
        _reg.set_implementation_base(
            inp.project_id,
            inp.task_group,
            inp.task_id,
            head_sha,
        )
    anchor = _reg.get_implementation_base(
        inp.project_id,
        inp.task_group,
        inp.task_id,
    ) or ""

    selected_stage = next(
        s for s in ctx_stage_map if s["stage_number"] == selected
    )
    try:
        stage_base = stage_targets.resolve_stage_base_commit(
            selected_stage,
            snapshot.done_rows,
            anchor_base_commit=anchor,
            candidate_base=head_sha,
            project_root=Path(task_worktree_path),
            plan_run_root=plan_run_root,
        )
    except stage_targets.StageTargetError as exc:
        raise _as_implementation_stage_error(exc) from exc

    try:
        prov = _worktree.provision_stage_worktree(
            project_root=Path(inp.project_root),
            project_id=inp.project_id,
            task_group_segment=task_group_segment,
            task_id_segment=task_id_segment,
            work_category=inp.work_category,
            stage_number=selected,
            base_commit=stage_base,
        )
    except RuntimeError as exc:
        from .git_reconcile import guidance

        hint = guidance(
            plan_run_root=plan_run_root,
            project_id=inp.project_id,
            task_group=inp.task_group,
            task_id=inp.task_id,
            work_category=inp.work_category,
        )
        raise ImplementationStageError(
            f"stage worktree provisioning failed: {exc}\n{hint}"
        ) from exc

    claim = StageRunClaim(
        stage=selected,
        worktree_path=prov.path,
        worktree_branch=prov.branch,
        worktree_base_ref=prov.base_ref,
        worktree_status=prov.status,
        worktree_note=prov.note,
        started_head_commit=prov.base_ref,
        design_prep_decision=design_prep_decision,
        concurrent_stages=concurrent_stages,
    )
    _record_stage_run_claim_started(task_key, plan_run_root, claim)
    return claim


def _record_stage_run_claim_started(
    task_key: str,
    plan_run_root: Path,
    claim: StageRunClaim,
) -> None:
    from .consumers import append_consumer

    now = _dt.datetime.now(_dt.timezone.utc).isoformat()
    append_consumer(
        plan_run_root,
        impl_task_key=task_key,
        stage=claim.stage,
        status="started",
        started_at=now,
        head_commit=claim.started_head_commit,
    )


def publish_stage_run_claim(
    inp: Any,
    ctx: dict[str, Any],
    ctx_stage_map: list[dict[str, Any]],
    claim: StageRunClaim,
) -> None:
    """Publish a Stage Run Claim into run context."""
    ctx["parsed_stage_map"] = ctx_stage_map
    ctx["effective_stages"] = [claim.stage]
    csv = str(claim.stage)
    ctx["EFFECTIVE_STAGES"] = csv
    ctx["CONCURRENT_RUN_STAGES"] = ",".join(str(s) for s in claim.concurrent_stages)
    ctx["STAGE_BATCH_DIRECTIVE"] = (
        f"- {IMPLEMENTATION_STAGE_HEADER} `{csv}`. "
        "Execute exactly this Stage Map stage — this is the authoritative scope. "
        "Do NOT recompute from `consumers.jsonl`; the runtime already selected "
        "and reserved this stage."
    )
    ctx["DESIGN_PREP_CONTEXT"] = claim.design_prep_decision.as_prompt_markdown()
    inp.stage = csv
    print(f"selected stages: {csv}", file=sys.stdout)

    if claim.worktree_status and not claim.worktree_status.startswith("skipped"):
        ctx["EXECUTOR_WORKTREE_PATH"] = claim.worktree_path
        ctx["EXECUTOR_WORKTREE_BRANCH"] = claim.worktree_branch
        ctx["EXECUTOR_WORKTREE_BASE_REF"] = claim.worktree_base_ref
        ctx["EXECUTOR_WORKTREE_STATUS"] = claim.worktree_status
        ctx["EXECUTOR_WORKTREE_NOTE"] = claim.worktree_note