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

Checks the SQLite DB for todo tasks before spawning Claude so we don't
waste money on empty runs.

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 json
import logging
import os
import re
import signal
import socket
import sqlite3
import subprocess
import sys
import threading
import time
import urllib.error
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) -> dict | None:
    """Make a synchronous HTTP request, returning parsed JSON or None on failure."""
    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=5) as resp:
            return json.loads(resp.read())
    except (urllib.error.URLError, json.JSONDecodeError, OSError) as exc:
        log.warning("Session API call failed (%s %s): %s", method, url, exc)
        return None


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) -> None:
        """Create the session record before the agent process starts."""
        url = f"{self._api_url}/api/sessions"
        result = _http_request(
            "POST",
            url,
            {"project_id": project_id, "loop_id": loop_id, "model": model},
        )
        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 link_task(self, task_id: int) -> None:
        """Associate this session with the task the agent picked up."""
        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})

    def event(self, event_type: str, tool_name: str = "", detail: str = "") -> None:
        """Append an event to the session."""
        if self.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 = "",
    ) -> None:
        """Update the session when the agent process exits."""
        if self.session_id is None:
            return
        payload: dict = {
            "status": status,
            "ended_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        }
        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)", self.session_id, status)

    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)


def _poll_discord_messages(project_id: str, api_url: str, api_key: str = "") -> None:
    """Fetch and log any pending inbound Discord messages between task cycles.

    Messages are acknowledged as 'processed' after logging so the next cycle
    sees a clean queue.  Future iterations may route non-trivial messages as
    agent context or trigger interrupt handling.

    Args:
        project_id: The project whose Discord configs to query.
        api_url: Base URL for the peon REST API.
        api_key: Optional API key for authenticated instances.
    """
    try:
        from peon_mcp.discord.loop_bridge import ack_message, handle_interrupt, poll_messages
    except ImportError:
        return

    messages = poll_messages(project_id, api_url, api_key=api_key)
    if not messages:
        return

    log.info("Discord: %d pending message(s) received.", len(messages))
    for msg in messages:
        author = msg.get("author_name", "unknown")
        channel = msg.get("channel_name") or msg.get("channel_id", "?")
        content = msg.get("content", "")
        is_interrupt = handle_interrupt(msg)

        if is_interrupt:
            log.warning(
                "Discord interrupt message from %s in #%s: %s",
                author, channel, content[:200],
            )
        else:
            log.info(
                "Discord message from %s in #%s: %s",
                author, channel, content[:200],
            )

        ack_message(msg["id"], api_url, api_key=api_key)


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


def ensure_project(db_path: str, project_id: str, repo_path: str, github_url: str = "") -> None:
    """Create the DB, tables, and project row if they don't already exist.

    Also updates the path column to match the current machine's filesystem
    and backfills github_url if the existing value is empty.
    """
    from peon_mcp.db import SCHEMA

    conn = sqlite3.connect(db_path)
    try:
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA foreign_keys=ON")
        conn.executescript(SCHEMA)
        conn.execute(
            "INSERT OR IGNORE INTO projects (id, path, github_url) VALUES (?, ?, ?)",
            (project_id, repo_path, github_url),
        )
        # Always update path for this machine
        conn.execute(
            "UPDATE projects SET path = ? WHERE id = ?",
            (repo_path, project_id),
        )
        # Backfill github_url if not already set
        if github_url:
            conn.execute(
                "UPDATE projects SET github_url = ? WHERE id = ? AND github_url = ''",
                (github_url, project_id),
            )
        conn.commit()
    finally:
        conn.close()
    log.info("Ready to work! Peon loop starting for project '%s'", project_id)


def ensure_db_schema(db_path: str) -> None:
    """Ensure the DB file and schema tables exist (without creating a project).

    Used in project-agnostic mode where we need the DB ready to receive
    tasks but don't know which project we'll work on yet.
    """
    from peon_mcp.db import SCHEMA

    Path(db_path).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    try:
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA foreign_keys=ON")
        conn.executescript(SCHEMA)
        conn.commit()
    finally:
        conn.close()


def find_next_task_any_project(db_path: str) -> str | None:
    """Find the project containing the highest-priority todo task across all projects.

    Returns the project_id or None if no todo tasks exist in any project.
    Uses the same priority ordering as next_task(): critical > high > medium > low,
    oldest first within the same priority.
    """
    try:
        conn = sqlite3.connect(db_path)
        try:
            cursor = conn.execute(
                "SELECT project_id FROM tasks"
                " WHERE status = 'todo'"
                " AND (scheduled_at IS NULL OR scheduled_at <= datetime('now'))"
                " ORDER BY CASE priority"
                "   WHEN 'critical' THEN 0"
                "   WHEN 'high' THEN 1"
                "   WHEN 'medium' THEN 2"
                "   WHEN 'low' THEN 3"
                " END, created_at ASC"
                " LIMIT 1"
            )
            row = cursor.fetchone()
            return row[0] if row else None
        finally:
            conn.close()
    except sqlite3.OperationalError as exc:
        log.warning("DB query failed (table may not exist): %s", exc)
        return None


