"""Agent loop CLI that spawns Claude to work through peon-mcp tasks.

Talks to the peon REST API (``PEON_API_URL``, default ``http://127.0.0.1:8420``)
to check for tasks, manage heartbeats, and coordinate agent runs.  No direct
SQLite access — works in Kubernetes where peon-loop and peon-api run as
separate pods.

Supports two modes:
  - Single-project: peon-loop <project_id> <repo_path> [options]
  - Project-agnostic: peon-loop [options]
    Picks up tasks from any project in the database, auto-cloning repos
    via github_url when needed.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import re
import shutil
import signal
import socket
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from datetime import datetime, timezone
from pathlib import Path
from shutil import which

log = logging.getLogger("peon-loop")


# ---------------------------------------------------------------------------
# Session tracking
# ---------------------------------------------------------------------------


def _get_api_url() -> str:
    return os.environ.get("PEON_API_URL", "http://127.0.0.1:8420").rstrip("/")


def _get_api_key() -> str:
    return os.environ.get("PEON_API_KEY", "")


def _http_request(
    method: str,
    url: str,
    payload: dict | None = None,
    *,
    timeout: int = 5,
    quiet_statuses: set[int] | None = None,
) -> dict | None:
    """Make a synchronous HTTP request, returning parsed JSON or None on failure.

    Args:
        method: HTTP method (GET, POST, PUT, DELETE).
        url: Fully qualified URL.
        payload: Optional JSON body.
        timeout: Request timeout in seconds (default 5).
        quiet_statuses: HTTP status codes that should not produce a warning log
                        (e.g. {409} for expected conflicts, {404} for optional lookups).
    """
    body = json.dumps(payload).encode("utf-8") if payload is not None else None
    headers: dict[str, str] = {"Content-Type": "application/json"}
    api_key = _get_api_key()
    if api_key:
        headers["X-API-Key"] = api_key
    req = urllib.request.Request(url, data=body, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read())
            _api_failure_counts.pop(url, None)  # Reset on success
            return data
    except urllib.error.HTTPError as exc:
        if quiet_statuses and exc.code in quiet_statuses:
            return None
        _api_failure_counts[url] = _api_failure_counts.get(url, 0) + 1
        if _api_failure_counts[url] <= _API_FAILURE_SUPPRESS_AFTER:
            log.warning("API call failed (%s %s): HTTP %d", method, url, exc.code)
        elif _api_failure_counts[url] == _API_FAILURE_SUPPRESS_AFTER + 1:
            log.warning("API call failed (%s %s): HTTP %d (suppressing further warnings for this URL)", method, url, exc.code)
        return None
    except (urllib.error.URLError, json.JSONDecodeError, OSError) as exc:
        _api_failure_counts[url] = _api_failure_counts.get(url, 0) + 1
        if _api_failure_counts[url] <= _API_FAILURE_SUPPRESS_AFTER:
            log.warning("API call failed (%s %s): %s", method, url, exc)
        elif _api_failure_counts[url] == _API_FAILURE_SUPPRESS_AFTER + 1:
            log.warning("API call failed (%s %s): %s (suppressing further warnings for this URL)", method, url, exc)
        return None


# Track consecutive API failures per URL to suppress repeated warnings
_api_failure_counts: dict[str, int] = {}
_API_FAILURE_SUPPRESS_AFTER = 5


def _should_self_update() -> bool:
    """Check if a newer version is deployed and this pod should restart.

    Compares the loop's own version against the API's reported version.
    If the API is running a newer version, the loop should exit so K8s
    restarts it with the new image.
    """
    try:
        import importlib.metadata
        my_version = importlib.metadata.version("peon-mcp")
    except Exception:
        return False  # Can't determine version, don't update

    api_url = _get_api_url()
    info = _http_request("GET", f"{api_url}/api/system-info", timeout=3)
    if not info:
        return False  # API unreachable, don't update

    api_version = info.get("version", "")
    if not api_version or api_version == my_version or api_version == "dev":
        return False  # Same version or unknown

    log.info(
        "Version mismatch: loop=%s, api=%s — will self-update after current task",
        my_version,
        api_version,
    )
    return True


class SessionTracker:
    """Tracks the lifecycle of a single Claude agent run via the sessions REST API.

    All methods are fire-and-forget: failures log a warning and the loop
    continues normally. Session tracking is observability, not a critical path.
    """

    def __init__(self, api_url: str):
        self._api_url = api_url
        self.session_id: int | None = None

    def start(
        self,
        project_id: str,
        loop_id: str,
        model: str,
        hostname: str = "",
        task_id: int | None = None,
    ) -> None:
        """Create the session record before the agent process starts.

        When *task_id* is provided (pre-claimed task), the session is created
        with the task already linked — no follow-up ``link_task()`` needed.
        """
        url = f"{self._api_url}/api/sessions"
        payload: dict = {
            "project_id": project_id,
            "loop_id": loop_id,
            "model": model,
            "hostname": hostname,
        }
        if task_id is not None:
            payload["task_id"] = task_id
        result = _http_request("POST", url, payload)
        if result and isinstance(result.get("id"), int):
            self.session_id = result["id"]
            log.info("Session tracking started (session_id=%d)", self.session_id)

    def event(self, event_type: str, tool_name: str = "", detail: str = "") -> None:
        """Append an event to the session."""
        if self.session_id is None:
            log.warning("SessionTracker.event() called but session_id is None")
            return
        url = f"{self._api_url}/api/sessions/{self.session_id}/events"
        _http_request(
            "POST",
            url,
            {"session_id": self.session_id, "event_type": event_type,
             "tool_name": tool_name, "detail": detail},
        )

    def finish(
        self,
        status: str,
        exit_code: int | None = None,
        tokens_input: int | None = None,
        tokens_output: int | None = None,
        cost_usd: float | None = None,
        error_summary: str = "",
        task_id: int | None = None,
    ) -> None:
        """Update the session when the agent process exits."""
        if self.session_id is None:
            log.warning("SessionTracker.finish() called but session_id is None")
            return
        payload: dict = {
            "status": status,
            "ended_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        }
        if task_id is not None:
            payload["task_id"] = task_id
        if exit_code is not None:
            payload["exit_code"] = exit_code
        if tokens_input is not None:
            payload["tokens_input"] = tokens_input
        if tokens_output is not None:
            payload["tokens_output"] = tokens_output
        if cost_usd is not None:
            payload["cost_usd"] = cost_usd
        if error_summary:
            payload["error_summary"] = error_summary
        url = f"{self._api_url}/api/sessions/{self.session_id}"
        _http_request("PUT", url, payload)
        log.info("Session tracking finished (session_id=%d, status=%s, task_id=%s)", self.session_id, status, task_id)

    def link_task(self, task_id: int) -> None:
        """Associate this session with a task as soon as the task is discovered.

        Called early in the agent run — typically within seconds of the agent
        claiming a task — rather than waiting until session teardown.
        This is a partial update (status remains unchanged).
        """
        if self.session_id is None:
            return
        url = f"{self._api_url}/api/sessions/{self.session_id}"
        _http_request("PUT", url, {"task_id": task_id})
        log.info("Session %d linked to task #%d (early link)", self.session_id, task_id)

    def cancel(self) -> None:
        """Mark this session as cancelled (called on loop shutdown)."""
        self.finish("cancelled")


def _parse_agent_log(log_file: Path) -> dict:
    """Best-effort extraction of token/cost stats from a Claude agent log.

    First scans for stream-json ``"type":"result"`` events (used when the agent
    runs with ``--output-format stream-json``).  Falls back to regex patterns
    for human-readable text output.

    Returns a dict with keys: tokens_input, tokens_output, cost_usd
    (all optional, omitted if not found).
    """
    stats: dict = {}
    try:
        lines = log_file.read_text(errors="replace").splitlines()

        # First pass: look for stream-json result event.
        # The Claude CLI emits a final {"type": "result", "total_cost_usd": ...,
        # "usage": {"input_tokens": ..., "output_tokens": ...}} line.
        for line in lines:
            stripped = line.strip()
            if not stripped:
                continue
            try:
                event = json.loads(stripped)
            except json.JSONDecodeError:
                continue
            if event.get("type") == "result":
                if "total_cost_usd" in event:
                    stats["cost_usd"] = float(event["total_cost_usd"])
                usage = event.get("usage") or {}
                if "input_tokens" in usage:
                    stats["tokens_input"] = int(usage["input_tokens"])
                if "output_tokens" in usage:
                    stats["tokens_output"] = int(usage["output_tokens"])
                break  # result event found; stop scanning

        if stats:
            return stats  # stream-json data found; skip regex fallback

        # Fallback: regex scan on last 100 lines for human-readable output.
        text = "\n".join(lines[-100:])

        # Match: "Total cost: $0.1234"
        m = re.search(r"[Tt]otal cost[:\s]+\$?([\d.]+)", text)
        if m:
            stats["cost_usd"] = float(m.group(1))

        # Match: "3,042 input tokens" or "3042 input tokens"
        m = re.search(r"([\d,]+)\s+input tokens?", text)
        if m:
            stats["tokens_input"] = int(m.group(1).replace(",", ""))

        # Match: "1,234 output tokens" or "1234 output tokens"
        m = re.search(r"([\d,]+)\s+output tokens?", text)
        if m:
            stats["tokens_output"] = int(m.group(1).replace(",", ""))

        # Alternate format: "cost of $0.34 (3,042 input tokens, 1,234 output tokens)"
        m = re.search(
            r"cost of \$?([\d.]+)\s+\(([\d,]+)\s+input tokens?,\s*([\d,]+)\s+output tokens?\)",
            text,
        )
        if m:
            stats["cost_usd"] = float(m.group(1))
            stats["tokens_input"] = int(m.group(2).replace(",", ""))
            stats["tokens_output"] = int(m.group(3).replace(",", ""))

    except Exception as exc:
        log.debug("Could not parse agent log for stats: %s", exc)

    return stats


def _create_fallback_work_log(api_url: str, project_id: str, task_id: int) -> None:
    """Create a minimal work log entry if the agent forgot to call log_work.

    Checks whether any work logs already exist for the task.  If none are
    found, fetches task metadata and creates a brief summary log so the
    audit trail is never empty for completed tasks.
    """
    # Skip if work logs already exist
    logs = _http_request("GET", f"{api_url}/api/tasks/{task_id}/logs")
    if logs and isinstance(logs, list) and len(logs) > 0:
        return

    # Fetch task metadata to build summary
    task_data = _http_request("GET", f"{api_url}/api/tasks/{task_id}")
    if not task_data:
        log.warning("Could not fetch task #%d metadata for fallback work log", task_id)
        return

    title = task_data.get("title", f"Task #{task_id}")
    pr_url = task_data.get("pr_url", "")
    lines_added = task_data.get("lines_added") or 0
    lines_deleted = task_data.get("lines_deleted") or 0

    summary_parts = [f"Completed: {title}"]
    if lines_added or lines_deleted:
        summary_parts.append(f"+{lines_added}/-{lines_deleted} lines")
    if pr_url:
        summary_parts.append(f"PR: {pr_url}")
    summary = " | ".join(summary_parts)

    result = _http_request(
        "POST",
        f"{api_url}/api/projects/{project_id}/logs",
        {"task_id": task_id, "summary": summary, "files_changed": "", "test_result": ""},
    )
    if result:
        log.info("Created fallback work log for task #%d", task_id)
    else:
        log.warning("Failed to create fallback work log for task #%d", task_id)


# ---------------------------------------------------------------------------
# DB pre-flight
# ---------------------------------------------------------------------------


def ensure_project(project_id: str, repo_path: str, github_url: str = "") -> None:
    """Ensure the project exists in the API, updating path and github_url.

    Creates the project if it doesn't exist (handles 409 Conflict on duplicate).
    Always updates the path for the current machine.
    """
    api_url = _get_api_url()
    # Try to create — 409 is expected if it already exists
    _http_request(
        "POST",
        f"{api_url}/api/projects",
        {"id": project_id, "path": repo_path, "github_url": github_url or ""},
        quiet_statuses={409},
    )
    # Always update path and optionally backfill github_url
    update_payload: dict = {"path": repo_path}
    if github_url:
        update_payload["github_url"] = github_url
    _http_request("PUT", f"{api_url}/api/projects/{project_id}", update_payload)
    log.info("Ready to work! Peon loop starting for project '%s'", project_id)


def ensure_api_ready() -> None:
    """Wait for the REST API to become available (up to 30s with backoff).

    Used on startup — in Kubernetes the API pod may still be starting.
    """
    api_url = _get_api_url()
    deadline = time.monotonic() + 30
    delay = 0.5
    while True:
        result = _http_request("GET", f"{api_url}/api/health", timeout=3, quiet_statuses={500})
        if result and result.get("status") == "ok":
            log.info("API is ready at %s", api_url)
            return
        if time.monotonic() >= deadline:
            log.error("API not reachable at %s after 30s — continuing anyway", api_url)
            return
        log.info("Waiting for API at %s (retrying in %.1fs)...", api_url, delay)
        time.sleep(delay)
        delay = min(delay * 2, 5)


def lookup_github_url(project_id: str) -> str:
    """Look up the github_url for a project via the REST API.

    Returns the URL string or empty string if not found.
    """
    api_url = _get_api_url()
    result = _http_request("GET", f"{api_url}/api/projects/{project_id}", quiet_statuses={404})
    if result:
        return result.get("github_url", "") or ""
    return ""


def resolve_repo_path(project_id: str) -> str:
    """Return the canonical local repo path: ``~/.peon/repos/{owner}/{repo}``."""
    return str(Path.home() / ".peon" / "repos" / project_id)


def _inject_github_token(url: str) -> str:
    """Inject a GitHub App installation token into an HTTPS GitHub URL.

    Returns the original URL unchanged if no token is available or the URL
    is not an HTTPS GitHub URL.
    """
    from peon_mcp.github_app import _parse_github_owner, get_github_token

    owner = _parse_github_owner(url)
    token = get_github_token(owner=owner)
    if not token:
        return url
    if not url.startswith("https://github.com/"):
        return url
    return url.replace("https://github.com/", f"https://x-access-token:{token}@github.com/", 1)


def clone_repo(github_url: str, target_path: str) -> None:
    """Clone a git repository to the target path.

    Uses a GitHub App installation token for authentication.  After cloning,
    the remote URL is reset to the *plain* (token-free) URL so that the token
    is never persisted in ``.git/config``.

    Raises:
        RuntimeError: If the clone command fails.
        FileNotFoundError: If git is not installed.
    """
    if not which("git"):
        raise FileNotFoundError("git is not installed or not on PATH")

    Path(target_path).parent.mkdir(parents=True, exist_ok=True)

    auth_url = _inject_github_token(github_url)
    # Log the *clean* URL so tokens never appear in logs
    log.info("Cloning %s → %s", github_url, target_path)
    result = subprocess.run(
        ["git", "clone", auth_url, target_path],
        capture_output=True,
        text=True,
        timeout=300,  # 5-minute timeout for large repos
    )
    if result.returncode != 0:
        # Clean up partial clone directory
        shutil.rmtree(target_path, ignore_errors=True)
        # Scrub token from stderr before raising
        stderr = result.stderr.strip().replace(auth_url, github_url)
        raise RuntimeError(
            f"git clone failed (exit {result.returncode}): {stderr}"
        )

    # Reset the remote to the plain URL so the token is not persisted,
    # then configure a credential helper that supplies it on demand.
    from peon_mcp.github_app import _parse_github_owner, get_github_token as _get_gh_token

    _owner = _parse_github_owner(github_url)
    token = _get_gh_token(owner=_owner)
    if token and auth_url != github_url:
        subprocess.run(
            ["git", "remote", "set-url", "origin", github_url],
            cwd=target_path, capture_output=True, text=True,
        )
        # Store-based credential helper so fetch/push keep working
        subprocess.run(
            ["git", "config", "credential.helper",
             "!f() { echo username=x-access-token; echo password=$GH_TOKEN; }; f"],
            cwd=target_path, capture_output=True, text=True,
        )
    log.info("Clone complete.")


def _cleanup_stale_sessions(project_id: str) -> None:
    """Mark any leftover 'running' sessions as 'failed' on loop startup.

    If a previous loop run was killed (SIGKILL, OOM, etc.) the session records
    are never finished and remain in status='running' forever.  This startup
    cleanup prevents stale sessions from polluting the session list.
    """
    api_url = _get_api_url()
    result = _http_request(
        "POST",
        f"{api_url}/api/sessions/cleanup-stale",
        {"project_id": project_id},
    )
    if result and result.get("updated", 0) > 0:
        log.info("Marked %d stale running session(s) as failed on startup", result["updated"])


# How long a waiting_for_input task may sit unanswered before being timed out.
_WAITING_FOR_INPUT_TIMEOUT_HOURS = 24


def _recover_stuck_tasks(project_id: str, current_hostname: str = "") -> None:
    """Recover in_progress tasks on loop startup, respecting other live loops.

    Only recovers tasks that belong to this host, to dead hosts, or to
    hosts with no heartbeat.  Tasks owned by a live loop on a different
    host are left alone.

    Also recovers waiting_for_input tasks where the human approval has not
    arrived within _WAITING_FOR_INPUT_TIMEOUT_HOURS hours.
    """
    api_url = _get_api_url()
    # Get all in_progress tasks for this project
    result = _http_request(
        "GET",
        f"{api_url}/api/projects/{project_id}/tasks?status=in_progress&limit=100",
    )

    if result and result.get("items"):
        # Build set of alive hosts from active heartbeats (90s TTL)
        alive_hosts: set[str] = set()
        loops_result = _http_request(
            "GET",
            f"{api_url}/api/loops?project_id={project_id}",
            quiet_statuses={404},
        )
        if loops_result and loops_result.get("loops"):
            for loop in loops_result["loops"]:
                h = loop.get("hostname", "")
                if h:
                    alive_hosts.add(h)

        count = 0
        for task in result["items"]:
            task_id = task["id"]
            task_title = task.get("title", "")
            started_at = task.get("started_at", "unknown")
            host = task.get("host", "") or ""

            # Decision: should we recover this task?
            if host and host != current_hostname and host in alive_hosts:
                # Another live loop owns this task — leave it alone
                log.debug(
                    "Skipping task #%d (%s) — owned by live host '%s'",
                    task_id, task_title, host,
                )
                continue

            # Recover: our own crashed task, dead host's task, or unowned task
            if host == current_hostname:
                reason_detail = "owned by this host (previous crash)"
            elif not host:
                reason_detail = "no host recorded (legacy task)"
            else:
                reason_detail = f"owned by dead host '{host}' (no active heartbeat)"

            reason = (
                f"Recovered on loop startup — task was stuck in 'in_progress' "
                f"(started_at={started_at}, host={host}). {reason_detail}."
            )
            resp = _http_request(
                "PUT",
                f"{api_url}/api/tasks/{task_id}",
                {"status": "timeout", "status_reason": reason},
            )
            if resp:
                count += 1
                log.warning(
                    "Recovered stuck task #%d (%s) → timeout (started_at=%s, host=%s, reason=%s)",
                    task_id, task_title, started_at, host, reason_detail,
                )

        if count:
            log.info("Recovered %d stuck in_progress task(s) on startup", count)

    # Recover waiting_for_input tasks where the human never responded
    _recover_waiting_for_input_tasks(project_id)


def _recover_waiting_for_input_tasks(project_id: str) -> None:
    """Time out waiting_for_input tasks that have exceeded the approval timeout.

    A task enters waiting_for_input when an agent creates an approval request.
    If the human never responds, the task stalls indefinitely. This function
    moves such tasks to timeout after _WAITING_FOR_INPUT_TIMEOUT_HOURS hours.
    """
    api_url = _get_api_url()
    result = _http_request(
        "GET",
        f"{api_url}/api/projects/{project_id}/tasks?status=waiting_for_input&limit=100",
    )
    if not result or not result.get("items"):
        return

    now = datetime.now(tz=timezone.utc)
    cutoff_hours = _WAITING_FOR_INPUT_TIMEOUT_HOURS
    count = 0

    for task in result["items"]:
        task_id = task["id"]
        task_title = task.get("title", "")
        updated_at_raw = task.get("updated_at") or task.get("started_at")
        if not updated_at_raw:
            continue

        try:
            updated_at = datetime.fromisoformat(
                updated_at_raw.replace("Z", "+00:00") if isinstance(updated_at_raw, str) else updated_at_raw.isoformat()
            )
        except (ValueError, AttributeError):
            log.debug("Could not parse updated_at for task #%d — skipping", task_id)
            continue

        hours_waiting = (now - updated_at).total_seconds() / 3600
        if hours_waiting < cutoff_hours:
            log.debug(
                "Task #%d (%s) has been waiting_for_input for %.1fh (< %dh threshold)",
                task_id, task_title, hours_waiting, cutoff_hours,
            )
            continue

        reason = (
            f"Human approval not received within {cutoff_hours}h "
            f"(task entered waiting_for_input at {updated_at_raw})."
        )
        resp = _http_request(
            "PUT",
            f"{api_url}/api/tasks/{task_id}",
            {"status": "timeout", "status_reason": reason},
        )
        if resp:
            count += 1
            log.warning(
                "Timed out waiting_for_input task #%d (%s) after %.1fh without human response",
                task_id, task_title, hours_waiting,
            )

    if count:
        log.info("Timed out %d waiting_for_input task(s) without human approval", count)


def has_todo_tasks(project_id: str, feature_id: int | None = None) -> tuple[int, int]:
    """Check the REST API for todo/in_progress tasks.

    Args:
        project_id: Project identifier.
        feature_id: Optional feature ID to filter tasks.

    Returns:
        (todo_count, in_progress_count).  Returns (0, 0) on API failure.
    """
    api_url = _get_api_url()

    # Build query params
    base_url = f"{api_url}/api/projects/{project_id}/tasks"
    feature_param = f"&feature_id={feature_id}" if feature_id is not None else ""

    # Check todo tasks
    todo_result = _http_request(
        "GET",
        f"{base_url}?status=todo&limit=1{feature_param}",
    )
    todo = todo_result.get("total", 0) if todo_result else 0

    # Check in_progress tasks
    ip_result = _http_request(
        "GET",
        f"{base_url}?status=in_progress&limit=1{feature_param}",
    )
    in_progress = ip_result.get("total", 0) if ip_result else 0

    if in_progress:
        log.warning(
            "%d task(s) stuck in 'in_progress' for project '%s'. "
            "Consider resetting them via peon-ui before running the loop.",
            in_progress,
            project_id,
        )

    return (todo, in_progress)


def get_in_progress_task(project_id: str) -> dict | None:
    """Get the most recently updated task that is in_progress for a project.

    Returns:
        Task dict with id, title, etc., or None if no in_progress task exists.
    """
    api_url = _get_api_url()
    result = _http_request(
        "GET",
        f"{api_url}/api/projects/{project_id}/tasks?status=in_progress&sort=updated&limit=1",
    )
    if result and result.get("items"):
        return result["items"][0]
    return None


def get_task_by_id(task_id: int) -> dict | None:
    """Get a task by its ID, regardless of status.

    Args:
        task_id: Task ID to fetch.

    Returns:
        Task dict with id, title, etc., or None if task doesn't exist.
    """
    api_url = _get_api_url()
    return _http_request("GET", f"{api_url}/api/tasks/{task_id}", quiet_statuses={404})


def _get_task_started_after(project_id: str, after_timestamp: str) -> dict | None:
    """Get the most recently started task for a project that started after a given time.

    Used to link a session to the task the agent picked up when the initial
    2-second poll missed the task assignment.

    Args:
        project_id: Project identifier.
        after_timestamp: ISO/SQLite timestamp string; only tasks with
                         ``started_at >= after_timestamp`` are considered.

    Returns:
        Task dict, or None if no matching task found.
    """
    api_url = _get_api_url()
    encoded_ts = urllib.parse.quote(after_timestamp)
    result = _http_request(
        "GET",
        f"{api_url}/api/projects/{project_id}/tasks?started_after={encoded_ts}&sort=started&limit=1",
    )
    if result and result.get("items"):
        return result["items"][0]
    return None


def _get_commit_base(repo_path: str) -> str:
    """Run ``git rev-parse HEAD`` in *repo_path* and return the SHA.

    Returns empty string on failure (not a git repo, git missing, etc.).
    """
    try:
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=repo_path,
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode == 0:
            return result.stdout.strip()
    except (subprocess.TimeoutExpired, OSError, FileNotFoundError) as exc:
        log.warning("Failed to get commit base from %s: %s", repo_path, exc)
    return ""


def _claim_next_task(
    project_id: str,
    commit_base: str,
    feature_id: int | None = None,
    hostname: str = "",
) -> dict | None:
    """Claim the next todo task for *project_id* before spawning an agent.

    Calls ``POST /api/projects/{project_id}/next-task`` which uses
    ``SELECT ... FOR UPDATE SKIP LOCKED`` to guarantee only one loop
    wins the claim.  Returns the full task dict on success, ``None``
    when no task is available (404) or on concurrency conflict (409).
    """
    api_url = _get_api_url()
    payload: dict = {"commit_base": commit_base, "host": hostname}
    if feature_id is not None:
        payload["feature_id"] = feature_id
    return _http_request(
        "POST",
        f"{api_url}/api/projects/{project_id}/next-task",
        payload,
        timeout=10,
        quiet_statuses={404, 409},
    )


def _claim_next_task_global(
    feature_id: int | None = None,
    hostname: str = "",
) -> dict | None:
    """Claim the next todo task across all projects.

    Calls ``POST /api/next-task`` (global endpoint).  Returns the full
    task dict on success, ``None`` when no task is available.
    """
    api_url = _get_api_url()
    payload: dict = {"host": hostname}
    if feature_id is not None:
        payload["feature_id"] = feature_id
    return _http_request(
        "POST",
        f"{api_url}/api/next-task",
        payload,
        timeout=10,
        quiet_statuses={404, 409},
    )


# ---------------------------------------------------------------------------
# Heartbeat
# ---------------------------------------------------------------------------


def write_heartbeat(
    loop_id: str,
    project_id: str | None,
    pid: int,
    hostname: str,
) -> None:
    """Upsert a heartbeat row via the REST API so the UI can show this loop as active."""
    api_url = _get_api_url()
    _http_request(
        "PUT",
        f"{api_url}/api/loops/{loop_id}",
        {
            "project_id": project_id,
            "pid": pid,
            "hostname": hostname,
        },
    )


def delete_heartbeat(loop_id: str) -> None:
    """Remove the heartbeat row on clean shutdown."""
    api_url = _get_api_url()
    _http_request("DELETE", f"{api_url}/api/loops/{loop_id}", quiet_statuses={404})


class HeartbeatThread(threading.Thread):
    """Background thread that writes heartbeats at regular intervals.

    This ensures the loop appears active even when the agent subprocess
    is running a long task.
    """

    def __init__(
        self,
        loop_id: str,
        project_id: str | None,
        pid: int,
        hostname: str,
        interval: int = 30,
    ):
        super().__init__(daemon=True, name="heartbeat-thread")
        self.loop_id = loop_id
        self.project_id = project_id
        self.pid = pid
        self.hostname = hostname
        self.interval = interval
        self._stop_event = threading.Event()
        self.consecutive_failures = 0
        self._max_consecutive_failures = 3

    def run(self) -> None:
        """Write heartbeats every interval seconds until stopped."""
        # Write immediate heartbeat on startup
        try:
            write_heartbeat(
                self.loop_id,
                self.project_id,
                self.pid,
                self.hostname,
            )
            self.consecutive_failures = 0
        except Exception as exc:
            self.consecutive_failures += 1
            log.warning("Initial heartbeat failed to write: %s", exc)

        # Continue with regular interval-based heartbeats
        while not self._stop_event.wait(self.interval):
            try:
                write_heartbeat(
                    self.loop_id,
                    self.project_id,
                    self.pid,
                    self.hostname,
                )
                self.consecutive_failures = 0
            except Exception as exc:
                self.consecutive_failures += 1
                log.warning("Heartbeat thread failed to write: %s", exc)

    @property
    def healthy(self) -> bool:
        """Returns True if the heartbeat thread hasn't exceeded max consecutive failures."""
        return self.consecutive_failures < self._max_consecutive_failures

    def stop(self) -> None:
        """Signal the thread to stop and wait for it to finish."""
        self._stop_event.set()
        self.join(timeout=5)


