"""Process management: child tracking, tree killing, streaming runner.

All subprocess lifecycle concerns — spawning, monitoring, killing, and
signal handling — live here.
"""

import os
import re
import signal
import subprocess
import sys
import threading
import time
import uuid
from pathlib import Path
from typing import Callable, List, Optional, Set

from common import CLAUDE_CMD, CODEX_CMD, CURSOR_CMD, OPENCODE_CMD, log

# ---------------------------------------------------------------------------
# Active child PID tracking (thread-safe)
# ---------------------------------------------------------------------------

_active_child_pids: List[int] = []
_active_child_lock = threading.Lock()
_last_streaming_pid: Optional[int] = None


def register_child_pid(pid: int) -> None:
    with _active_child_lock:
        _active_child_pids.append(pid)


def unregister_child_pid(pid: int) -> None:
    with _active_child_lock:
        try:
            _active_child_pids.remove(pid)
        except ValueError:
            pass


def get_last_agent_pid() -> Optional[int]:
    """Return the PID of the most recent ``run_streaming`` child process."""
    return _last_streaming_pid


def _kill_all_children() -> None:
    with _active_child_lock:
        pids = list(_active_child_pids)
    for pid in pids:
        kill_tree(pid)


# ---------------------------------------------------------------------------
# Process tree helpers
# ---------------------------------------------------------------------------


def _get_descendants(pid: int) -> List[int]:
    try:
        result = subprocess.run(
            ["pgrep", "-P", str(pid)],
            capture_output=True, text=True, timeout=5,
        )
        children = [int(p) for p in result.stdout.strip().split() if p.strip()]
    except Exception:
        return []
    descendants = list(children)
    for child in children:
        descendants.extend(_get_descendants(child))
    return descendants


def kill_tree(pid: int) -> None:
    """Send SIGTERM then SIGKILL to *pid* and all its descendants."""
    pids = _get_descendants(pid) + [pid]
    for p in pids:
        try:
            os.kill(p, signal.SIGTERM)
        except OSError:
            pass
    time.sleep(2)
    for p in pids:
        try:
            os.kill(p, signal.SIGKILL)
        except OSError:
            pass


# ---------------------------------------------------------------------------
# Signal handlers
# ---------------------------------------------------------------------------


def install_signal_handlers(log_fh) -> None:
    """Install SIGTERM/SIGINT handlers that kill all tracked children."""
    def _handler(signum, _frame):
        signame = signal.Signals(signum).name
        log(log_fh, f"SIGNAL {signame} received — killing all child processes")
        _kill_all_children()
        log(log_fh, "Child processes killed. Exiting.")
        log_fh.flush()
        sys.exit(128 + signum)
    signal.signal(signal.SIGTERM, _handler)
    signal.signal(signal.SIGINT, _handler)


# ---------------------------------------------------------------------------
# Daemon lifecycle
# ---------------------------------------------------------------------------


def kill_previous_daemon(project_dir: Path) -> None:
    """Kill a previously running daemon for this project, if any."""
    pid_file = project_dir / ".autobuilder" / "daemon.pid"
    if not pid_file.exists():
        return
    try:
        old_pid = int(pid_file.read_text(encoding="utf-8").strip())
    except (ValueError, OSError):
        return
    if old_pid == os.getpid():
        return
    try:
        os.kill(old_pid, 0)
    except OSError:
        pass
    else:
        print(f"Killing previous daemon (pid {old_pid}) and its descendants...")
        kill_tree(old_pid)
        time.sleep(3)


def cleanup_dev_servers() -> None:
    """No-op. Dev servers are intentionally left running across iterations.

    Vite/webpack HMR picks up code changes automatically, so restarting is
    unnecessary.  The old implementation used ``pkill -f "vite"`` which
    matched *all* processes whose command line contained "vite" — including
    Cursor agent processes (the prompt text is a CLI argument and mentions
    "Vite").  This caused cross-agent kills when multiple daemons ran in
    parallel.
    """
    pass