def lookup_github_url(db_path: str, project_id: str) -> str:
    """Look up the github_url for a project from the database.

    Returns the URL string or empty string if not found.
    """
    try:
        conn = sqlite3.connect(db_path)
        try:
            cursor = conn.execute(
                "SELECT github_url FROM projects WHERE id = ?", (project_id,)
            )
            row = cursor.fetchone()
            return row[0] if row else ""
        finally:
            conn.close()
    except sqlite3.OperationalError:
        return ""


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


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

    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)
    log.info("Cloning %s → %s", github_url, target_path)
    result = subprocess.run(
        ["git", "clone", github_url, target_path],
        capture_output=True,
        text=True,
        timeout=300,  # 5-minute timeout for large repos
    )
    if result.returncode != 0:
        raise RuntimeError(
            f"git clone failed (exit {result.returncode}): {result.stderr.strip()}"
        )
    log.info("Clone complete.")


def run_subagent_maintenance(db_path: str) -> None:
    """Run sub-agent timeout and stale-cleanup checks on each loop cycle.

    Opens its own sqlite3 connection (consistent with loop.py patterns),
    calls :func:`~peon_mcp.subagents.manager.monitor_timeouts` to kill and mark
    processes that exceeded their timeout, then calls
    :func:`~peon_mcp.subagents.manager.cleanup_stale` to handle sub-agents that
    got stuck in ``pending`` or ``running`` without a watchdog.

    Errors are logged but never propagate — maintenance failures must not
    disrupt the normal agent dispatch cycle.
    """
    try:
        from peon_mcp.subagents.manager import cleanup_stale, monitor_timeouts

        conn = sqlite3.connect(db_path)
        try:
            conn.execute("PRAGMA foreign_keys=ON")
            timed_out = monitor_timeouts(conn)
            cleaned = cleanup_stale(conn)
            if timed_out or cleaned:
                log.info(
                    "Sub-agent maintenance: %d timed out, %d stale cleaned up",
                    timed_out,
                    cleaned,
                )
        finally:
            conn.close()
    except Exception as exc:
        log.warning("Sub-agent maintenance failed (non-fatal): %s", exc)