# ---------------------------------------------------------------------------
# Binary discovery
# ---------------------------------------------------------------------------


def find_peon_mcp_binary() -> str:
    """Locate the peon-mcp binary.

    Checks next to the current Python interpreter first (same venv), then
    falls back to $PATH.
    """
    venv_bin = Path(sys.executable).parent / "peon-mcp"
    if venv_bin.is_file():
        return str(venv_bin)

    found = which("peon-mcp")
    if found:
        return found

    raise FileNotFoundError(
        "Could not find 'peon-mcp' binary. "
        "Is peon-mcp installed in this environment?"
    )


def find_claude_binary() -> str:
    """Locate the claude CLI binary."""
    found = which("claude")
    if found:
        return found
    raise FileNotFoundError(
        "Could not find 'claude' CLI. "
        "Install it from https://claude.ai/download"
    )


# ---------------------------------------------------------------------------
# MCP config & prompt
# ---------------------------------------------------------------------------


def build_mcp_config(peon_bin: str) -> str:
    """Build inline JSON for --mcp-config.

    Passes API connection settings as env vars so the MCP server can reach
    the REST API via PeonAPIClient.
    """
    env: dict[str, str] = {}
    api_url = os.environ.get("PEON_API_URL")
    if api_url:
        env["PEON_API_URL"] = api_url
    api_key = os.environ.get("PEON_API_KEY")
    if api_key:
        env["PEON_API_KEY"] = api_key
    config = {
        "mcpServers": {
            "peon-mcp": {
                "command": peon_bin,
                "env": env,
            }
        }
    }
    return json.dumps(config)