# ---------------------------------------------------------------------------
# Browser cleanup — multi-builder safe
# ---------------------------------------------------------------------------

_BROWSER_CMD_PATTERNS_OWN = [
    "agent-browser/dist/daemon",
    "agent-browser/bin/cli",
    "playwright/chromium",
    "playwright/.local-browsers",
    ".cache/ms-playwright",
]

_BROWSER_CMD_PATTERNS_BROAD = [
    "--remote-debugging-pipe",
    "--remote-debugging-port",
]

_BROWSER_CMD_PATTERNS = _BROWSER_CMD_PATTERNS_OWN + _BROWSER_CMD_PATTERNS_BROAD

BROWSER_MAX_AGE_SECONDS = 3 * 3600  # 3 h safety net


def close_browser_session(session_name: str, log_fh=None) -> bool:
    """Gracefully close an agent-browser session by name.

    Returns True if the close command succeeded (or agent-browser is not
    installed).  This is the preferred way to shut down browser processes —
    it lets the daemon clean up properly instead of force-killing PIDs.
    """
    try:
        result = subprocess.run(
            ["agent-browser", "--session", session_name, "close"],
            capture_output=True, text=True, timeout=15,
        )
        if result.returncode == 0:
            if log_fh:
                log(log_fh, f"Browser session '{session_name}' closed.")
            return True
        if log_fh:
            log(log_fh, f"Browser close returned {result.returncode}: "
                        f"{result.stderr.strip() or result.stdout.strip()}")
    except FileNotFoundError:
        pass
    except subprocess.TimeoutExpired:
        if log_fh:
            log(log_fh, f"Browser close timed out for session '{session_name}'.")
    except Exception as exc:
        if log_fh:
            log(log_fh, f"Browser close error: {exc}")
    return False


def _parse_etime(etime: str) -> int:
    """Convert ``ps`` etime format ``[[dd-]hh:]mm:ss`` to total seconds."""
    etime = etime.strip()
    days = 0
    if "-" in etime:
        day_part, etime = etime.split("-", 1)
        days = int(day_part)
    parts = [int(p) for p in etime.split(":")]
    if len(parts) == 3:
        return days * 86400 + parts[0] * 3600 + parts[1] * 60 + parts[2]
    if len(parts) == 2:
        return days * 86400 + parts[0] * 60 + parts[1]
    return days * 86400 + parts[0]


def _list_browser_processes(
    patterns: Optional[List[str]] = None,
) -> List[tuple]:
    """Return ``[(pid, age_seconds)]`` for browser-automation processes.

    *patterns* defaults to all patterns (own + broad).  Pass
    ``_BROWSER_CMD_PATTERNS_OWN`` to restrict to agent-browser / playwright
    only (avoids matching the user's personal Chrome).
    """
    if patterns is None:
        patterns = _BROWSER_CMD_PATTERNS
    try:
        r = subprocess.run(
            ["ps", "-A", "-o", "pid,etime,command"],
            capture_output=True, text=True, timeout=10,
        )
    except Exception:
        return []
    results = []
    for line in r.stdout.splitlines()[1:]:
        line = line.strip()
        if not line:
            continue
        if not any(pat in line for pat in patterns):
            continue
        parts = line.split(None, 2)
        if len(parts) < 2:
            continue
        try:
            results.append((int(parts[0]), _parse_etime(parts[1])))
        except (ValueError, IndexError):
            continue
    return results


def snapshot_browser_pids() -> Set[int]:
    """Return the set of currently running browser-automation PIDs.

    Call this **before** the agent starts.  After the agent exits, pass the
    result to :func:`cleanup_browsers` so it can diff against the current
    state and kill only browsers spawned during *this* iteration.
    """
    return {pid for pid, _ in _list_browser_processes()}