def _cleanup_stale_sessions(db_path: str, 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.
    """
    try:
        conn = sqlite3.connect(db_path)
        try:
            conn.execute(
                "UPDATE agent_sessions"
                " SET status = 'failed',"
                "     error_summary = 'Marked failed on loop startup (prior process was killed)',"
                "     updated_at = CURRENT_TIMESTAMP"
                " WHERE project_id = ? AND status = 'running'",
                (project_id,),
            )
            count = conn.execute("SELECT changes()").fetchone()[0]
            conn.commit()
            if count:
                log.info("Marked %d stale running session(s) as failed on startup", count)
        finally:
            conn.close()
    except sqlite3.Error as exc:
        log.warning("Failed to cleanup stale sessions: %s", exc)


def cleanup_stale_worktrees(db_path: str, project_id: str, repo_path: str) -> None:
    """Clean up worktrees for tasks that are done, cancelled, or timed out.

    This is a startup cleanup to remove any worktrees that were left behind
    from previous runs.
    """
    try:
        conn = sqlite3.connect(db_path)
        try:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(
                "SELECT id, worktree_path, status FROM tasks"
                " WHERE project_id = ? AND worktree_path != ''"
                " AND status IN ('done', 'cancelled', 'timeout')",
                (project_id,),
            )
            tasks = [dict(row) for row in cursor.fetchall()]
        finally:
            conn.close()

        if not tasks:
            log.debug("No stale worktrees to clean up")
            return

        main_repo_path = Path(repo_path).resolve()
        log.info("Found %d task(s) with worktrees to clean up", len(tasks))
        for task in tasks:
            worktree_path = Path(task["worktree_path"]).expanduser().resolve()
            if worktree_path == main_repo_path:
                log.debug(
                    "Skipping worktree cleanup for task #%d — path is the main repo: %s",
                    task["id"],
                    worktree_path,
                )
                continue
            if not worktree_path.exists():
                log.debug("Worktree %s already removed", worktree_path)
                continue

            try:
                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 for task #%d at %s", task["id"], worktree_path)
                else:
                    log.warning(
                        "Failed to remove worktree for task #%d at %s: %s",
                        task["id"],
                        worktree_path,
                        result.stderr.strip(),
                    )
            except Exception as exc:
                log.error(
                    "Error removing worktree for task #%d at %s: %s",
                    task["id"],
                    worktree_path,
                    exc,
                )
    except sqlite3.Error as exc:
        log.warning("Failed to query for stale worktrees: %s", exc)


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

    Uses stdlib sqlite3 (not async aiosqlite) for a lightweight pre-flight
    check before spawning the full agent.

    Args:
        db_path: Path to SQLite database.
        project_id: Project identifier.
        feature_id: Optional feature ID to filter tasks.

    Returns:
        (todo_count, in_progress_count).  Returns (0, 0) if the DB or table
        doesn't exist.
    """
    try:
        conn = sqlite3.connect(db_path)
        try:
            query = (
                "SELECT status, COUNT(*) FROM tasks"
                " WHERE project_id = ?"
                " AND (status = 'in_progress'"
                "      OR (status = 'todo'"
                "          AND (scheduled_at IS NULL OR scheduled_at <= datetime('now'))))"
            )
            params = [project_id]

            if feature_id is not None:
                query += " AND feature_id = ?"
                params.append(feature_id)

            query += " GROUP BY status"

            cursor = conn.execute(query, params)
            counts = dict(cursor.fetchall())
        finally:
            conn.close()
    except sqlite3.OperationalError as exc:
        # Table doesn't exist yet, etc.
        log.warning("DB query failed (table may not exist): %s", exc)
        return (0, 0)

    todo = counts.get("todo", 0)
    in_progress = counts.get("in_progress", 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(db_path: str, 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.
    """
    try:
        conn = sqlite3.connect(db_path)
        try:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(
                "SELECT * FROM tasks"
                " WHERE project_id = ? AND status = 'in_progress'"
                " ORDER BY updated_at DESC"
                " LIMIT 1",
                (project_id,),
            )
            row = cursor.fetchone()
            if row:
                return dict(row)
            return None
        finally:
            conn.close()
    except sqlite3.OperationalError as exc:
        log.warning("DB query failed (table may not exist): %s", exc)
        return None


def evaluate_schedules_sync(db_path: str) -> list[dict]:
    """Synchronously evaluate due schedules and create any pending tasks.

    Wraps the async evaluate_schedules() from the scheduling module so the
    synchronous peon-loop can call it without an existing event loop.

    Args:
        db_path: Path to the SQLite database file.

    Returns:
        List of created task dicts, or an empty list on error.
    """
    import asyncio

    import aiosqlite

    from peon_mcp.scheduling.evaluator import evaluate_schedules

    async def _run() -> list[dict]:
        async with aiosqlite.connect(db_path) as db:
            db.row_factory = aiosqlite.Row
            await db.execute("PRAGMA journal_mode=WAL")
            await db.execute("PRAGMA foreign_keys=ON")
            return await evaluate_schedules(db)

    try:
        return asyncio.run(_run())
    except Exception as exc:
        log.warning("Schedule evaluation failed: %s", exc)
        return []


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

    Args:
        db_path: Path to SQLite database.
        task_id: Task ID to fetch.

    Returns:
        Task dict with id, title, etc., or None if task doesn't exist.
    """
    try:
        conn = sqlite3.connect(db_path)
        try:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(
                "SELECT * FROM tasks WHERE id = ?",
                (task_id,),
            )
            row = cursor.fetchone()
            if row:
                return dict(row)
            return None
        finally:
            conn.close()
    except sqlite3.OperationalError as exc:
        log.warning("DB query failed (table may not exist): %s", exc)
        return None


def _get_task_started_after(db_path: str, 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:
        db_path: Path to SQLite database.
        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.
    """
    try:
        conn = sqlite3.connect(db_path)
        try:
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(
                "SELECT * FROM tasks"
                " WHERE project_id = ? AND started_at >= ?"
                " ORDER BY started_at DESC"
                " LIMIT 1",
                (project_id, after_timestamp),
            )
            row = cursor.fetchone()
            return dict(row) if row else None
        finally:
            conn.close()
    except sqlite3.OperationalError as exc:
        log.warning("DB query failed (table may not exist): %s", exc)
        return None


# ---------------------------------------------------------------------------
# Tool policy
# ---------------------------------------------------------------------------

# Map peon-mcp policy tool names to Claude CLI tool names.
_POLICY_TOOL_TO_CLAUDE: dict[str, str] = {
    "read": "Read",
    "write": "Write",
    "edit": "Edit",
    "glob": "Glob",
    "grep": "Grep",
    "bash": "Bash",
    "git": "Bash",
    "task": "Task",
    "webfetch": "WebFetch",
    "websearch": "WebSearch",
    "notebookedit": "NotebookEdit",
}


def get_policy_for_project(
    db_path: str,
    project_id: str,
    task_id: int | None = None,
    policy_name_override: str | None = None,
) -> dict | None:
    """Get the effective tool policy for a project/task.

    Resolution order:
    1. Explicit policy name override (--policy flag)
    2. Task-level tool_policy_id (if task_id is given)
    3. Project default policy (is_default = 1)

    Returns a raw policy dict or None if no policy applies.
    """
    try:
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        try:
            if policy_name_override:
                cursor = conn.execute(
                    "SELECT * FROM tool_policies WHERE project_id = ? AND name = ?",
                    (project_id, policy_name_override),
                )
                row = cursor.fetchone()
                if row:
                    return dict(row)
                log.warning(
                    "Policy '%s' not found for project '%s'", policy_name_override, project_id
                )
                return None

            if task_id is not None:
                cursor = conn.execute(
                    "SELECT tp.* FROM tool_policies tp"
                    " JOIN tasks t ON t.tool_policy_id = tp.id"
                    " WHERE t.id = ?",
                    (task_id,),
                )
                row = cursor.fetchone()
                if row:
                    return dict(row)

            cursor = conn.execute(
                "SELECT * FROM tool_policies WHERE project_id = ? AND is_default = 1",
                (project_id,),
            )
            row = cursor.fetchone()
            if row:
                return dict(row)

            return None
        finally:
            conn.close()
    except sqlite3.OperationalError as exc:
        log.debug("Policy query failed (table may not exist): %s", exc)
        return None


def build_policy_flags(policy_row: dict | None) -> list[str]:
    """Convert a raw policy row to Claude CLI --allowedTools / --disallowedTools flags.

    Returns an empty list when no policy applies or when the policy is unrestricted.
    """
    if policy_row is None:
        return []

    try:
        from peon_mcp.policies.resolver import resolve_policy

        resolved = resolve_policy(policy_row)
    except ImportError:
        log.debug("Policy resolver not available; skipping policy flags")
        return []

    # exec_security='deny' disables Bash entirely regardless of allowed_tools
    if resolved.exec_security == "deny":
        return ["--disallowedTools", "Bash"]

    # Unrestricted profile with no explicit overrides → no flags needed
    if resolved.is_unrestricted:
        return []

    if not resolved.allowed_tools:
        return []

    # Map policy tool names to Claude CLI tool names, deduplicating
    claude_tools: list[str] = []
    seen: set[str] = set()
    for tool in resolved.allowed_tools:
        claude_name = _POLICY_TOOL_TO_CLAUDE.get(tool.lower(), tool.capitalize())
        if claude_name not in seen:
            claude_tools.append(claude_name)
            seen.add(claude_name)

    return ["--allowedTools", ",".join(claude_tools)]


def build_policy_context(policy_row: dict | None, worktree_path: str | None = None) -> str:
    """Build a human-readable policy constraints string to inject into the agent prompt.

    Returns an empty string when no policy applies.
    """
    if policy_row is None:
        return ""

    try:
        from peon_mcp.policies.resolver import resolve_policy

        resolved = resolve_policy(policy_row)
    except ImportError:
        return ""

    lines = [f"## Tool Policy: {resolved.policy_name}", ""]

    if resolved.is_unrestricted:
        lines.append("You have unrestricted tool access (no policy restrictions apply).")
    else:
        if resolved.allowed_tools:
            lines.append(
                f"You are restricted to the following tools: {', '.join(resolved.allowed_tools)}"
            )
        if resolved.denied_tools:
            lines.append(f"The following tools are explicitly blocked: {', '.join(resolved.denied_tools)}")

    if resolved.fs_workspace_only:
        if worktree_path:
            lines.append(f"Filesystem operations are limited to: {worktree_path}")
        else:
            lines.append("Filesystem operations are limited to the task worktree directory.")

    if resolved.exec_security == "deny":
        lines.append("Shell command execution (Bash) is disabled for this session.")
    elif resolved.exec_security == "allowlist":
        lines.append("Shell commands require approval for any commands not on the standard allowlist.")

    return "\n".join(lines)


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


def write_heartbeat(
    db_path: str,
    loop_id: str,
    project_id: str,
    pid: int,
    hostname: str,
    discord_status: str = "",
) -> None:
    """Upsert a heartbeat row so the UI can show this loop as active."""
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "INSERT OR REPLACE INTO loop_heartbeats"
            " (id, project_id, pid, hostname, discord_status, started_at, last_heartbeat)"
            " VALUES (?, ?, ?, ?, ?, COALESCE("
            "   (SELECT started_at FROM loop_heartbeats WHERE id = ?), CURRENT_TIMESTAMP"
            " ), CURRENT_TIMESTAMP)",
            (loop_id, project_id, pid, hostname, discord_status, loop_id),
        )
        conn.commit()
    finally:
        conn.close()


def delete_heartbeat(db_path: str, loop_id: str) -> None:
    """Remove the heartbeat row on clean shutdown."""
    try:
        conn = sqlite3.connect(db_path)
        conn.execute("DELETE FROM loop_heartbeats WHERE id = ?", (loop_id,))
        conn.commit()
        conn.close()
    except Exception:
        pass  # best-effort cleanup


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,
        db_path: str,
        loop_id: str,
        project_id: str,
        pid: int,
        hostname: str,
        interval: int = 30,
    ):
        super().__init__(daemon=True, name="heartbeat-thread")
        self.db_path = db_path
        self.loop_id = loop_id
        self.project_id = project_id
        self.pid = pid
        self.hostname = hostname
        self.interval = interval
        self.discord_status: str = ""
        self._stop_event = threading.Event()

    def run(self) -> None:
        """Write heartbeats every interval seconds until stopped."""
        # Write immediate heartbeat on startup
        try:
            write_heartbeat(
                self.db_path,
                self.loop_id,
                self.project_id,
                self.pid,
                self.hostname,
                discord_status=self.discord_status,
            )
        except Exception as exc:
            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.db_path,
                    self.loop_id,
                    self.project_id,
                    self.pid,
                    self.hostname,
                    discord_status=self.discord_status,
                )
            except Exception as exc:
                log.warning("Heartbeat thread failed to write: %s", exc)

    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 load_pruning_config(db_path: str, project_id: str) -> dict:
    """Load the context pruning config for a project from the DB.

    Returns a dict with at minimum:
      - enabled (bool)
      - default_max_chars (int)
      - ttl_minutes (int)

    Falls back to sensible defaults if the project has no config set.
    """
    defaults = {"enabled": True, "default_max_chars": 10000, "ttl_minutes": 30}
    try:
        conn = sqlite3.connect(db_path)
        try:
            row = conn.execute(
                "SELECT context_pruning FROM projects WHERE id = ?", (project_id,)
            ).fetchone()
        finally:
            conn.close()
    except sqlite3.Error:
        return defaults

    if row is None or not row[0]:
        return defaults

    try:
        stored = json.loads(row[0])
        return {**defaults, **stored}
    except (json.JSONDecodeError, TypeError):
        return defaults


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, db_path: str) -> str:
    """Build inline JSON for --mcp-config.

    Passes the resolved DB path as an env var so the MCP server uses the
    same database file as the pre-flight check.
    """
    config = {
        "mcpServers": {
            "peon-mcp": {
                "command": peon_bin,
                "env": {"PEON_DB_PATH": db_path},
            }
        }
    }
    return json.dumps(config)


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

    When ``task`` is provided, appends task-type-specific instructions for
    ``fix`` and ``verify`` task types from the review→fix→verify pipeline.
    """
    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"
        "\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}) 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"
        "\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"
        f"Note: All actual development work will be done in a git worktree, NOT in the main repo at {repo_path}. "
        "The execute_task prompt will provide detailed worktree instructions.\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 from a code review pipeline. Apply specific findings from reviewers, "
        "keep scope tight, and create a PR with focused commits.\n"
        "- `verify`: A VERIFY task — the final gate in a review→fix→verify pipeline. Run the full "
        "test suite and only mark done if everything passes.\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 from a code review pipeline. Your job is to:\n"
                "1. Read the task description carefully — it contains specific findings from code reviewers\n"
                "2. Checkout the PR branch referenced in the review session\n"
                "3. Apply ONLY the fixes described in the findings — keep scope tight\n"
                "4. Each fix should be a focused commit: `fix: <category> in <file_path>`\n"
                "5. Do NOT expand scope beyond what the findings describe\n"
                "6. After applying fixes, create a PR (or push to existing branch)\n"
                "7. Run tests to verify your fix doesn't break anything\n"
                "8. Mark the task done with the PR URL\n"
                "\n"
                "Remember: You are the repair phase of a review\u2192fix\u2192verify pipeline.\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 stage of a review\u2192fix\u2192verify pipeline.\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 via run_tests()\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."
            )
    if policy_context:
        prompt += "\n\n" + policy_context
    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 cleanup_worktree(task: dict | None, repo_path: str) -> None:
    """Remove the worktree for a task if it exists.

    Only cleans up worktrees for tasks that have a PR URL or are in terminal
    states (cancelled, timeout). This prevents premature cleanup of work that
    hasn't been pushed to GitHub yet.

    Args:
        task: Task dict with worktree_path field
        repo_path: Path to the main repository
    """
    if not task or not task.get("worktree_path"):
        return

    # Skip cleanup if task is 'done' but has no PR URL - work not pushed yet
    if task.get("status") == "done" and not task.get("pr_url"):
        log.warning(
            "Skipping worktree cleanup for task #%d - marked done but no PR created. "
            "Worktree preserved at %s for manual push/PR creation.",
            task["id"],
            task["worktree_path"],
        )
        return

    worktree_path = Path(task["worktree_path"]).expanduser().resolve()

    # Never attempt to remove the main repo as a worktree
    if worktree_path == Path(repo_path).resolve():
        log.debug(
            "Skipping worktree cleanup for task #%d — path is the main repo: %s",
            task["id"],
            worktree_path,
        )
        return

    # Check if the worktree path exists
    if not worktree_path.exists():
        log.debug("Worktree path %s does not exist, skipping cleanup", worktree_path)
        return

    try:
        log.info("Cleaning up worktree at %s", worktree_path)
        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("Successfully removed worktree for task #%d", task["id"])
        else:
            log.warning(
                "Failed to remove worktree at %s: %s",
                worktree_path,
                result.stderr.strip(),
            )
    except subprocess.TimeoutExpired:
        log.error("Timeout while removing worktree at %s", worktree_path)
    except Exception as exc:
        log.error("Error removing worktree at %s: %s", worktree_path, exc)


def _mark_current_task_timeout(db_path: str, project_id: str) -> None:
    """Mark the currently in_progress task for this project as 'timeout'."""
    current_task = get_in_progress_task(db_path, project_id)
    if current_task:
        try:
            conn = sqlite3.connect(db_path)
            try:
                conn.execute(
                    "UPDATE tasks SET status = 'timeout', updated_at = CURRENT_TIMESTAMP"
                    " WHERE id = ? AND status = 'in_progress'",
                    (current_task["id"],),
                )
                conn.commit()
                log.info("Task #%d status updated to 'timeout'", current_task["id"])
            finally:
                conn.close()
        except sqlite3.Error as exc:
            log.error("Failed to update task status to timeout: %s", exc)


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.policies.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,
    db_path: str,
    project_id: str,
    extra_env: dict[str, str] | None = None,
    loop_detection: bool = True,
    loop_id: str = "",
    model: str = "",
    session_tracker: SessionTracker | None = None,
) -> int:
    """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 the exit code.

    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.
        db_path: SQLite database path, passed to the circuit-breaker logger.
        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.
    """
    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
    if session_tracker is not None:
        session_tracker.start(project_id, loop_id, model)
        session_tracker.event("start", detail=f"Spawning agent for project {project_id}")

    # 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 {})}

    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))

        # Give the agent a moment to start and claim a task
        time.sleep(2)

        # Check which task the agent picked up and save the task ID for later cleanup
        task = get_in_progress_task(db_path, project_id)
        task_id = None
        if task:
            task_id = task["id"]
            log.info("Peon is working on task #%d: %s", task_id, task["title"])
            if session_tracker is not None:
                session_tracker.link_task(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.policies.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

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

            _mark_current_task_timeout(db_path, project_id)

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

    # If the circuit breaker fired, mark the still-in-progress task as timed out
    if circuit_breaker_fired.is_set():
        _mark_current_task_timeout(db_path, project_id)

    # Safety net: unconditionally re-fetch the task the agent worked on.
    # The initial 2-second poll often misses the assignment (agents take 20-30s
    # to call next_task).  After proc.wait() we try two strategies:
    #  1. Still in_progress (agent crashed without marking done).
    #  2. Any task started after agent_start_ts (agent completed normally).
    if task_id is None:
        task = get_in_progress_task(db_path, project_id)
        if not task:
            task = _get_task_started_after(db_path, project_id, agent_start_ts)
        if task:
            task_id = task["id"]
            log.info(
                "Late task link: peon was working on task #%d: %s",
                task_id,
                task.get("title", ""),
            )
    if session_tracker is not None and task_id is not None:
        session_tracker.link_task(task_id)

    # 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 timed out after {timeout}s",
                **log_stats,
            )
        elif proc.returncode == 0:
            session_tracker.finish("completed", exit_code=0, **log_stats)
        else:
            session_tracker.finish(
                "failed",
                exit_code=proc.returncode,
                error_summary=f"Agent exited with code {proc.returncode}",
                **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)

    # Re-fetch task data before cleanup to ensure we have current status/worktree_path/pr_url
    # Use task_id captured above instead of querying by status, since the agent may have
    # already marked the task as done/cancelled/timeout
    if task_id is not None:
        task = get_task_by_id(db_path, task_id)
        cleanup_worktree(task, repo_path)
    else:
        log.warning("No task ID captured, skipping worktree cleanup")

    _child_proc = None
    pid_file.unlink(missing_ok=True)
    log.info("Job's done! Agent exited with code %d", proc.returncode)
    return proc.returncode


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