def build_mcp_config_dict(peon_bin: str) -> dict:
    """Build MCP server config as a dict for the Claude Agent SDK.

    Unlike ``build_mcp_config()`` which returns a JSON string for the CLI
    ``--mcp-config`` flag, this returns a Python dict suitable for the
    ``mcp_servers`` parameter of ``ClaudeAgentOptions``.
    """
    env: dict[str, str] = {}
    api_url = os.environ.get("PEON_API_URL")
    if api_url:
        env["PEON_API_URL"] = api_url
    api_key = os.environ.get("PEON_API_KEY")
    if api_key:
        env["PEON_API_KEY"] = api_key
    return {
        "peon-mcp": {
            "command": peon_bin,
            "env": env,
        }
    }


_IMAGE_CONTENT_TYPES = {"image/png", "image/jpeg", "image/gif", "image/webp"}


def _fetch_and_download_attachments(task_id: int, download_dir: Path) -> list[dict]:
    """Fetch attachment metadata and download files for a task.

    Returns a list of dicts with keys: id, filename, content_type, size_bytes,
    local_path. Attachments that fail to download are silently skipped.
    Returns an empty list if S3 is not configured (HTTP 501) or no attachments exist.
    """
    api_url = _get_api_url()
    attachments = _http_request(
        "GET",
        f"{api_url}/api/tasks/{task_id}/attachments",
        quiet_statuses={404, 501},
    )
    if not attachments or not isinstance(attachments, list):
        return []

    download_dir.mkdir(parents=True, exist_ok=True)
    results = []
    seen_filenames: set[str] = set()

    for att in attachments:
        att_id = att.get("id")
        filename = att.get("filename") or f"attachment_{att_id}"
        content_type = att.get("content_type", "application/octet-stream")
        size_bytes = att.get("size_bytes", 0)
        download_url = att.get("download_url", "")

        if not download_url:
            log.warning("Attachment %s (%s) has no download_url, skipping", att_id, filename)
            continue

        # Deduplicate filenames
        local_name = filename
        if local_name in seen_filenames:
            stem, _, suffix = filename.rpartition(".")
            local_name = f"{stem}_{att_id}.{suffix}" if suffix else f"{filename}_{att_id}"
        seen_filenames.add(local_name)
        local_path = download_dir / local_name

        try:
            req = urllib.request.Request(download_url)
            with urllib.request.urlopen(req, timeout=30) as resp:
                data = resp.read()
            local_path.write_bytes(data)
            results.append(
                {
                    "id": att_id,
                    "filename": filename,
                    "content_type": content_type,
                    "size_bytes": size_bytes,
                    "local_path": str(local_path),
                }
            )
            log.info("Downloaded attachment %s -> %s (%d bytes)", filename, local_path, len(data))
        except (urllib.error.URLError, OSError) as exc:
            log.warning("Failed to download attachment %s (%s): %s", att_id, filename, exc)

    return results