def cleanup_browsers(
    agent_pid: int,
    log_fh=None,
    pre_snapshot: Optional[Set[int]] = None,
) -> int:
    """Post-iteration browser cleanup (safe for multiple builders).

    Two complementary strategies:

    1. **Snapshot diff** — if *pre_snapshot* is provided (a set of browser
       PIDs captured before the agent started), kill any browser-automation
       process that appeared *after* the snapshot.  These were spawned
       during our iteration and are almost certainly ours.  This replaces
       the old descendant-walk approach which silently failed because the
       agent process is already dead by the time cleanup runs (children get
       reparented to PID 1).

    2. **Age-based safety net** — kill any browser-automation process on the
       host that has been alive longer than ``BROWSER_MAX_AGE_SECONDS``
       (3 h).  This catches leaks from *any* builder instance (including
       previous runs that crashed without cleanup).

    Returns the total number of processes killed.
    """
    killed = 0

    # --- Strategy 1: browsers spawned during this iteration ---
    if pre_snapshot is not None:
        current_browsers = _list_browser_processes()
        new_pids = {pid for pid, _ in current_browsers} - pre_snapshot
        for pid in sorted(new_pids):
            try:
                os.kill(pid, 0)
            except OSError:
                continue
            try:
                kill_tree(pid)
                killed += 1
                if log_fh:
                    log(log_fh, f"Browser cleanup: killed pid {pid} "
                                f"(spawned during iteration)")
            except Exception:
                continue

    # --- Strategy 2: stale agent-browser / playwright only (>= 3 h) ---
    # Uses narrow patterns so we never kill the user's personal Chrome.
    own_procs = _list_browser_processes(patterns=_BROWSER_CMD_PATTERNS_OWN)
    for pid, age in own_procs:
        if age >= BROWSER_MAX_AGE_SECONDS:
            try:
                os.kill(pid, 0)
            except OSError:
                continue
            try:
                kill_tree(pid)
                killed += 1
                if log_fh:
                    log(log_fh, f"Browser cleanup: killed pid {pid} "
                                f"(stale — running for {age // 3600}h {(age % 3600) // 60}m)")
            except Exception:
                continue

    return killed


# ---------------------------------------------------------------------------
# Subprocess runner with streaming + hard timeout
# ---------------------------------------------------------------------------


def run_streaming(
    cmd: List[str],
    cwd: str,
    log_fh,
    hard_timeout: int,
    label: str = "",
    done_callback=None,
    done_grace: int = 120,
    line_callback: Optional[Callable[[str], None]] = None,
) -> int:
    """Run *cmd*, stream its output, and enforce a hard timeout.

    Returns the process exit code, or 124 on timeout.
    If *done_callback* returns ``True`` and the process doesn't exit within
    *done_grace* seconds, the process is killed and 0 is returned.
    If *line_callback* is provided, it is invoked with each streamed line
    (without trailing newline). Exceptions raised by it are swallowed so
    output is never blocked.
    """
    # Strip PORT to prevent the dashboard's port from leaking into project
    # dev servers (Vite/Next.js read process.env.PORT to choose their port).
    agent_env = dict(os.environ)
    agent_env.pop("PORT", None)

    global _last_streaming_pid
    proc = subprocess.Popen(
        cmd, cwd=cwd, env=agent_env,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        text=True, bufsize=1,
    )
    _last_streaming_pid = proc.pid
    register_child_pid(proc.pid)

    deadline = time.monotonic() + hard_timeout
    timed_out = False
    work_done_kill = False

    def watchdog():
        nonlocal timed_out, work_done_kill
        done_since: Optional[float] = None
        while proc.poll() is None:
            if time.monotonic() > deadline:
                timed_out = True
                log(log_fh, f"WATCHDOG: hard timeout ({hard_timeout}s) for {label} — killing pid {proc.pid}")
                kill_tree(proc.pid)
                return
            if done_callback is not None:
                try:
                    is_done = done_callback()
                except Exception:
                    is_done = False
                if is_done:
                    if done_since is None:
                        done_since = time.monotonic()
                        log(log_fh, f"WATCHDOG: {label} work complete — waiting {done_grace}s for clean exit")
                    elif time.monotonic() - done_since > done_grace:
                        work_done_kill = True
                        log(log_fh, f"WATCHDOG: {label} hanging after work complete — killing pid {proc.pid}")
                        kill_tree(proc.pid)
                        return
                else:
                    done_since = None
            time.sleep(10)

    t = threading.Thread(target=watchdog, daemon=True)
    t.start()

    try:
        for line in proc.stdout:
            fh_line = line if line.endswith("\n") else line + "\n"
            log_fh.write(fh_line)
            log_fh.flush()
            try:
                sys.stdout.write(line)
                sys.stdout.flush()
            except (BrokenPipeError, OSError):
                pass
            if line_callback is not None:
                try:
                    line_callback(line.rstrip("\n"))
                except Exception:
                    pass
        proc.wait()
    except KeyboardInterrupt:
        kill_tree(proc.pid)
        log(log_fh, "INTERRUPTED by user (Ctrl+C)")
        raise
    finally:
        unregister_child_pid(proc.pid)

    t.join(timeout=5)

    if work_done_kill:
        return 0
    if timed_out:
        return 124
    return proc.returncode