def main_loop(args: argparse.Namespace) -> int:
    """Core loop: check for tasks, spawn agent, repeat."""
    from peon_mcp.db import get_db_path

    if args.db_path is not None:
        db_path = str(Path(args.db_path).resolve())
    else:
        db_path = get_db_path()

    # Determine mode
    project_agnostic = args.project_id is None

    if project_agnostic:
        repo_path = None
        log_dir = Path.home() / ".peon" / "logs"
        prompt = None
        # Ensure DB schema exists so tasks can be added
        ensure_db_schema(db_path)
        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 DB and project exist so tasks can be added immediately
        ensure_project(db_path, args.project_id, repo_path, getattr(args, "github_url", None) or "")
        # Clean up any stale worktrees from previous runs
        cleanup_stale_worktrees(db_path, args.project_id, repo_path)
        # Mark any sessions left in 'running' state by a killed process as failed
        _cleanup_stale_sessions(db_path, args.project_id)
        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()
    claude_bin = find_claude_binary()
    log.info("peon-mcp binary: %s", peon_bin)
    log.info("claude binary: %s", claude_bin)

    mcp_config = build_mcp_config(peon_bin, db_path)

    global _active_loop_id, _active_db_path, _heartbeat_thread
    loop_pid = os.getpid()
    loop_hostname = socket.gethostname()
    _active_loop_id = loop_id
    _active_db_path = db_path
    log.info("Loop instance %s (pid=%d, host=%s)", loop_id, loop_pid, loop_hostname)

    if not project_agnostic:
        # Start background heartbeat thread immediately for single-project mode
        _heartbeat_thread = HeartbeatThread(
            db_path, loop_id, args.project_id, loop_pid, loop_hostname, interval=30
        )
        _heartbeat_thread.start()
        log.info("Heartbeat thread started (interval=30s)")

        # Start Discord bot thread if --discord flag is set
        if getattr(args, "discord", False):
            global _discord_loop_thread
            try:
                from peon_mcp.discord.loop_bridge import DiscordLoopThread
                _discord_loop_thread = DiscordLoopThread(
                    db_path=db_path, project_id=args.project_id
                )
                _discord_loop_thread.start()
                log.info("Discord loop thread started for project '%s'", args.project_id)
            except ImportError as exc:
                log.warning("Discord integration unavailable (missing deps): %s", exc)

    # In project-agnostic mode, heartbeat starts lazily in _run_loop_iterations

    iteration = 0
    try:
        return _run_loop_iterations(
            args, iteration, loop_id, loop_pid, loop_hostname,
            db_path, peon_bin, claude_bin, mcp_config, prompt,
            repo_path, log_dir, pid_file
        )
    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,
    db_path: 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)

        # Run sub-agent maintenance on every cycle (non-fatal if it fails)
        run_subagent_maintenance(db_path)

        if project_agnostic:
            # --- Project-agnostic mode: discover work dynamically ---
            log.info("Something need doing? Scanning all projects for tasks...")
            project_id = find_next_task_any_project(db_path)
            if project_id is None:
                if args.once:
                    log.info("No todo tasks in any project.")
                    return 0
                log.info("Nothing to do... peon wait. Sleeping %ds...", args.sleep)
                time.sleep(args.sleep)
                continue

            log.info("Zug zug! Found work in project '%s'", project_id)
            iter_repo_path = resolve_repo_path(project_id)

            # Clone repo if it doesn't exist locally
            if not Path(iter_repo_path).is_dir():
                github_url = getattr(args, "github_url", None) or lookup_github_url(db_path, 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 update_project tool 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

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

            # Start or update heartbeat for current project
            if _heartbeat_thread is None:
                _heartbeat_thread = HeartbeatThread(
                    db_path, loop_id, project_id, loop_pid, loop_hostname, interval=30
                )
                _heartbeat_thread.start()
                log.info("Heartbeat thread started (interval=30s)")
            else:
                _heartbeat_thread.project_id = project_id

        else:
            # --- Single-project mode ---
            project_id = args.project_id
            iter_repo_path = repo_path
            iter_prompt = prompt
            iter_log_dir = log_dir
            iter_pid_file = pid_file

            log.info("Something need doing? Checking for tasks...")
            todo, in_progress = has_todo_tasks(db_path, project_id, args.feature_id)
            log.info("Zug zug! Found %d task(s), %d in progress", todo, in_progress)

            if todo == 0 and in_progress == 0:
                if args.once:
                    log.info("Nothing to do... peon wait. Sleeping %ds...", 0)
                    return 0
                log.info("Nothing to do... peon wait. Sleeping %ds...", args.sleep)
                time.sleep(args.sleep)
                continue

            if todo == 0 and in_progress > 0:
                log.warning(
                    "No todo tasks but %d task(s) stuck in_progress. "
                    "Reset them via peon-ui or update_task.",
                    in_progress,
                )
                if args.once:
                    return 0
                log.info("Sleeping %ds...", args.sleep)
                time.sleep(args.sleep)
                continue

        # Evaluate schedules before spawning the agent so any due templates
        # become visible tasks that the agent can immediately pick up.
        created = evaluate_schedules_sync(db_path)
        for item in created:
            log.info(
                "Schedule '%s' created task: %s",
                item.get("template_id", "?"),
                item.get("title", "?"),
            )

        # Resolve tool policy for this project/iteration and build policy flags + context
        policy_name_override = getattr(args, "policy", None)
        policy_row = get_policy_for_project(db_path, project_id, policy_name_override=policy_name_override)
        policy_flags = build_policy_flags(policy_row)
        policy_context = build_policy_context(policy_row)

        # Rebuild the prompt with policy context for this iteration
        iter_prompt_with_policy = build_prompt(project_id, iter_repo_path, args.feature_id, policy_context)

        # Build the command
        cmd = [
            claude_bin,
            "-p", iter_prompt_with_policy,
            "--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)])
        if policy_flags:
            cmd.extend(policy_flags)

        # Load context pruning config and pass limits as env vars so the MCP
        # server and execute_task prompt can surface them to the agent.
        pruning = load_pruning_config(db_path, project_id)
        pruning_env: dict[str, str] = {}
        if pruning.get("enabled", True):
            pruning_env["PEON_MAX_TOOL_RESULT_CHARS"] = str(pruning["default_max_chars"])
            pruning_env["PEON_TOOL_RESULT_TTL_MINUTES"] = str(pruning["ttl_minutes"])
            log.debug(
                "Context pruning enabled: max_chars=%s ttl_minutes=%s",
                pruning["default_max_chars"],
                pruning["ttl_minutes"],
            )

        use_loop_detection = not getattr(args, "no_loop_detection", False)
        global _active_session_tracker
        _active_session_tracker = SessionTracker(_get_api_url())
        run_agent(
            cmd, iter_repo_path, iter_log_dir, iter_pid_file, args.timeout, db_path, project_id,
            loop_detection=use_loop_detection, extra_env=pruning_env,
            loop_id=loop_id, model=args.model, session_tracker=_active_session_tracker,
        )
        _active_session_tracker = None

        # --- Discord: update heartbeat status and poll inbound messages ---
        if getattr(args, "discord", False) and _discord_loop_thread is not None:
            # Refresh Discord status in the heartbeat
            discord_status = _discord_loop_thread.get_discord_status()
            if _heartbeat_thread is not None:
                _heartbeat_thread.discord_status = discord_status

            # Poll pending inbound messages
            _poll_discord_messages(project_id, _get_api_url(), _get_api_key())

        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

        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_db_path: str | None = None
_heartbeat_thread: HeartbeatThread | None = None
_active_session_tracker: SessionTracker | None = None
_discord_loop_thread = None  # DiscordLoopThread | 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, _discord_loop_thread

    if _discord_loop_thread is not None:
        log.info("Stopping Discord loop thread...")
        try:
            _discord_loop_thread.stop()
        except Exception as exc:
            log.warning("Error stopping Discord loop thread: %s", exc)
        _discord_loop_thread = None

    if _heartbeat_thread:
        _heartbeat_thread.stop()
        _heartbeat_thread = None

    if _active_loop_id and _active_db_path:
        delete_heartbeat(_active_db_path, _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)


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


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

    Returns the repository name (e.g., '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 the last path component (repo name)
        if ":" in remote_url:
            # SSH format: git@github.com:user/repo
            repo_name = remote_url.split(":")[-1].split("/")[-1]
        else:
            # HTTPS format: https://github.com/user/repo
            repo_name = remote_url.rstrip("/").split("/")[-1]

        return repo_name if repo_name else 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="SQLite DB path (default: $PEON_DB_PATH or ~/.peon/<project_id>/peon_tasks.db).",
    )
    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(
        "--policy",
        default=None,
        help="Override the project default tool policy by name for this run.",
    )
    parser.add_argument(
        "--no-loop-detection",
        action="store_true",
        default=False,
        help="Disable loop detection (useful for debugging).",
    )
    parser.add_argument(
        "--discord",
        action="store_true",
        default=False,
        help="Enable Discord bot integration (opt-in). Starts bots for enabled configs.",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    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)

    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()