def build_prompt(
    project_id: str,
    repo_path: str,
    feature_id: int | None = None,
    task: dict | None = None,
    pre_claimed_task: dict | None = None,
    downloaded_attachments: list[dict] | None = None,
) -> str:
    """Build the agent prompt — mirrors the start_work prompt from server.py.

    When ``pre_claimed_task`` is provided, the prompt tells the agent that
    the task has already been claimed and it should NOT call ``next_task``.

    When ``task`` (or ``pre_claimed_task``) is provided, appends
    task-type-specific instructions for ``fix``, ``verify``, etc.
    """
    # If pre_claimed_task is provided, use it as the task for type-specific instructions
    if pre_claimed_task is not None:
        task = pre_claimed_task

    feature_clause = f", feature_id={feature_id}" if feature_id is not None else ""
    prompt = (
        "You are a Peon — a hardworking Orc laborer from Warcraft. You speak in Peon mannerisms "
        "('Zug zug', 'Job's done!', 'Work work', 'Dabu'). Your work logs should reflect this persona, "
        "but your actual code and commits must remain professional and high quality. The persona lives "
        "in your summaries and logs, NOT in the code you write for the project.\n"
        "\n"
        f"You are a development agent about to start work on project '{project_id}' "
        f"located at {repo_path}.\n"
        "\n"
        "Follow these steps in order:\n"
        "\n"
        f"1. Call get_project(project_id='{project_id}'). If the project does not "
        f"exist, call create_project(project_id='{project_id}', path='{repo_path}') "
        "to register it.\n"
        "\n"
        f"2. **Recall project memories** — Call recall_memories(project_id='{project_id}') to load "
        "context from previous agents. Review the returned memories for environment quirks, "
        "patterns, gotchas, and test insights before starting work.\n"
        "\n"
        f"3. **Ensure main is clean** — Change to {repo_path}. Run `git status`. "
        "If there are untracked or uncommitted changes, run `git stash --include-untracked` "
        "to stash them. The main repo should remain clean since all work happens in worktrees.\n"
    )

    if pre_claimed_task is not None:
        # Task already claimed by the loop — skip git rev-parse and next_task
        task_id = pre_claimed_task.get("id", "?")
        task_title = pre_claimed_task.get("title", "")
        task_desc = pre_claimed_task.get("description", "")
        feature_branch = pre_claimed_task.get("feature_branch", "")
        fb_clause = f", feature_branch='{feature_branch}'" if feature_branch else ""
        prompt += (
            "\n"
            f"4. **Task #{task_id} has been pre-claimed for you.** Do NOT call next_task(). "
            f"The task is already marked as in_progress.\n"
            f"   - **Title**: {task_title}\n"
            f"   - **Description**: {task_desc}\n"
            "\n"
            f"5. Use the execute_task prompt with project_id='{project_id}', "
            f"task_id={task_id}, title='{task_title}', "
            f"description (from above){fb_clause}. "
            "This ensures the task branches from and PRs to the correct feature branch.\n"
        )
    else:
        prompt += (
            "\n"
            f"4. Run `git rev-parse HEAD` in {repo_path} to capture the current commit SHA.\n"
            "\n"
            f"5. Call next_task(project_id='{project_id}', commit_base='<sha>'{feature_clause}, host='{socket.gethostname()}') passing "
            "the SHA from step 4. This marks the task as in_progress, records the base "
            "commit, and returns the task details.\n"
            "\n"
            f"6. Use the execute_task prompt with project_id='{project_id}', the task's "
            "title and description, and if the task response includes a 'feature_branch' field, "
            "pass it as the feature_branch parameter. This ensures the task branches from and PRs "
            "to the correct feature branch.\n"
        )

    prompt += (
        "\n"
        "**IMPORTANT**: Once you start work, you MUST use update_task_progress() frequently "
        "throughout the task to keep stakeholders informed. Update progress after every major "
        "step (reading files, writing code, running tests, etc.). Aim for at least 3-4 progress "
        "updates per task.\n"
        "\n"
        "Note: You are already running in an isolated git worktree automatically created by peon-loop. "
        "The execute_task prompt will provide detailed instructions for working in this worktree.\n"
        "\n"
        "**MEMORY**: If you discover something useful during your work (environment quirks, "
        "non-obvious patterns, test gotchas, dependency issues), save it with "
        f"save_memory(project_id='{project_id}', category='...', content='...'). "
        "Keep memories concise and actionable — future agents will benefit.\n"
        "\n"
        "**TASK TYPES**: The task you pick up may have a `task_type` field indicating special "
        "pipeline context. Always read the task description fully — it contains the authoritative "
        "instructions for the specific work. Task types you may encounter:\n"
        "- `fix`: A FIX task. Apply specific findings, keep scope tight, and create a PR with "
        "focused commits.\n"
        "- `verify`: A VERIFY task — the final gate. Run the full test suite and only mark done "
        "if everything passes.\n"
        "- `review`: A PR REVIEW task — either conduct a code review (approve/request-changes) "
        "or address reviewer feedback on an existing PR.\n"
        "- Standard tasks (no task_type): Normal development tasks — follow the execute_task prompt."
    )
    if task:
        task_type = task.get("task_type")
        if task_type == "fix":
            prompt += (
                "\n\n"
                "## Fix Task Instructions\n"
                "You are working on a FIX task. Your job is to:\n"
                "1. Read the task description carefully — it contains specific findings to address\n"
                "2. Apply ONLY the fixes described in the findings — keep scope tight\n"
                "3. Each fix should be a focused commit: `fix: <category> in <file_path>`\n"
                "4. Do NOT expand scope beyond what the findings describe\n"
                "5. After applying fixes, create a PR (or push to existing branch)\n"
                "6. Run tests to verify your fix doesn't break anything\n"
                "7. Mark the task done with the PR URL\n"
                "\n"
                "Tight scope. Surgical fixes. No refactoring beyond what's needed."
            )
        elif task_type == "verify":
            prompt += (
                "\n\n"
                "## Verify Task Instructions\n"
                "You are working on a VERIFY task — the final gate.\n"
                "1. Checkout the branch where fixes were applied\n"
                "2. Pull/merge all fix PRs into the branch\n"
                "3. Run the full test suite\n"
                "4. If tests pass: mark this task done\n"
                "5. If tests fail:\n"
                "   - Identify which fix introduced the regression\n"
                "   - Report the failure in your work log\n"
                "   - Do NOT mark this task done — leave it for retry\n"
                "\n"
                "You are the gate. Only pass if everything is green."
            )
        elif task_type == "incident":
            log_query = task.get("log_query", "")
            sample_lines = task.get("sample_log_lines", "")
            occurrence_count = task.get("occurrence_count", 0)
            parts = [
                "\n\n"
                "## Incident Task Instructions\n"
                "You are working on an INCIDENT task — a production issue that needs investigation and a fix.\n"
            ]
            if occurrence_count:
                parts.append(f"- **Occurrences**: {occurrence_count}\n")
            if log_query:
                parts.append(f"- **Log query**: `{log_query}`\n")
            if sample_lines:
                parts.append(f"- **Sample log lines**:\n```\n{sample_lines}\n```\n")
            parts.append(
                "\n"
                "1. Read the task description and any attached log data\n"
                "2. Identify the root cause from the logs and code\n"
                "3. Implement a targeted fix\n"
                "4. Add or update tests to cover the failure scenario\n"
                "5. Run the full test suite\n"
                "6. Create a PR and mark the task done\n"
            )
            prompt += "".join(parts)
        elif task_type == "review":
            task_title = task.get("title", "")
            pr_url = task.get("pr_url", "")
            if task_title.startswith("Code review PR #"):
                prompt += (
                    "\n\n"
                    "## Code Review Task Instructions\n"
                    "You are conducting a code review of a pull request. Your job is to:\n"
                    "1. Read the task description — it contains the PR details\n"
                    "2. Fetch the PR diff and details using the commands in the description\n"
                    "3. Review the code for bugs, security issues, performance, and maintainability\n"
                    "4. Submit a **formal GitHub review** using `--approve` or `--request-changes`\n"
                    "   - **Approve** if the code is correct, safe, and ready to merge\n"
                    "   - **Request changes** if there are bugs, security issues, or logic errors\n"
                    "   - Do NOT use `--comment` — it does not count as a formal review\n"
                    "5. Mark the task done once your review is submitted\n"
                    "\n"
                    "Be thorough — review every changed file, not just the obvious ones."
                )
            else:
                prompt += (
                    "\n\n"
                    "## PR Feedback Task Instructions\n"
                    "A GitHub reviewer has requested changes on a pull request. Your job is to:\n"
                    "1. Read the task description — it contains the reviewer's comments\n"
                    f"2. Fetch the full review comments: `gh pr view {pr_url} --comments` "
                    "and `gh api` for inline review comments\n"
                    "3. Check out the PR branch (it's already set on this task)\n"
                    "4. Address each piece of feedback from the reviewer\n"
                    "5. Commit your fixes and push to the EXISTING PR branch — do NOT create a new PR\n"
                    "6. Run tests to make sure nothing is broken\n"
                    "7. Mark the task done once all feedback is addressed\n"
                    "\n"
                    "Be thorough — read every inline comment, not just the review summary."
                )
    if downloaded_attachments:
        n = len(downloaded_attachments)
        prompt += (
            f"\n\n**Attachments**: This task has {n} attachment(s) downloaded to your workspace:\n"
        )
        for att in downloaded_attachments:
            filename = att.get("filename", "unknown")
            content_type = att.get("content_type", "")
            local_path = att.get("local_path", "")
            is_image = content_type in _IMAGE_CONTENT_TYPES
            kind = "image" if is_image else content_type
            prompt += f"- `{local_path}` ({filename}, {kind})\n"
        prompt += (
            "\nThese files are accessible to your tools. "
            "View images with your image viewing capability; read text files with Read."
        )
    return prompt


# ---------------------------------------------------------------------------
# Process management
# ---------------------------------------------------------------------------


def kill_previous(pid_file: Path) -> None:
    """Kill a previously-spawned agent if its PID file exists and it's alive."""
    if not pid_file.exists():
        return

    try:
        pid = int(pid_file.read_text().strip())
    except (ValueError, OSError):
        pid_file.unlink(missing_ok=True)
        return

    try:
        os.kill(pid, 0)  # Check if alive
    except OSError:
        # Process is gone
        pid_file.unlink(missing_ok=True)
        return

    log.info("Killing previous agent process %d", pid)
    try:
        os.kill(pid, signal.SIGTERM)
        # Give it a moment to exit
        for _ in range(10):
            time.sleep(1)
            try:
                os.kill(pid, 0)
            except OSError:
                break
        else:
            log.warning("Process %d did not exit after SIGTERM, sending SIGKILL", pid)
            os.kill(pid, signal.SIGKILL)
    except OSError:
        pass
    finally:
        pid_file.unlink(missing_ok=True)


def _mark_task_timeout(task_id: int, reason: str, task_title: str = "") -> None:
    """Mark a specific task as 'timeout' with a reason.

    Args:
        task_id: The task ID to mark as timed out.
        reason: Human-readable explanation of why the task timed out.
        task_title: Optional task title for log messages.
    """
    api_url = _get_api_url()
    log.warning(
        "TIMEOUT task #%d%s — reason: %s",
        task_id,
        f" ({task_title})" if task_title else "",
        reason,
    )
    result = _http_request(
        "PUT",
        f"{api_url}/api/tasks/{task_id}",
        {"status": "timeout", "status_reason": reason},
    )
    if result:
        log.info("Task #%d status updated to 'timeout' (reason persisted)", task_id)
    else:
        log.error("Failed to update task #%d status to timeout", task_id)


class _StreamProcessorThread(threading.Thread):
    """Background thread that tees agent stdout to a log file while detecting loops.

    Feeds parsed tool call events to a ``LoopDetector``. Expects stdout to emit
    newline-delimited JSON events (``stream-json`` format from the Claude CLI).
    On a circuit-breaker trigger, calls ``on_circuit_breaker`` and stops consuming
    output.
    """

    def __init__(
        self,
        stdout,
        log_file,
        detector,
        on_circuit_breaker,
        session_tracker: SessionTracker | None = None,
    ) -> None:
        super().__init__(daemon=True, name="loop-detector")
        self._stdout = stdout
        self._log_file = log_file
        self._detector = detector
        self._on_circuit_breaker = on_circuit_breaker
        self._session_tracker = session_tracker
        # Buffer for accumulating partial tool_use input JSON per block index
        self._pending: dict[int, dict] = {}

    def run(self) -> None:
        try:
            from peon_mcp.loop_detection import LoopSeverity
        except ImportError:
            # Drain stdout to log without detection
            for raw_line in self._stdout:
                self._log_file.write(raw_line)
                self._log_file.flush()
            return

        try:
            for raw_line in self._stdout:
                self._log_file.write(raw_line)
                self._log_file.flush()

                line = raw_line.strip()
                if not line:
                    continue

                try:
                    event = json.loads(line)
                except json.JSONDecodeError:
                    continue

                tool_name, tool_params = self._process_event(event)
                if tool_name is None:
                    continue

                if self._session_tracker is not None:
                    self._session_tracker.event("tool_call", tool_name=tool_name)

                result = self._detector.record_tool_call(tool_name, tool_params or {})
                if result.detected and result.severity == LoopSeverity.CIRCUIT_BREAKER:
                    self._on_circuit_breaker()
                    break
        except Exception as exc:
            log.debug("Stream processor encountered error: %s", exc)

    def _process_event(self, event: dict) -> tuple[str | None, dict | None]:
        """Extract a complete (tool_name, params) pair from a stream-json event.

        Tool calls span three event types:
          1. ``content_block_start`` with ``content_block.type == "tool_use"``
             — provides the tool name and block index.
          2. ``content_block_delta`` with ``delta.type == "input_json_delta"``
             — accumulates partial parameter JSON.
          3. ``content_block_stop`` — signals that the block is complete;
             we parse the accumulated JSON and return the full call.

        Handles the alternative ``assistant`` envelope format as well.
        """
        event_type = event.get("type")

        if event_type == "content_block_start":
            block = event.get("content_block", {})
            if block.get("type") == "tool_use":
                idx = event.get("index", 0)
                self._pending[idx] = {"name": block.get("name", ""), "json": ""}
            return None, None

        if event_type == "content_block_delta":
            delta = event.get("delta", {})
            if delta.get("type") == "input_json_delta":
                idx = event.get("index", 0)
                if idx in self._pending:
                    self._pending[idx]["json"] += delta.get("partial_json", "")
            return None, None

        if event_type == "content_block_stop":
            idx = event.get("index", 0)
            if idx in self._pending:
                entry = self._pending.pop(idx)
                params: dict = {}
                if entry["json"]:
                    try:
                        params = json.loads(entry["json"])
                    except json.JSONDecodeError:
                        pass
                return entry["name"], params
            return None, None

        # Fallback: some Claude CLI versions wrap events in an "assistant" envelope
        if event_type == "assistant":
            msg = event.get("message", {})
            for block in msg.get("content", []):
                if isinstance(block, dict) and block.get("type") == "tool_use":
                    return block.get("name", ""), block.get("input", {})

        return None, None