# ---------------------------------------------------------------------------
# Executor session adapter (chained, multi-phase agent calls)
# ---------------------------------------------------------------------------

# UUID v4 with optional surrounding markers — used to scrape the codex banner.
_UUID_RE = re.compile(
    r"\b([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-"
    r"[0-9a-fA-F]{4}-[0-9a-fA-F]{12})\b"
)
_CODEX_SESSION_HINTS = ("session id", "session_id", "session:", "conversation id")

# OpenCode session IDs are not UUIDs — they're typed prefixes like
# ``ses_65b3acf58ffeLSa4dfj1RVoPpW``. We scrape them out of the
# ``--format json`` event stream, where every event carries
# ``"sessionID":"ses_..."``. The body is permissive (alphanumerics,
# underscores, dashes) so we don't break if opencode's id scheme
# evolves; the ``ses_`` prefix keeps us from matching codex UUIDs.
_OPENCODE_SESSION_RE = re.compile(
    r'"sessionID"\s*:\s*"(ses_[A-Za-z0-9_-]+)"'
)


def open_chat_session(executor: str, cwd: str, log_fh) -> Optional[str]:
    """Create a new chat session for *executor* and return its id.

    For ``cursor``: invokes ``cursor-agent create-chat`` (with stdin closed
    so the CLI exits immediately) and returns the printed UUID.

    For ``codex``: returns ``None`` because codex creates the session
    implicitly on the first ``codex exec`` call. The caller should later
    capture the session UUID from the streamed banner via
    :func:`scan_codex_session_line`.

    For ``claude``: pre-generates a UUID locally; the first phase call
    will materialise the session via ``--session-id <uuid>`` and
    subsequent calls reuse it via ``--resume <uuid>``.

    For ``opencode``: returns ``None`` because opencode creates the
    session implicitly on the first ``opencode run`` call. The caller
    should later capture the typed session id (``ses_…``) from the
    ``--format json`` event stream via
    :func:`scan_opencode_session_line`.

    Returns ``None`` if the call fails. Callers should fall back to
    non-resumed phase invocations in that case.
    """
    if executor == "cursor":
        try:
            result = subprocess.run(
                [CURSOR_CMD, "agent", "create-chat"],
                cwd=cwd,
                stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                timeout=30,
            )
        except (subprocess.TimeoutExpired, FileNotFoundError, OSError) as exc:
            log(log_fh, f"open_chat_session(cursor) failed: {exc}")
            return None
        if result.returncode != 0:
            log(log_fh,
                f"open_chat_session(cursor) exit {result.returncode}: "
                f"{result.stdout.strip()[:200]}")
            return None
        match = _UUID_RE.search(result.stdout or "")
        if match:
            return match.group(1)
        log(log_fh,
            f"open_chat_session(cursor) returned no UUID: "
            f"{(result.stdout or '').strip()[:200]}")
        return None

    if executor == "codex":
        # Codex creates the session on first exec; nothing to do up front.
        return None

    if executor == "claude":
        # Claude Code accepts a caller-provided UUID via --session-id on
        # the first invocation; we just generate one locally.
        return str(uuid.uuid4())

    if executor == "opencode":
        # OpenCode mints the session id on the first ``run`` call and
        # echoes it on every JSON event. We scrape it from the stream.
        return None

    log(log_fh, f"open_chat_session: unknown executor {executor!r}")
    return None


