"""
In-memory registry for tracking running agent processes.

This module provides thread-safe tracking of running agents that complements
the database storage. It tracks runtime information like PIDs and process handles
that shouldn't be persisted.
"""

from __future__ import annotations

import logging
import re
import threading
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any

logger = logging.getLogger(__name__)

# Event callback type - (event_type, run_id, data)
EventCallback = Callable[[str, str, dict[str, Any]], None]


@dataclass
class RunningAgent:
    """
    In-memory record of a running agent process.

    Tracks runtime state that isn't appropriate for database storage.
    """

    run_id: str
    """Agent run ID (matches database record)."""

    session_id: str
    """Child session ID for this agent."""

    parent_session_id: str
    """Parent session that spawned this agent."""

    mode: str
    """Execution mode: in_process, terminal, embedded, headless."""

    started_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    """When the agent started running."""

    # Process tracking (for terminal/embedded/headless modes)
    pid: int | None = None
    """Process ID if running externally."""

    master_fd: int | None = None
    """PTY master file descriptor (embedded mode only)."""

    terminal_type: str | None = None
    """Terminal type (ghostty, iterm, etc.) for terminal mode."""

    # State tracking
    provider: str = "claude"
    """LLM provider being used."""

    workflow_name: str | None = None
    """Workflow being executed, if any."""

    worktree_id: str | None = None
    """Associated worktree, if any."""

    # In-process agent tracking
    task: Any | None = None
    """Async task object for in-process agents (asyncio.Task)."""

    monitor_task: Any | None = None
    """Background monitoring task for headless agents (asyncio.Task)."""

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "run_id": self.run_id,
            "session_id": self.session_id,
            "parent_session_id": self.parent_session_id,
            "mode": self.mode,
            "started_at": self.started_at.isoformat(),
            "pid": self.pid,
            "master_fd": self.master_fd,
            "terminal_type": self.terminal_type,
            "provider": self.provider,
            "workflow_name": self.workflow_name,
            "worktree_id": self.worktree_id,
            "has_task": self.task is not None,
        }