def run_agent(
    cmd: list[str],
    repo_path: str,
    log_dir: Path,
    pid_file: Path,
    timeout: int,
    project_id: str,
    extra_env: dict[str, str] | None = None,
    loop_detection: bool = True,
    loop_id: str = "",
    model: str = "",
    hostname: str = "",
    session_tracker: SessionTracker | None = None,
    github_url: str = "",
    pre_claimed_task_id: int | None = None,
    use_nyx_identity: bool = False,
) -> tuple[int, int | None]:
    """Spawn the claude agent subprocess.

    Logs stdout/stderr to a timestamped file. Writes a PID file so
    subsequent iterations can kill a stale agent. Returns a tuple of
    (exit_code, task_id) where task_id is the ID of the task this agent
    claimed, or None if no task was identified.

    When ``loop_detection`` is True, adds ``--output-format stream-json`` to
    the command, reads the streaming output in a background thread to detect
    repetitive tool-call patterns, and terminates the agent on circuit breaker.

    Args:
        cmd: Base command list to execute (e.g. ``["claude", "-p", ...]``).
        repo_path: Working directory for the subprocess.
        log_dir: Directory where timestamped agent log files are written.
        pid_file: Path to write the child process PID for stale-process cleanup.
        timeout: Maximum wall-clock seconds before the subprocess is killed.
        project_id: Project identifier, passed to the circuit-breaker logger.
        extra_env: Additional environment variables to set for the agent process.
                   Merged on top of the current process environment.
        loop_detection: If True, enable streaming JSON output parsing and
                        circuit-breaker termination on repetitive tool calls.
        use_nyx_identity: If True, inject Nyx's GitHub App token (``GH_TOKEN``)
                          instead of the standard peon token.  Falls back to the
                          peon token when Nyx is not configured.
    """
    global _child_proc
    log_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    log_file_path = log_dir / f"agent-{timestamp}.log"

    log.info("Work work! Sending peon to the mines...")
    log.info("Agent log: %s", log_file_path)

    # Record the start time before spawning so we can find the task the agent
    # claimed even if the initial 2-second poll misses it.
    agent_start_ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")

    # Create the session record before spawning the agent.
    # When pre_claimed_task_id is set, pass it at creation time so the session
    # starts with the task already linked — no follow-up PUT needed.
    if session_tracker is not None:
        session_tracker.start(
            project_id, loop_id, model,
            hostname=hostname, task_id=pre_claimed_task_id,
        )
        detail = f"Spawning agent for project {project_id}"
        if pre_claimed_task_id is not None:
            detail += f" (pre-claimed task #{pre_claimed_task_id})"
        session_tracker.event("start", detail=detail)

    # When loop detection is active, stream output through a pipe so we can
    # parse tool call events in real time.
    actual_cmd = list(cmd)
    if loop_detection:
        actual_cmd += ["--output-format", "stream-json", "--verbose"]

    # Inject session_id into the prompt so the agent can pass it to execute_task
    if session_tracker is not None and session_tracker.session_id is not None:
        try:
            p_idx = actual_cmd.index("-p")
            actual_cmd[p_idx + 1] += (
                f"\n\n**Session tracking**: Your session_id is {session_tracker.session_id}. "
                f"When calling the execute_task prompt, pass session_id={session_tracker.session_id} "
                "to enable session event tracking."
            )
        except (ValueError, IndexError):
            pass

    circuit_breaker_fired = threading.Event()
    agent_env = {**os.environ, **(extra_env or {})}

    # Inject a GitHub token into the agent subprocess environment.
    # get_github_token() already checks expiry (with a 10-minute buffer)
    # and only generates a new token if the cached one is near-expired,
    # so there's no need to invalidate the cache here.
    #
    # For code-review tasks (use_nyx_identity=True), prefer Nyx's identity
    # so that review comments appear under Nyx's GitHub App avatar.  Falls
    # back to the standard peon token when Nyx is not configured.
    from peon_mcp.github_app import (
        _parse_github_owner,
        get_github_token as _get_gh_token,
        get_nyx_github_token as _get_nyx_token,
        is_nyx_app_auth as _nyx_auth_available,
    )

    _owner = _parse_github_owner(github_url) if github_url else None
    if use_nyx_identity and _nyx_auth_available():
        _fresh_token = _get_nyx_token(owner=_owner)
        if not _fresh_token:
            # Nyx configured but token generation failed — fall back to peon token
            _fresh_token = _get_gh_token(owner=_owner)
    else:
        _fresh_token = _get_gh_token(owner=_owner)
    if _fresh_token:
        agent_env["GH_TOKEN"] = _fresh_token

    # Pass project_id and session_id as env vars so hooks can find the current task
    agent_env["PEON_PROJECT_ID"] = project_id
    if session_tracker is not None and session_tracker.session_id is not None:
        agent_env["PEON_SESSION_ID"] = str(session_tracker.session_id)
    if pre_claimed_task_id is not None:
        agent_env["PEON_TASK_ID"] = str(pre_claimed_task_id)

    with open(log_file_path, "w") as log_file:
        if loop_detection:
            proc = subprocess.Popen(
                actual_cmd,
                cwd=repo_path,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
                env=agent_env,
            )
        else:
            proc = subprocess.Popen(
                actual_cmd,
                cwd=repo_path,
                stdout=log_file,
                stderr=subprocess.STDOUT,
                env=agent_env,
            )

        _child_proc = proc
        pid_file.write_text(str(proc.pid))

        task_id = pre_claimed_task_id

        # Start stream processor thread for loop detection
        stream_processor: _StreamProcessorThread | None = None
        if loop_detection and proc.stdout is not None:
            try:
                from peon_mcp.loop_detection import LoopDetector

                loop_detector = LoopDetector()

                def _on_circuit_breaker() -> None:
                    circuit_breaker_fired.set()
                    log.error("Loop detection circuit breaker! Terminating agent process.")
                    try:
                        proc.terminate()
                    except OSError:
                        pass

                stream_processor = _StreamProcessorThread(
                    stdout=proc.stdout,
                    log_file=log_file,
                    detector=loop_detector,
                    on_circuit_breaker=_on_circuit_breaker,
                    session_tracker=session_tracker,
                )
                stream_processor.start()
            except ImportError:
                log.debug("Loop detection module not available; disabling for this run")
                loop_detection = False

        # Start an early task-link thread: polls for the task the agent claimed
        # and links the session as soon as it's found, rather than waiting until
        # teardown.  Only needed when task was NOT pre-claimed (backward compat
        # path where the agent calls next_task itself).
        _stop_early_link = threading.Event()
        early_link_thread = None

        if pre_claimed_task_id is None:
            def _early_link_worker() -> None:
                nonlocal task_id
                for wait_secs in [5, 10, 15, 30, 60]:
                    if _stop_early_link.wait(timeout=wait_secs):
                        return  # Agent exited or task already linked
                    if task_id is not None:
                        return  # Already linked (e.g. by session_link.py hook)
                    try:
                        found = get_in_progress_task(project_id)
                        if not found:
                            found = _get_task_started_after(project_id, agent_start_ts)
                        if found:
                            task_id = found["id"]
                            if session_tracker is not None:
                                session_tracker.link_task(task_id)
                            return
                    except Exception as exc:
                        log.debug("Early task-link poll failed (non-fatal): %s", exc)

            early_link_thread = threading.Thread(
                target=_early_link_worker, daemon=True, name="early-task-link"
            )
            early_link_thread.start()

        timed_out = False
        agent_wall_start = time.monotonic()
        try:
            proc.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            timed_out = True
            elapsed = int(time.monotonic() - agent_wall_start)
            log.warning("Peon took too long! Giving up after %ds (wall-clock: %ds)", timeout, elapsed)
            proc.terminate()
            try:
                proc.wait(timeout=10)
            except subprocess.TimeoutExpired:
                log.warning("Agent did not exit after SIGTERM, sending SIGKILL")
                proc.kill()
                proc.wait()

        # Wait for the stream processor to drain remaining output
        if stream_processor is not None:
            stream_processor.join(timeout=5)

        # Stop the early task-link thread now that the agent has exited
        _stop_early_link.set()
        if early_link_thread is not None:
            early_link_thread.join(timeout=5)

    # --- Task discovery ---
    # Identify which task THIS agent claimed.  Uses _get_task_started_after
    # (scoped to tasks started after this agent spawned) rather than
    # get_in_progress_task (which could return a task owned by a concurrent
    # agent in a different loop instance).
    if task_id is None:
        try:
            task = _get_task_started_after(project_id, agent_start_ts)
            if task:
                task_id = task["id"]
                log.info(
                    "Post-process task discovery: peon was working on task #%d: %s",
                    task_id,
                    task.get("title", ""),
                )
        except Exception as exc:
            log.warning("Post-process task discovery failed (non-fatal): %s", exc)

    # --- Timeout / circuit-breaker marking ---
    # Only mark the specific task THIS agent was working on.  Previous code
    # used _mark_current_task_timeout (which finds ANY in_progress task for
    # the project) and could accidentally timeout a task owned by a concurrent
    # agent in another loop pod.
    elapsed = int(time.monotonic() - agent_wall_start)
    if task_id is not None:
        if timed_out:
            _mark_task_timeout(
                task_id,
                reason=f"Agent exceeded wall-clock timeout ({timeout}s limit, ran for {elapsed}s)",
            )
        elif circuit_breaker_fired.is_set():
            _mark_task_timeout(
                task_id,
                reason=f"Loop detection circuit breaker fired after {elapsed}s — agent was stuck in repetitive tool-call pattern",
            )
    elif timed_out or circuit_breaker_fired.is_set():
        log.warning(
            "Agent timed out or circuit breaker fired, but could not identify "
            "which task was being worked on — skipping timeout marking to avoid "
            "affecting tasks owned by other agents."
        )

    # Record session outcome
    if session_tracker is not None:
        log_stats = _parse_agent_log(log_file_path)
        if timed_out:
            session_tracker.finish(
                "timeout",
                exit_code=proc.returncode,
                error_summary=f"Agent exceeded wall-clock timeout ({timeout}s limit, ran for {elapsed}s)",
                task_id=task_id,
                **log_stats,
            )
        elif circuit_breaker_fired.is_set():
            session_tracker.finish(
                "timeout",
                exit_code=proc.returncode,
                error_summary=f"Circuit breaker fired after {elapsed}s — repetitive tool-call pattern detected",
                task_id=task_id,
                **log_stats,
            )
        elif proc.returncode == 0:
            session_tracker.finish("completed", exit_code=0, task_id=task_id, **log_stats)
        else:
            session_tracker.finish(
                "failed",
                exit_code=proc.returncode,
                error_summary=f"Agent exited with code {proc.returncode} after {elapsed}s",
                task_id=task_id,
                **log_stats,
            )

    # Create a fallback work log if the agent skipped calling log_work
    if task_id is not None:
        _create_fallback_work_log(_get_api_url(), project_id, task_id)

    _child_proc = None
    pid_file.unlink(missing_ok=True)
    elapsed = int(time.monotonic() - agent_wall_start)
    log.info(
        "Job's done! Agent exited with code %d after %ds (timeout_limit=%ds, timed_out=%s, circuit_breaker=%s)",
        proc.returncode, elapsed, timeout, timed_out, circuit_breaker_fired.is_set(),
    )
    return proc.returncode, task_id