def build_phase_cmd(
    executor: str,
    chat_id: Optional[str],
    model: Optional[str],
    prompt: str,
    *,
    is_new_session: bool = False,
) -> List[str]:
    """Return the CLI command for a single phase invocation.

    *chat_id* is the conversation identifier returned by
    :func:`open_chat_session` (cursor / claude) or scraped from the previous
    phase's output (codex). When ``None``, the command is built without
    resume — the agent will start a fresh conversation.

    *is_new_session* tells executors that pre-generate the chat id whether
    this is the first call (``True``) or a continuation (``False``).
    Currently only used by the ``claude`` branch: ``True`` →
    ``--session-id <uuid>``; ``False`` → ``--resume <uuid>``.
    """
    if executor == "codex":
        if chat_id:
            cmd = [
                CODEX_CMD, "exec", "resume", chat_id,
                "--dangerously-bypass-approvals-and-sandbox",
            ]
        else:
            cmd = [CODEX_CMD, "exec",
                   "--dangerously-bypass-approvals-and-sandbox"]
        if model:
            cmd.extend(["--model", model])
        cmd.append(prompt)
        return cmd

    if executor == "claude":
        cmd = [CLAUDE_CMD, "--print", "--dangerously-skip-permissions"]
        if chat_id:
            if is_new_session:
                cmd.extend(["--session-id", chat_id])
            else:
                cmd.extend(["--resume", chat_id])
        if model:
            cmd.extend(["--model", model])
        cmd.append(prompt)
        return cmd

    if executor == "opencode":
        # ``--format json`` is required so we can scrape the session id
        # off the event stream; ``--dangerously-skip-permissions`` makes
        # the call headless (every tool action is auto-approved).
        cmd = [
            OPENCODE_CMD, "run",
            "--format", "json",
            "--dangerously-skip-permissions",
        ]
        if chat_id:
            cmd.extend(["--session", chat_id])
        if model:
            cmd.extend(["--model", model])
        cmd.append(prompt)
        return cmd

    # Default: cursor
    cmd = [CURSOR_CMD, "agent", "--print", "--yolo", "--approve-mcps",
           "--disable-auto-update", "--endless-retries"]
    if chat_id:
        cmd.extend(["--resume", chat_id])
    if model:
        cmd.extend(["--model", model])
    cmd.append(prompt)
    return cmd


def scan_codex_session_line(line: str) -> Optional[str]:
    """Return a session UUID if *line* looks like a codex session banner.

    Codex prints the session UUID in its startup banner (e.g.
    ``session id: 7f7116f0-...``). We accept any line that mentions a
    session/conversation hint and contains a UUID.
    """
    lower = line.lower()
    if not any(hint in lower for hint in _CODEX_SESSION_HINTS):
        return None
    match = _UUID_RE.search(line)
    return match.group(1) if match else None


def scan_opencode_session_line(line: str) -> Optional[str]:
    """Return an opencode session id (``ses_…``) if *line* contains one.

    With ``opencode run --format json`` every event line is an nd-JSON
    object that carries ``"sessionID":"ses_..."``. We grab it via
    regex rather than full JSON parsing so a truncated or split line
    still yields the id as soon as the first complete event arrives.
    """
    match = _OPENCODE_SESSION_RE.search(line)
    return match.group(1) if match else None