class RunningAgentRegistry:
    """
    Thread-safe registry for tracking running agents.

    This registry tracks agents that are currently executing, whether
    in-process or in external processes (terminal/headless). It provides:

    - Thread-safe add/get/remove operations
    - Lookup by run_id, session_id, or parent_session_id
    - PID-based lookup for process management
    - Cleanup of stale entries

    Example:
        >>> registry = RunningAgentRegistry()
        >>> agent = RunningAgent(
        ...     run_id="ar-123",
        ...     session_id="sess-456",
        ...     parent_session_id="sess-parent",
        ...     mode="terminal",
        ...     pid=12345,
        ... )
        >>> registry.add(agent)
        >>> registry.get("ar-123")
        RunningAgent(...)
        >>> registry.remove("ar-123")
    """

    def __init__(self) -> None:
        """Initialize the registry with an empty agents dict and lock."""
        self._agents: dict[str, RunningAgent] = {}
        self._lock = threading.RLock()
        self._logger = logger
        self._event_callbacks: list[EventCallback] = []
        self._event_callbacks_lock = threading.Lock()

    def add_event_callback(self, callback: EventCallback) -> None:
        """
        Add an event callback for agent lifecycle events.

        Callbacks are invoked when agents are added or removed.

        Args:
            callback: Function that receives (event_type, run_id, data)
        """
        with self._event_callbacks_lock:
            self._event_callbacks.append(callback)

    def emit_event(self, event_type: str, run_id: str, data: dict[str, Any]) -> None:
        """
        Emit a custom event to all registered callbacks.

        Args:
            event_type: Type of event (e.g., terminal_output)
            run_id: Agent run ID
            data: Additional event data
        """
        self._emit_event(event_type, run_id, data)

    def _emit_event(self, event_type: str, run_id: str, data: dict[str, Any]) -> None:
        """
        Emit an event to all registered callbacks.

        Args:
            event_type: Type of event (agent_started, agent_completed, etc.)
            run_id: Agent run ID
            data: Additional event data
        """
        # Take a snapshot of callbacks under lock, then iterate outside lock
        with self._event_callbacks_lock:
            callbacks = list(self._event_callbacks)
        for callback in callbacks:
            try:
                callback(event_type, run_id, data)
            except Exception as e:
                self._logger.warning(f"Event callback error: {e}")

    def add(self, agent: RunningAgent) -> None:
        """
        Add a running agent to the registry.

        Args:
            agent: The running agent to track.
        """
        with self._lock:
            self._agents[agent.run_id] = agent
            self._logger.debug(
                f"Registered running agent {agent.run_id} (mode={agent.mode}, pid={agent.pid})"
            )
        # Emit event outside lock
        self._emit_event(
            "agent_started",
            agent.run_id,
            {
                "session_id": agent.session_id,
                "parent_session_id": agent.parent_session_id,
                "mode": agent.mode,
                "provider": agent.provider,
                "pid": agent.pid,
            },
        )

    def get(self, run_id: str) -> RunningAgent | None:
        """
        Get a running agent by run ID.

        Args:
            run_id: The agent run ID.

        Returns:
            The RunningAgent if found, None otherwise.
        """
        with self._lock:
            return self._agents.get(run_id)

    def remove(self, run_id: str, status: str = "completed") -> RunningAgent | None:
        """
        Remove a running agent from the registry.

        Args:
            run_id: The agent run ID to remove.
            status: Final status (completed, failed, cancelled, timeout).

        Returns:
            The removed RunningAgent if found, None otherwise.
        """
        with self._lock:
            agent = self._agents.pop(run_id, None)
            if agent:
                self._logger.debug(f"Unregistered running agent {run_id}")
        # Emit event outside lock
        if agent:
            self._emit_event(
                f"agent_{status}",
                run_id,
                {
                    "session_id": agent.session_id,
                    "parent_session_id": agent.parent_session_id,
                    "mode": agent.mode,
                    "provider": agent.provider,
                },
            )
        return agent

    def kill(
        self,
        run_id: str,
        signal_name: str = "TERM",
        timeout: float = 5.0,
    ) -> dict[str, Any]:
        """
        Kill a running agent process.

        Strategy varies by mode:
        - headless: Direct signal to tracked PID
        - terminal: Check terminal_context for PID, fallback to pgrep
        - embedded: Close PTY fd + signal
        - in_process: Cancel asyncio task

        Args:
            run_id: Agent run ID
            signal_name: Signal without SIG prefix (TERM, KILL)
            timeout: Seconds before escalating TERM → KILL

        Returns:
            Dict with success status and details
        """
        import os
        import signal
        import subprocess  # nosec B404 - subprocess needed for process management
        import time

        agent = self.get(run_id)
        if not agent:
            return {"success": False, "error": "Agent not found in registry"}

        # Handle in_process mode (asyncio.Task)
        if agent.mode == "in_process" and agent.task:
            agent.task.cancel()
            self.remove(run_id, status="cancelled")
            return {"success": True, "message": "Cancelled in-process task"}

        # For terminal mode, find PID via multiple strategies
        target_pid = agent.pid
        found_via = "registry"

        if agent.mode == "terminal" and agent.session_id:
            # Strategy 1: Check session's terminal_context (Claude hooks)
            try:
                from gobby.storage.database import LocalDatabase
                from gobby.storage.sessions import LocalSessionManager

                db = LocalDatabase()
                session_mgr = LocalSessionManager(db)
                session = session_mgr.get(agent.session_id)
                if session and session.terminal_context:
                    ctx_pid = session.terminal_context.get("parent_pid")
                    if ctx_pid:
                        target_pid = int(ctx_pid)
                        found_via = "terminal_context"
                        self._logger.info(f"Found PID from session terminal_context: {target_pid}")
            except Exception as e:
                self._logger.debug(f"terminal_context lookup failed: {e}")

            # Strategy 2: pgrep fallback (for Codex/Gemini without hooks)
            if found_via == "registry" or not target_pid:
                # Validate session_id format (UUID or safe identifier) to prevent injection
                session_id_pattern = re.compile(r"^[a-zA-Z0-9_-]+$")
                if not session_id_pattern.match(agent.session_id):
                    self._logger.warning(
                        f"Invalid session_id format, skipping pgrep: {agent.session_id}"
                    )
                else:
                    try:
                        # Use -- to prevent pgrep from interpreting pattern as options
                        result = subprocess.run(  # nosec B603 B607 - pgrep with validated session_id
                            ["pgrep", "-f", "--", f"session-id {agent.session_id}"],
                            capture_output=True,
                            text=True,
                            timeout=5.0,
                        )
                        if result.returncode == 0 and result.stdout.strip():
                            pids = result.stdout.strip().split("\n")
                            if len(pids) == 1:
                                target_pid = int(pids[0])
                                found_via = "pgrep"
                                self._logger.info(f"Found PID via pgrep: {target_pid}")
                            else:
                                # Multiple PIDs found - need to disambiguate
                                self._logger.warning(
                                    f"pgrep returned {len(pids)} PIDs for session "
                                    f"{agent.session_id}: {pids}"
                                )
                                # Inspect each candidate to find the correct one
                                matched_pid = None
                                for pid_str in pids:
                                    try:
                                        candidate_pid = int(pid_str)
                                        # Query the process command line to verify
                                        ps_result = subprocess.run(  # nosec B603 B607 - ps with numeric PID
                                            ["ps", "-p", str(candidate_pid), "-o", "args="],
                                            capture_output=True,
                                            text=True,
                                            timeout=2.0,
                                        )
                                        if ps_result.returncode == 0:
                                            cmdline = ps_result.stdout.strip()
                                            # Verify it's actually the agent process
                                            # (contains session-id and matches expected CLI)
                                            # Check both provider name and "claude" (for cursor/windsurf/copilot)
                                            is_matched = agent.provider in cmdline.lower()
                                            if not is_matched and agent.provider in (
                                                "cursor",
                                                "windsurf",
                                                "copilot",
                                            ):
                                                is_matched = "claude" in cmdline.lower()

                                            if (
                                                f"session-id {agent.session_id}" in cmdline
                                                and is_matched
                                            ):
                                                if matched_pid is not None:
                                                    # Multiple matches - ambiguous
                                                    self._logger.error(
                                                        f"Ambiguous PID match: both {matched_pid} "
                                                        f"and {candidate_pid} match session "
                                                        f"{agent.session_id}"
                                                    )
                                                    matched_pid = None
                                                    break
                                                matched_pid = candidate_pid
                                    except (ValueError, subprocess.TimeoutExpired):
                                        continue
                                if matched_pid is not None:
                                    target_pid = matched_pid
                                    found_via = "pgrep_disambiguated"
                                    self._logger.info(
                                        f"Disambiguated PID via ps inspection: {target_pid}"
                                    )
                                else:
                                    self._logger.error(
                                        f"Could not disambiguate PIDs for session "
                                        f"{agent.session_id}: {pids}"
                                    )
                    except Exception as e:
                        self._logger.warning(f"pgrep fallback failed: {e}")

        if not target_pid:
            return {"success": False, "error": "No target PID found"}

        # Check if process is alive
        try:
            os.kill(target_pid, 0)
        except ProcessLookupError:
            self.remove(run_id, status="completed")
            return {
                "success": True,
                "message": f"Process {target_pid} already dead",
                "already_dead": True,
            }
        except PermissionError:
            return {"success": False, "error": f"No permission to signal PID {target_pid}"}

        # Close PTY if embedded mode
        if agent.master_fd is not None:
            try:
                os.close(agent.master_fd)
            except OSError:
                pass

        # Send signal
        sig = getattr(signal, f"SIG{signal_name}", signal.SIGTERM)
        try:
            os.kill(target_pid, sig)
        except ProcessLookupError:
            self.remove(run_id, status="completed")
            return {
                "success": True,
                "message": "Process died during signal",
                "already_dead": True,
            }

        # Wait for termination with optional SIGKILL escalation
        if signal_name == "TERM" and timeout > 0:
            deadline = time.time() + timeout
            while time.time() < deadline:
                try:
                    os.kill(target_pid, 0)
                    time.sleep(0.1)
                except ProcessLookupError:
                    break
            else:
                # Still alive - escalate to SIGKILL
                try:
                    os.kill(target_pid, signal.SIGKILL)
                    self._logger.info(f"Escalated to SIGKILL for PID {target_pid}")
                except ProcessLookupError:
                    pass

        self.remove(run_id, status="killed")
        return {
            "success": True,
            "message": f"Sent SIG{signal_name} to PID {target_pid}",
            "pid": target_pid,
            "signal": signal_name,
            "found_via": found_via,
        }

    def get_by_session(self, session_id: str) -> RunningAgent | None:
        """
        Get a running agent by its child session ID.

        Args:
            session_id: The child session ID.

        Returns:
            The RunningAgent if found, None otherwise.
        """
        with self._lock:
            for agent in self._agents.values():
                if agent.session_id == session_id:
                    return agent
            return None

    def get_by_pid(self, pid: int) -> RunningAgent | None:
        """
        Get a running agent by its process ID.

        Args:
            pid: The process ID.

        Returns:
            The RunningAgent if found, None otherwise.
        """
        with self._lock:
            for agent in self._agents.values():
                if agent.pid == pid:
                    return agent
            return None

    def list_by_parent(self, parent_session_id: str) -> list[RunningAgent]:
        """
        List all running agents for a parent session.

        Args:
            parent_session_id: The parent session ID.

        Returns:
            List of running agents spawned by this parent.
        """
        with self._lock:
            return [
                agent
                for agent in self._agents.values()
                if agent.parent_session_id == parent_session_id
            ]

    def list_by_mode(self, mode: str) -> list[RunningAgent]:
        """
        List all running agents by execution mode.

        Args:
            mode: Execution mode (in_process, terminal, embedded, headless).

        Returns:
            List of running agents with this mode.
        """
        with self._lock:
            return [agent for agent in self._agents.values() if agent.mode == mode]

    def list_all(self) -> list[RunningAgent]:
        """
        List all running agents.

        Returns:
            List of all running agents (copy of current state).
        """
        with self._lock:
            return list(self._agents.values())

    def count(self) -> int:
        """
        Get the number of running agents.

        Returns:
            Count of running agents.
        """
        with self._lock:
            return len(self._agents)

    def count_by_parent(self, parent_session_id: str) -> int:
        """
        Count running agents for a parent session.

        Args:
            parent_session_id: The parent session ID.

        Returns:
            Count of running agents for this parent.
        """
        with self._lock:
            return sum(
                1 for agent in self._agents.values() if agent.parent_session_id == parent_session_id
            )

    def cleanup_by_pids(self, dead_pids: set[int]) -> list[RunningAgent]:
        """
        Remove agents whose PIDs are no longer running.

        This should be called periodically by a cleanup process that
        checks which PIDs are still alive.

        Args:
            dead_pids: Set of PIDs that are no longer running.

        Returns:
            List of agents that were removed.
        """
        removed: list[RunningAgent] = []
        with self._lock:
            for run_id, agent in list(self._agents.items()):
                if agent.pid and agent.pid in dead_pids:
                    self._agents.pop(run_id)
                    removed.append(agent)
                    self._logger.info(f"Cleaned up agent {run_id} with dead PID {agent.pid}")
        # Emit events outside lock for each removed agent
        for agent in removed:
            self._emit_event(
                "agent_completed",
                agent.run_id,
                {
                    "session_id": agent.session_id,
                    "parent_session_id": agent.parent_session_id,
                    "mode": agent.mode,
                    "provider": agent.provider,
                    "cleanup_reason": "dead_pid",
                },
            )
        return removed

    def cleanup_stale(self, max_age_seconds: float = 3600.0) -> list[RunningAgent]:
        """
        Remove agents that have been running longer than max_age.

        Args:
            max_age_seconds: Maximum age in seconds before cleanup (default: 1 hour).

        Returns:
            List of agents that were removed.
        """
        now = datetime.now(UTC)
        removed: list[RunningAgent] = []
        with self._lock:
            for run_id, agent in list(self._agents.items()):
                age = (now - agent.started_at).total_seconds()
                if age > max_age_seconds:
                    self._agents.pop(run_id)
                    removed.append(agent)
                    self._logger.info(f"Cleaned up stale agent {run_id} (age={age:.0f}s)")
        # Emit events outside lock for each removed agent
        for agent in removed:
            self._emit_event(
                "agent_timeout",
                agent.run_id,
                {
                    "session_id": agent.session_id,
                    "parent_session_id": agent.parent_session_id,
                    "mode": agent.mode,
                    "provider": agent.provider,
                    "cleanup_reason": "stale",
                },
            )
        return removed

    def clear(self) -> int:
        """
        Clear all running agents from the registry.

        Returns:
            Number of agents that were cleared.
        """
        with self._lock:
            count = len(self._agents)
            self._agents.clear()
            self._logger.info(f"Cleared {count} running agents from registry")
            return count


# Global singleton instance
_default_registry: RunningAgentRegistry | None = None
_registry_lock = threading.Lock()


def get_running_agent_registry() -> RunningAgentRegistry:
    """
    Get the global running agent registry singleton.

    Returns:
        The shared RunningAgentRegistry instance.
    """
    global _default_registry
    if _default_registry is None:
        with _registry_lock:
            if _default_registry is None:
                _default_registry = RunningAgentRegistry()
    return _default_registry