# ---------------------------------------------------------------------------
# Claude Agent SDK path  (--use-sdk)
# ---------------------------------------------------------------------------


def _make_directive_hook(loop_id: str, api_url: str):
    """Create a PreToolUse hook that checks for directives from Grukk.

    Returns an async callable compatible with the Claude Agent SDK hook system.
    Polls ``GET /api/loops/{loop_id}/directives?status=pending`` before each
    tool call.  Handles pause (blocking), stop (block tool), answer (log),
    and redirect (block tool) directive types.
    """

    async def check_directives(input_data, tool_use_id, context):
        directives = _http_request(
            "GET",
            f"{api_url}/api/loops/{loop_id}/directives?status=pending",
        )
        if not directives:
            return {}

        for d in (directives if isinstance(directives, list) else []):
            directive_id = d.get("id")
            directive_type = d.get("directive_type")

            # Acknowledge the directive
            _http_request(
                "PATCH",
                f"{api_url}/api/directives/{directive_id}",
                {"status": "acknowledged"},
            )

            if directive_type == "pause":
                # Block until a resume directive arrives
                while True:
                    await asyncio.sleep(5)
                    resume = _http_request(
                        "GET",
                        f"{api_url}/api/loops/{loop_id}/directives?status=pending",
                    )
                    if resume and any(
                        r.get("directive_type") == "resume" for r in resume
                    ):
                        for r in resume:
                            if r.get("directive_type") == "resume":
                                _http_request(
                                    "PATCH",
                                    f"{api_url}/api/directives/{r['id']}",
                                    {"status": "acknowledged"},
                                )
                        break

            elif directive_type == "stop":
                return {"decision": "block", "reason": "Stopped by Grukk"}

            elif directive_type == "answer":
                log.info("Received answer from Grukk: %s", d.get("payload", {}))

            elif directive_type == "redirect":
                log.info("Redirect directive received: %s", d.get("payload", {}))
                return {"decision": "block", "reason": "Redirected by Grukk"}

        return {}

    return check_directives


def _make_heartbeat_hook(
    loop_id: str,
    project_id: str,
    pid: int,
    hostname: str,
    session_tracker: SessionTracker | None,
):
    """Create a PostToolUse hook that writes heartbeats and captures session events.

    Throttles heartbeat writes to at most once every 30 seconds.
    Also records tool_call events on the session tracker when available.
    """

    _last_heartbeat = [0.0]  # mutable container for closure
    HEARTBEAT_INTERVAL = 30  # seconds

    async def heartbeat_and_events(input_data, tool_use_id, context):
        now = time.time()

        # Write heartbeat if enough time has passed
        if now - _last_heartbeat[0] >= HEARTBEAT_INTERVAL:
            write_heartbeat(loop_id, project_id, pid, hostname)
            _last_heartbeat[0] = now

        # Capture session event
        if session_tracker is not None:
            tool_name = "unknown"
            if isinstance(input_data, dict):
                tool_name = input_data.get("tool_name", "unknown")
            session_tracker.event("tool_call", tool_name=tool_name)

        return {}

    return heartbeat_and_events


async def _heartbeat_watchdog(
    loop_id: str,
    project_id: str,
    pid: int,
    hostname: str,
    interval: int = 30,
    stop_event: asyncio.Event | None = None,
) -> None:
    """Fallback heartbeat writer for when no tool calls are happening.

    Runs as an asyncio task alongside the ``query()`` call.  Writes a
    heartbeat every *interval* seconds until *stop_event* is set or the
    task is cancelled.
    """
    while not (stop_event and stop_event.is_set()):
        await asyncio.sleep(interval)
        if not (stop_event and stop_event.is_set()):
            write_heartbeat(loop_id, project_id, pid, hostname)


async def run_agent_sdk(
    prompt: str,
    repo_path: str,
    log_dir: Path,
    timeout: int,
    project_id: str,
    model: str = "claude-sonnet-4-20250514",
    loop_id: str = "",
    hostname: str = "",
    session_tracker: SessionTracker | None = None,
    pre_claimed_task_id: int | None = None,
    peon_bin: str = "",
    mcp_config_dict: dict | None = None,
    worktree_name: str = "",
    max_budget: float | None = None,
    attachments_dir: Path | None = None,
    github_url: str = "",
    use_nyx_identity: bool = False,
) -> tuple[int, int | None]:
    """Run the Claude agent via the Agent SDK instead of a CLI subprocess.

    This is the ``--use-sdk`` alternative to ``run_agent()``.  It produces
    the same observable behaviour: task gets claimed, work gets done,
    session tracked, heartbeats written, worktree cleaned up.

    Returns ``(exit_code, task_id)`` — 0 on success, 1 on error/timeout.
    """
    from claude_agent_sdk import (
        ClaudeAgentOptions,
        HookMatcher,
        ResultMessage,
        query,
    )

    log_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    log_file_path = log_dir / f"agent-sdk-{timestamp}.log"

    log.info("Work work! Sending peon to the mines (SDK mode)...")
    log.info("Agent log: %s", log_file_path)

    api_url = _get_api_url()
    pid = os.getpid()
    task_id = pre_claimed_task_id

    # Create session record before starting the agent
    if session_tracker is not None:
        session_tracker.start(
            project_id,
            loop_id,
            model,
            hostname=hostname,
            task_id=pre_claimed_task_id,
        )
        detail = f"Spawning agent (SDK) for project {project_id}"
        if pre_claimed_task_id is not None:
            detail += f" (pre-claimed task #{pre_claimed_task_id})"
        session_tracker.event("start", detail=detail)

    # Inject session_id into the prompt
    final_prompt = prompt
    if session_tracker is not None and session_tracker.session_id is not None:
        final_prompt += (
            f"\n\n**Session tracking**: Your session_id is {session_tracker.session_id}. "
            f"When calling the execute_task prompt, pass session_id={session_tracker.session_id} "
            "to enable session event tracking."
        )

    # Build MCP server config
    if mcp_config_dict is None and peon_bin:
        mcp_config_dict = build_mcp_config_dict(peon_bin)

    # Build hooks
    directive_hook = _make_directive_hook(loop_id, api_url)
    heartbeat_hook = _make_heartbeat_hook(
        loop_id, project_id, pid, hostname, session_tracker
    )

    # Inject a GitHub token into the agent environment
    agent_env: dict[str, str] = {}
    from peon_mcp.github_app import (
        _parse_github_owner,
        get_github_token as _get_gh_token,
        get_nyx_github_token as _get_nyx_token,
        is_nyx_app_auth as _nyx_auth_available,
    )

    _owner = _parse_github_owner(github_url) if github_url else None
    if use_nyx_identity and _nyx_auth_available():
        _fresh_token = _get_nyx_token(owner=_owner)
        if not _fresh_token:
            _fresh_token = _get_gh_token(owner=_owner)
    else:
        _fresh_token = _get_gh_token(owner=_owner)
    if _fresh_token:
        agent_env["GH_TOKEN"] = _fresh_token

    agent_env["PEON_PROJECT_ID"] = project_id
    if session_tracker is not None and session_tracker.session_id is not None:
        agent_env["PEON_SESSION_ID"] = str(session_tracker.session_id)
    if pre_claimed_task_id is not None:
        agent_env["PEON_TASK_ID"] = str(pre_claimed_task_id)

    # Build add_dirs for attachment access
    add_dirs: list[str] = []
    if attachments_dir and attachments_dir.exists():
        add_dirs.append(str(attachments_dir))

    # Open a stderr log file so we can see Claude CLI errors
    stderr_log_path = log_dir / f"agent-sdk-stderr-{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}.log"
    stderr_file = open(stderr_log_path, "w")
    log.info("Agent SDK stderr log: %s", stderr_log_path)

    options = ClaudeAgentOptions(
        allowed_tools=["Read", "Write", "Edit", "Bash", "Glob", "Grep", "Agent"],
        model=model,
        mcp_servers=mcp_config_dict or {},
        permission_mode="bypassPermissions",
        cwd=repo_path,
        env=agent_env,
        add_dirs=add_dirs,
        max_budget_usd=max_budget,
        stderr=stderr_file,
        hooks={
            "PreToolUse": [HookMatcher(matcher=".*", hooks=[directive_hook])],
            "PostToolUse": [HookMatcher(matcher=".*", hooks=[heartbeat_hook])],
        },
    )

    # Add worktree arg via extra_args if specified
    if worktree_name:
        options.extra_args = {"worktree": worktree_name}

    # Start heartbeat watchdog
    stop_watchdog = asyncio.Event()
    watchdog_task = asyncio.create_task(
        _heartbeat_watchdog(loop_id, project_id, pid, hostname, stop_event=stop_watchdog)
    )

    # Write initial heartbeat
    write_heartbeat(loop_id, project_id, pid, hostname)

    exit_code = 0
    timed_out = False
    agent_wall_start = time.monotonic()
    result_msg: ResultMessage | None = None

    try:
        async def _run_query():
            nonlocal result_msg
            with open(log_file_path, "w") as log_file:
                async for message in query(prompt=final_prompt, options=options):
                    # Log all messages for debugging
                    log_file.write(f"{type(message).__name__}: {message}\n")
                    log_file.flush()

                    if isinstance(message, ResultMessage):
                        result_msg = message

        await asyncio.wait_for(_run_query(), timeout=timeout)

    except asyncio.TimeoutError:
        timed_out = True
        elapsed = int(time.monotonic() - agent_wall_start)
        log.warning(
            "Peon took too long! Giving up after %ds (wall-clock: %ds)",
            timeout,
            elapsed,
        )
        exit_code = 1

    except Exception as exc:
        elapsed = int(time.monotonic() - agent_wall_start)
        log.error("Agent SDK error after %ds: %s", elapsed, exc, exc_info=True)
        exit_code = 1

    finally:
        # Stop the heartbeat watchdog
        stop_watchdog.set()
        watchdog_task.cancel()
        try:
            await watchdog_task
        except asyncio.CancelledError:
            pass

        # Close stderr log and dump contents on error
        stderr_file.close()
        if exit_code != 0 and stderr_log_path.exists():
            stderr_contents = stderr_log_path.read_text().strip()
            if stderr_contents:
                log.error("Agent SDK stderr:\n%s", stderr_contents[-2000:])

    elapsed = int(time.monotonic() - agent_wall_start)

    # Check result for errors
    if result_msg is not None:
        if result_msg.is_error:
            exit_code = 1

    # --- Timeout marking ---
    if task_id is not None and timed_out:
        _mark_task_timeout(
            task_id,
            reason=f"Agent (SDK) exceeded wall-clock timeout ({timeout}s limit, ran for {elapsed}s)",
        )

    # Record session outcome
    if session_tracker is not None:
        stats: dict = {}
        if result_msg is not None:
            if result_msg.total_cost_usd is not None:
                stats["cost_usd"] = result_msg.total_cost_usd
            usage = result_msg.usage or {}
            if "input_tokens" in usage:
                stats["tokens_input"] = int(usage["input_tokens"])
            if "output_tokens" in usage:
                stats["tokens_output"] = int(usage["output_tokens"])

        if timed_out:
            session_tracker.finish(
                "timeout",
                exit_code=exit_code,
                error_summary=f"Agent (SDK) exceeded wall-clock timeout ({timeout}s limit, ran for {elapsed}s)",
                task_id=task_id,
                **stats,
            )
        elif exit_code == 0:
            session_tracker.finish("completed", exit_code=0, task_id=task_id, **stats)
        else:
            error_summary = f"Agent (SDK) exited with error after {elapsed}s"
            if result_msg is not None and result_msg.result:
                error_summary += f": {result_msg.result[:200]}"
            session_tracker.finish(
                "failed",
                exit_code=exit_code,
                error_summary=error_summary,
                task_id=task_id,
                **stats,
            )

    # Create a fallback work log if the agent skipped calling log_work
    if task_id is not None:
        _create_fallback_work_log(api_url, project_id, task_id)

    log.info(
        "Job's done! Agent (SDK) finished with code %d after %ds (timeout_limit=%ds, timed_out=%s)",
        exit_code,
        elapsed,
        timeout,
        timed_out,
    )
    return exit_code, task_id


