"""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.

Usage:
    peon-loop <project_id> <repo_path> [options]
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import signal
import socket
import sqlite3
import subprocess
import sys
import threading
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from shutil import which

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


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


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

    conn = sqlite3.connect(db_path)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    conn.executescript(SCHEMA)
    conn.execute(
        "INSERT OR IGNORE INTO projects (id, path) VALUES (?, ?)",
        (project_id, repo_path),
    )
    conn.commit()
    conn.close()
    log.info("Ready to work! Peon loop starting for project '%s'", project_id)


def has_todo_tasks(db_path: str, project_id: str) -> 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.

    Returns:
        (todo_count, in_progress_count).  Returns (0, 0) if the DB or table
        doesn't exist.
    """
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.execute(
            "SELECT status, COUNT(*) FROM tasks"
            " WHERE project_id = ? AND status IN ('todo', 'in_progress')"
            " GROUP BY status",
            (project_id,),
        )
        counts = dict(cursor.fetchall())
        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)


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


def write_heartbeat(
    db_path: str,
    loop_id: str,
    project_id: str,
    pid: int,
    hostname: 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, started_at, last_heartbeat)"
            " VALUES (?, ?, ?, ?, COALESCE("
            "   (SELECT started_at FROM loop_heartbeats WHERE id = ?), CURRENT_TIMESTAMP"
            " ), CURRENT_TIMESTAMP)",
            (loop_id, project_id, pid, hostname, 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._stop_event = threading.Event()

    def run(self) -> None:
        """Write heartbeats every interval seconds until stopped."""
        while not self._stop_event.wait(self.interval):
            try:
                write_heartbeat(
                    self.db_path,
                    self.loop_id,
                    self.project_id,
                    self.pid,
                    self.hostname,
                )
            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 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) -> str:
    """Build the agent prompt — mirrors the start_work prompt from server.py."""
    return (
        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. **Clean working tree** — Run `git status` in {repo_path}. If there are "
        "untracked or uncommitted changes, run `git stash --include-untracked` to "
        "stash them and call log_work to record that you stashed pre-existing changes "
        "(include the stash ref in the summary).\n"
        "\n"
        f"3. Run `git rev-parse HEAD` in {repo_path} to capture the current commit SHA.\n"
        "\n"
        f"4. Call next_task(project_id='{project_id}', commit_base='<sha>') passing "
        "the SHA from step 3. This marks the task as in_progress, records the base "
        "commit, and returns the task details.\n"
        "\n"
        f"5. Use the execute_task prompt with project_id='{project_id}' and the task's "
        "title and description to get your full working instructions, then follow them."
    )


# ---------------------------------------------------------------------------
# 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 run_agent(
    cmd: list[str],
    repo_path: str,
    log_dir: Path,
    pid_file: Path,
    timeout: int,
) -> 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.
    """
    log_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    log_file = log_dir / f"agent-{timestamp}.log"

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

    with open(log_file, "w") as fout:
        proc = subprocess.Popen(
            cmd,
            cwd=repo_path,
            stdout=fout,
            stderr=subprocess.STDOUT,
        )

    pid_file.write_text(str(proc.pid))

    try:
        proc.wait(timeout=timeout)
    except subprocess.TimeoutExpired:
        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()

    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()
    repo_path = str(Path(args.repo_path).resolve())
    log_dir = Path(repo_path) / ".peon-logs"
    pid_file = log_dir / "agent.pid"

    # Ensure DB and project exist so tasks can be added immediately
    ensure_project(db_path, args.project_id, repo_path)

    # 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)
    prompt = build_prompt(args.project_id, repo_path)

    global _active_loop_id, _active_db_path, _heartbeat_thread
    loop_id = str(uuid.uuid4())
    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)

    # Start background heartbeat thread
    _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)")

    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:
        if _heartbeat_thread:
            log.info("Stopping heartbeat thread...")
            _heartbeat_thread.stop()
            _heartbeat_thread = None


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,
    repo_path: str,
    log_dir: Path,
    pid_file: Path,
) -> int:
    """Execute the main loop iterations (extracted for cleanup handling)."""
    while True:
        iteration += 1
        log.info("--- Iteration %d ---", iteration)

        log.info("Something need doing? Checking for tasks...")
        todo, in_progress = has_todo_tasks(db_path, args.project_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

        # Build the command
        cmd = [
            claude_bin,
            "-p", prompt,
            "--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)])

        kill_previous(pid_file)
        run_agent(cmd, repo_path, log_dir, pid_file, args.timeout)

        if args.once:
            log.info("--once flag set, exiting after single iteration.")
            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


def _handle_sigint(signum: int, frame) -> None:
    """Clean up child process and heartbeat on Ctrl-C."""
    log.info("Me go now... *waves*")
    if _heartbeat_thread:
        _heartbeat_thread.stop()
    if _active_loop_id and _active_db_path:
        delete_heartbeat(_active_db_path, _active_loop_id)
    if _child_proc and _child_proc.poll() is None:
        _child_proc.terminate()
        try:
            _child_proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            _child_proc.kill()
    sys.exit(130)


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


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",
        help="Project identifier (e.g. 'peon-mcp').",
    )
    parser.add_argument(
        "repo_path",
        help="Absolute path to the repo on disk.",
    )
    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(
        "--verbose",
        action="store_true",
        help="Enable debug logging.",
    )
    parser.add_argument(
        "--permission-mode",
        default="bypassPermissions",
        help="Claude permission mode (default: bypassPermissions).",
    )
    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",
    )

    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*")
        if _heartbeat_thread:
            _heartbeat_thread.stop()
        if _active_loop_id and _active_db_path:
            delete_heartbeat(_active_db_path, _active_loop_id)
        sys.exit(130)

    sys.exit(code)


if __name__ == "__main__":
    main()