def cleanup_worktree(repo_path: str, worktree_name: str) -> None:
    """Remove a specific git worktree by name.

    Uses 'git worktree remove --force' to clean up the worktree directory
    and its entry in .git/worktrees. Non-fatal on failure.
    """
    try:
        worktree_path = Path(repo_path) / ".claude" / "worktrees" / worktree_name
        if not worktree_path.exists():
            return
        result = subprocess.run(
            ["git", "worktree", "remove", "--force", str(worktree_path)],
            cwd=repo_path,
            capture_output=True,
            text=True,
            timeout=30,
        )
        if result.returncode == 0:
            log.info("Cleaned up worktree: %s", worktree_name)
        else:
            log.debug("Worktree removal failed (may already be gone): %s", result.stderr.strip())
    except (subprocess.TimeoutExpired, OSError) as exc:
        log.debug("Worktree cleanup error for %s: %s", worktree_name, exc)


def cleanup_stale_worktrees(repo_path: str) -> None:
    """Prune stale worktree entries and remove old peon worktree directories.

    Called at loop startup and after each agent run to prevent disk space leaks.
    """
    try:
        # First, prune any stale entries git knows about
        subprocess.run(
            ["git", "worktree", "prune"],
            cwd=repo_path,
            capture_output=True,
            text=True,
            timeout=30,
        )

        # Then scan for leftover peon-* directories in .claude/worktrees/
        worktrees_dir = Path(repo_path) / ".claude" / "worktrees"
        if not worktrees_dir.exists():
            return

        for entry in worktrees_dir.iterdir():
            if entry.is_dir() and entry.name.startswith("peon-"):
                try:
                    subprocess.run(
                        ["git", "worktree", "remove", "--force", str(entry)],
                        cwd=repo_path,
                        capture_output=True,
                        text=True,
                        timeout=30,
                    )
                    log.info("Removed stale worktree: %s", entry.name)
                except (subprocess.TimeoutExpired, OSError):
                    pass
    except (subprocess.TimeoutExpired, OSError) as exc:
        log.debug("Stale worktree cleanup error: %s", exc)


# ---------------------------------------------------------------------------
# Main loop
# ---------------------------------------------------------------------------


def main_loop(args: argparse.Namespace) -> int:
    """Core loop: check for tasks, spawn agent, repeat."""
    # Determine mode
    project_agnostic = args.project_id is None

    # Wait for the REST API to be reachable before doing anything
    ensure_api_ready()

    global _active_loop_id, _active_api_url, _heartbeat_thread
    loop_pid = os.getpid()
    loop_hostname = socket.gethostname()

    if project_agnostic:
        repo_path = None
        log_dir = Path.home() / ".peon" / "logs"
        prompt = None
        log.info("Project-agnostic mode — scanning all projects for work")
    else:
        repo_path = str(Path(args.repo_path).resolve())
        log_dir = Path(repo_path) / ".peon-logs"
        # Ensure project exists in the API
        ensure_project(args.project_id, repo_path, getattr(args, "github_url", None) or "")
        # Mark any sessions left in 'running' state by a killed process as failed
        _cleanup_stale_sessions(args.project_id)
        _recover_stuck_tasks(args.project_id, current_hostname=loop_hostname)
        # Clean up stale worktrees from previous runs
        cleanup_stale_worktrees(repo_path)
        prompt = build_prompt(args.project_id, repo_path, args.feature_id)

    # Generate unique loop ID first so we can use it for the PID file
    loop_id = str(uuid.uuid4())
    pid_file = log_dir / f"agent-{loop_id}.pid"

    # Find binaries up front
    peon_bin = find_peon_mcp_binary()
    use_sdk = getattr(args, "use_sdk", False)
    if use_sdk:
        claude_bin = ""  # Not needed for SDK mode
        log.info("Using Claude Agent SDK (--use-sdk)")
    else:
        claude_bin = find_claude_binary()
    log.info("peon-mcp binary: %s", peon_bin)
    if claude_bin:
        log.info("claude binary: %s", claude_bin)

    mcp_config = build_mcp_config(peon_bin)
    _active_loop_id = loop_id
    _active_api_url = _get_api_url()
    try:
        import importlib.metadata
        _loop_version = importlib.metadata.version("peon-mcp")
        log.info("peon-loop version: %s", _loop_version)
    except Exception:
        log.info("peon-loop version: unknown (dev mode)")
    log.info("Loop instance %s (pid=%d, host=%s)", loop_id, loop_pid, loop_hostname)

    # Start background heartbeat thread immediately in both modes.
    # In project-agnostic mode, project_id starts as None and is updated
    # dynamically when the loop picks up work from a specific project.
    _heartbeat_thread = HeartbeatThread(
        loop_id, args.project_id, loop_pid, loop_hostname, interval=30
    )
    _heartbeat_thread.start()
    log.info("Heartbeat thread started (interval=30s)")

    iteration = 0
    try:
        return _run_loop_iterations(
            args, iteration, loop_id, loop_pid, loop_hostname,
            peon_bin, claude_bin, mcp_config, prompt,
            repo_path, log_dir, pid_file
        )
    except Exception as exc:
        log.error("Main loop crashed: %s", exc, exc_info=True)
        # Mark this loop's in-progress task as timeout so it's not stuck.
        # Use host filtering to avoid accidentally timing out tasks owned
        # by a concurrent loop instance in another pod.
        if args.project_id:
            current_task = get_in_progress_task(args.project_id)
            if current_task and current_task.get("host") == loop_hostname:
                _mark_task_timeout(
                    current_task["id"],
                    reason=f"Main loop crashed with unhandled exception: {type(exc).__name__}: {exc}",
                    task_title=current_task.get("title", ""),
                )
        return 1
    finally:
        log.info("Stopping heartbeat thread...")
        _cleanup(terminate_child=False)


def _run_loop_iterations(
    args: argparse.Namespace,
    iteration: int,
    loop_id: str,
    loop_pid: int,
    loop_hostname: str,
    peon_bin: str,
    claude_bin: str,
    mcp_config: str,
    prompt: str | None,
    repo_path: str | None,
    log_dir: Path,
    pid_file: Path,
) -> int:
    """Execute the main loop iterations (extracted for cleanup handling)."""
    project_agnostic = repo_path is None
    global _heartbeat_thread

    while True:
        iteration += 1
        log.info("--- Iteration %d ---", iteration)

        # -----------------------------------------------------------------
        # Phase 1: Claim a task BEFORE spawning an agent.
        # This prevents duplicate agent spawns when multiple loops run
        # concurrently — only the loop that wins the DB-level lock
        # (FOR UPDATE SKIP LOCKED) will proceed to spawn.
        # -----------------------------------------------------------------
        claimed_task: dict | None = None

        if project_agnostic:
            # --- Project-agnostic mode: claim globally, then resolve repo ---
            log.info("Something need doing? Scanning all projects for tasks...")
            claimed_task = _claim_next_task_global(
                feature_id=args.feature_id,
                hostname=loop_hostname,
            )
            if claimed_task is None:
                if args.once:
                    log.info("No todo tasks in any project.")
                    return 0
                if _should_self_update():
                    log.info("No tasks and new version available — restarting for upgrade.")
                    return 0
                log.info("Nothing to do... peon wait. Sleeping %ds...", args.sleep)
                time.sleep(args.sleep)
                continue

            project_id = claimed_task.get("project_id", "")
            log.info(
                "Zug zug! Claimed task #%d (%s) in project '%s'",
                claimed_task.get("id", 0), claimed_task.get("title", ""), project_id,
            )
            iter_repo_path = resolve_repo_path(project_id)
            github_url = getattr(args, "github_url", None) or ""

            # Clone repo if it doesn't exist locally
            if not Path(iter_repo_path).is_dir():
                github_url = github_url or lookup_github_url(project_id)
                if not github_url:
                    log.error(
                        "Repo not found at %s and no github_url available for project '%s'. "
                        "Set github_url via the REST API or --github-url flag.",
                        iter_repo_path,
                        project_id,
                    )
                    if args.once:
                        return 1
                    time.sleep(args.sleep)
                    continue
                try:
                    clone_repo(github_url, iter_repo_path)
                except (RuntimeError, FileNotFoundError) as exc:
                    log.error("Failed to clone repo for project '%s': %s", project_id, exc)
                    if args.once:
                        return 1
                    time.sleep(args.sleep)
                    continue

            iter_github_url = github_url or ""
            ensure_project(project_id, iter_repo_path, getattr(args, "github_url", None) or "")
            iter_log_dir = Path(iter_repo_path) / ".peon-logs"
            iter_pid_file = iter_log_dir / f"agent-{loop_id}.pid"

            # Update commit_base now that we have the repo
            commit_base = _get_commit_base(iter_repo_path)
            if commit_base:
                _http_request(
                    "PUT",
                    f"{_get_api_url()}/api/tasks/{claimed_task['id']}",
                    {"commit_base": commit_base},
                )

            # Update heartbeat with current project
            _heartbeat_thread.project_id = project_id

        else:
            # --- Single-project mode: claim from this project ---
            project_id = args.project_id
            iter_repo_path = repo_path
            iter_log_dir = log_dir
            iter_pid_file = pid_file
            iter_github_url = getattr(args, "github_url", None) or ""

            log.info("Something need doing? Checking for tasks...")
            commit_base = _get_commit_base(iter_repo_path)
            claimed_task = _claim_next_task(
                project_id, commit_base,
                feature_id=args.feature_id,
                hostname=loop_hostname,
            )

            if claimed_task is None:
                # Use has_todo_tasks for informational logging only
                todo, in_progress = has_todo_tasks(project_id, args.feature_id)
                if args.once:
                    if todo == 0 and in_progress == 0:
                        log.info("Nothing to do... peon wait.")
                    elif in_progress > 0:
                        log.info("No claimable tasks (%d in_progress).", in_progress)
                    else:
                        log.info("Could not claim a task.")
                    return 0
                if _should_self_update():
                    log.info("No tasks and new version available — restarting for upgrade.")
                    return 0
                if todo == 0 and in_progress == 0:
                    log.info("Nothing to do... peon wait. Sleeping %ds...", args.sleep)
                elif in_progress > 0:
                    log.info(
                        "No claimable tasks (%d in_progress). Sleeping %ds...",
                        in_progress, args.sleep,
                    )
                else:
                    log.info("Could not claim a task. Sleeping %ds...", args.sleep)
                time.sleep(args.sleep)
                continue

            log.info(
                "Zug zug! Claimed task #%d: %s",
                claimed_task.get("id", 0), claimed_task.get("title", ""),
            )

        # -----------------------------------------------------------------
        # Phase 2: Fetch attachments, build prompt, spawn agent
        # -----------------------------------------------------------------
        claimed_task_id = claimed_task.get("id") if claimed_task else None
        attachments_dir: Path | None = None
        downloaded_attachments: list[dict] = []
        if claimed_task_id is not None:
            attachments_dir = iter_log_dir / "attachments" / str(claimed_task_id)
            downloaded_attachments = _fetch_and_download_attachments(
                claimed_task_id, attachments_dir
            )

        iter_prompt_final = build_prompt(
            project_id, iter_repo_path, args.feature_id,
            pre_claimed_task=claimed_task,
            downloaded_attachments=downloaded_attachments or None,
        )

        # Generate a unique worktree name for this iteration; Claude creates the
        # worktree at <repo>/.claude/worktrees/<name> with branch worktree-<name>.
        worktree_name = f"peon-{loop_id[:8]}-{iteration}"

        # Build the command
        cmd = [
            claude_bin,
            "--worktree", worktree_name,
            "-p", iter_prompt_final,
            "--mcp-config", mcp_config,
            "--permission-mode", args.permission_mode,
            "--model", args.model,
        ]
        if args.max_budget is not None:
            cmd.extend(["--max-budget", str(args.max_budget)])
        # Grant the agent read access to any downloaded attachments
        if attachments_dir and attachments_dir.exists() and downloaded_attachments:
            cmd.extend(["--add-dir", str(attachments_dir)])

        use_loop_detection = not getattr(args, "no_loop_detection", False)
        # Use Nyx's GitHub identity for code-review tasks so that her review
        # comments appear under Nyx's GitHub App avatar.
        is_code_review = (
            claimed_task is not None
            and claimed_task.get("task_type") == "review"
            and claimed_task.get("title", "").startswith("Code review PR #")
        )
        global _active_session_tracker
        _active_session_tracker = SessionTracker(_get_api_url())

        if getattr(args, "use_sdk", False):
            _exit_code, agent_task_id = asyncio.run(run_agent_sdk(
                prompt=iter_prompt_final,
                repo_path=iter_repo_path,
                log_dir=iter_log_dir,
                timeout=args.timeout,
                project_id=project_id,
                model=args.model,
                loop_id=loop_id,
                hostname=loop_hostname,
                session_tracker=_active_session_tracker,
                pre_claimed_task_id=claimed_task_id,
                peon_bin=peon_bin,
                worktree_name=worktree_name,
                max_budget=args.max_budget,
                attachments_dir=attachments_dir,
                github_url=iter_github_url,
                use_nyx_identity=is_code_review,
            ))
        else:
            _exit_code, agent_task_id = run_agent(
                cmd, iter_repo_path, iter_log_dir, iter_pid_file, args.timeout, project_id,
                loop_detection=use_loop_detection,
                loop_id=loop_id, model=args.model, hostname=loop_hostname,
                session_tracker=_active_session_tracker,
                github_url=iter_github_url,
                pre_claimed_task_id=claimed_task_id,
                use_nyx_identity=is_code_review,
            )
        _active_session_tracker = None

        # Post-agent verification: if the agent claimed a task but left it
        # in_progress (crashed, ran out of budget, etc.), mark it as timeout.
        # Only check the SPECIFIC task this agent claimed — never look up
        # arbitrary in_progress tasks, which could belong to a concurrent
        # agent in another loop pod.
        if agent_task_id is not None:
            post_task = get_task_by_id(agent_task_id)
            if post_task and post_task.get("status") == "in_progress":
                reason = (
                    f"Agent process exited (exit_code={_exit_code}) "
                    f"but task remained in_progress — agent did not transition task to done/cancelled before exiting"
                )
                log.warning(
                    "Task #%d (%s) still in_progress after agent exit — marking as timeout. "
                    "This typically means the agent crashed, ran out of budget, or failed to "
                    "complete its work. Check the agent log for details.",
                    agent_task_id,
                    post_task.get("title", ""),
                )
                _mark_task_timeout(
                    agent_task_id,
                    reason=reason,
                    task_title=post_task.get("title", ""),
                )

        # Clean up this iteration's worktree and downloaded attachments
        if iter_repo_path:
            cleanup_worktree(iter_repo_path, worktree_name)
        if attachments_dir and attachments_dir.exists():
            shutil.rmtree(attachments_dir, ignore_errors=True)

        if args.once:
            log.info("--once flag set, exiting after single iteration.")
            return 0

        if args.max_iterations is not None and iteration >= args.max_iterations:
            log.info("Reached max iterations (%d), peon is done!", args.max_iterations)
            return 0

        if _should_self_update():
            log.info("New version available — peon going for upgrade. Zug zug!")
            return 0  # Clean exit, K8s restarts with new image

        log.info("Zug zug, sleeping %ds before next task...", args.sleep)
        time.sleep(args.sleep)


# ---------------------------------------------------------------------------
# Signal handling
# ---------------------------------------------------------------------------

_child_proc: subprocess.Popen | None = None
_active_loop_id: str | None = None
_active_api_url: str | None = None
_heartbeat_thread: HeartbeatThread | None = None
_active_session_tracker: SessionTracker | None = None


def _cleanup(terminate_child: bool = False) -> None:
    """Perform cleanup tasks for graceful shutdown.

    Args:
        terminate_child: Whether to terminate the child process (if active).
    """
    global _heartbeat_thread, _active_session_tracker

    if _heartbeat_thread:
        _heartbeat_thread.stop()
        _heartbeat_thread = None

    if _active_loop_id:
        delete_heartbeat(_active_loop_id)

    if _active_session_tracker is not None:
        _active_session_tracker.cancel()
        _active_session_tracker = None

    if terminate_child and _child_proc and _child_proc.poll() is None:
        _child_proc.terminate()
        try:
            _child_proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            _child_proc.kill()


def _handle_sigint(_signum: int, _frame) -> None:
    """Clean up child process and heartbeat on Ctrl-C."""
    log.info("Me go now... *waves*")
    _cleanup(terminate_child=True)
    sys.exit(130)


def _handle_sigterm(signum: int, _frame) -> None:
    """Handle SIGTERM for graceful K8s pod shutdown."""
    sig_name = signal.Signals(signum).name
    log.warning("Received %s — initiating graceful shutdown...", sig_name)
    _cleanup(terminate_child=True)
    sys.exit(128 + signum)


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------


def get_git_repo_name(repo_path: str) -> str | None:
    """Extract the owner/repo identifier from git remote URL.

    Returns the identifier in 'owner/repo' format
    (e.g., 'pyrabbit/peon-mcp' from 'git@github.com:pyrabbit/peon-mcp.git')
    or None if unable to detect.
    """
    try:
        result = subprocess.run(
            ["git", "remote", "get-url", "origin"],
            cwd=repo_path,
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode != 0:
            return None

        remote_url = result.stdout.strip()
        if not remote_url:
            return None

        # Handle various git URL formats:
        # - git@github.com:user/repo.git
        # - https://github.com/user/repo.git
        # - https://github.com/user/repo

        # Remove .git suffix if present
        if remote_url.endswith(".git"):
            remote_url = remote_url[:-4]

        # Extract owner/repo
        if ":" in remote_url and remote_url.startswith("git@"):
            # SSH format: git@github.com:owner/repo
            path = remote_url.split(":")[-1]
            parts = path.split("/")
            if len(parts) >= 2:
                return f"{parts[-2]}/{parts[-1]}"
        else:
            # HTTPS format: https://github.com/owner/repo
            parts = remote_url.rstrip("/").split("/")
            if len(parts) >= 2:
                return f"{parts[-2]}/{parts[-1]}"

        return None
    except (subprocess.TimeoutExpired, subprocess.SubprocessError, OSError):
        return None


def detect_git_repo_path() -> str | None:
    """Detect if current directory is a git repository.

    Returns the absolute path to the git repository root, or None if not in a git repo.
    """
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--show-toplevel"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0:
            return result.stdout.strip()
        return None
    except (subprocess.TimeoutExpired, subprocess.SubprocessError, OSError):
        return None


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="peon-loop",
        description="Agent loop that spawns Claude to work through peon-mcp tasks.",
    )
    parser.add_argument(
        "project_id",
        nargs="?",
        default=None,
        help="Project identifier (e.g. 'peon-mcp'). Defaults to git repo name if omitted.",
    )
    parser.add_argument(
        "repo_path",
        nargs="?",
        default=None,
        help="Absolute path to the repo on disk. Defaults to current directory if omitted.",
    )
    parser.add_argument(
        "--db-path",
        default=None,
        help="Deprecated. Use DATABASE_URL env var instead. The loop itself uses the REST API.",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=1800,
        help="Max seconds per agent run (default: 1800).",
    )
    parser.add_argument(
        "--sleep",
        type=int,
        default=30,
        help="Seconds between iterations (default: 30).",
    )
    parser.add_argument(
        "--model",
        default="sonnet",
        help="Claude model to use (default: sonnet).",
    )
    parser.add_argument(
        "--max-budget",
        type=float,
        default=None,
        help="Max USD budget per agent invocation (optional).",
    )
    parser.add_argument(
        "--once",
        action="store_true",
        help="Run a single iteration then exit.",
    )
    parser.add_argument(
        "--max-iterations",
        type=int,
        default=None,
        help="Maximum number of iterations to run before exiting (default: unlimited).",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable debug logging.",
    )
    parser.add_argument(
        "--permission-mode",
        default="bypassPermissions",
        help="Claude permission mode (default: bypassPermissions).",
    )
    parser.add_argument(
        "--feature-id",
        type=int,
        default=None,
        help="Work on tasks within a specific feature (optional).",
    )
    parser.add_argument(
        "--github-url",
        default=None,
        help="GitHub repository URL for cloning repos that don't exist locally.",
    )
    parser.add_argument(
        "--no-loop-detection",
        action="store_true",
        default=False,
        help="Disable loop detection (useful for debugging).",
    )
    parser.add_argument(
        "--use-sdk",
        action="store_true",
        default=False,
        help="Use Claude Agent SDK instead of CLI subprocess.",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    # Handle install-hooks subcommand before argparse
    effective_argv = argv if argv is not None else sys.argv[1:]
    if effective_argv and effective_argv[0] == "install-hooks":
        from peon_mcp.hooks.install import main as install_main
        install_main()
        return

    args = parse_args(argv)

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    # Auto-detect repo_path if not provided
    if args.repo_path is None:
        detected_path = detect_git_repo_path()
        if detected_path is not None:
            args.repo_path = detected_path
            log.info("Auto-detected repo path: %s", args.repo_path)

    # Auto-detect project_id if not provided but repo_path is available
    if args.project_id is None and args.repo_path is not None:
        detected_name = get_git_repo_name(args.repo_path)
        if detected_name is not None:
            args.project_id = detected_name
            log.info("Auto-detected project ID: %s", args.project_id)

    # If still no project_id, enter project-agnostic mode
    if args.project_id is None:
        log.info("No project specified — running in project-agnostic worker mode.")
        log.info("Will pick up tasks from any project in the database.")
    elif args.repo_path is None:
        # Have project_id but no repo_path — use default location
        args.repo_path = resolve_repo_path(args.project_id)
        log.info("Using default repo path: %s", args.repo_path)

    signal.signal(signal.SIGINT, _handle_sigint)
    signal.signal(signal.SIGTERM, _handle_sigterm)

    try:
        code = main_loop(args)
    except FileNotFoundError as exc:
        log.error("%s", exc)
        sys.exit(1)
    except KeyboardInterrupt:
        log.info("Me go now... *waves*")
        _cleanup(terminate_child=False)
        sys.exit(130)

    sys.exit(code)


if __name__ == "__main__":
    main()
