"""Sub-agent lifecycle manager and limit enforcement.

All functions use synchronous sqlite3 so they can be called directly from
the peon-loop (which is entirely synchronous) without asyncio overhead.
"""

from __future__ import annotations

import logging
import os
import signal
import sqlite3
import time
from dataclasses import dataclass

log = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data types
# ---------------------------------------------------------------------------


@dataclass
class SubAgent:
    """Represents a row in the sub_agents table."""

    id: int
    project_id: str
    label: str
    status: str
    depth: int
    idempotency_key: str = ""
    model: str = ""
    parent_session_id: int | None = None
    parent_task_id: int | None = None
    pid: int | None = None
    result_summary: str = ""
    tokens_used: int = 0
    timeout_seconds: int = 600
    started_at: str = ""
    ended_at: str | None = None
    created_at: str = ""


@dataclass
class SpawnRequest:
    """Parameters for registering a new sub-agent spawn."""

    project_id: str
    label: str
    idempotency_key: str = ""
    model: str = ""
    parent_task_id: int | None = None
    parent_session_id: int | None = None
    depth: int = 1
    timeout_seconds: int = 600


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _row_to_sub_agent(row: sqlite3.Row) -> SubAgent:
    """Convert a sqlite3.Row to a SubAgent dataclass."""
    d = dict(row)
    return SubAgent(
        id=d["id"],
        project_id=d["project_id"],
        label=d["label"],
        status=d["status"],
        depth=d["depth"],
        idempotency_key=d.get("idempotency_key") or "",
        model=d.get("model") or "",
        parent_session_id=d.get("parent_session_id"),
        parent_task_id=d.get("parent_task_id"),
        pid=d.get("pid"),
        result_summary=d.get("result_summary") or "",
        tokens_used=d.get("tokens_used") or 0,
        timeout_seconds=d.get("timeout_seconds") or 600,
        started_at=d.get("started_at") or "",
        ended_at=d.get("ended_at"),
        created_at=d.get("created_at") or "",
    )


def _try_kill_pid(pid: int, sigterm_timeout: int = 5) -> bool:
    """Attempt to kill a process by PID.

    Sends SIGTERM first, waits up to `sigterm_timeout` seconds, then sends
    SIGKILL if the process has not exited.

    Returns:
        True if the process was killed or was already gone, False on error.
    """
    try:
        os.kill(pid, 0)  # Check if process exists
    except OSError:
        return True  # Already gone

    try:
        log.info("Sending SIGTERM to sub-agent process %d", pid)
        os.kill(pid, signal.SIGTERM)
        deadline = time.monotonic() + sigterm_timeout
        while time.monotonic() < deadline:
            time.sleep(0.5)
            try:
                os.kill(pid, 0)
            except OSError:
                log.debug("Process %d exited after SIGTERM", pid)
                return True
        # Still alive — escalate to SIGKILL
        log.warning("Process %d did not exit after SIGTERM, sending SIGKILL", pid)
        os.kill(pid, signal.SIGKILL)
        return True
    except OSError as exc:
        log.warning("Failed to kill process %d: %s", pid, exc)
        return False


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def validate_spawn(
    db: sqlite3.Connection,
    project_id: str,
    parent_task_id: int | None,
    depth: int,
) -> tuple[bool, str]:
    """Check whether a new sub-agent spawn is within project limits.

    Checks (in order):
    1. ``max_subagent_depth`` — reject if ``depth`` exceeds the limit.
    2. ``max_subagents_per_task`` — count active sub-agents for this task.
    3. ``max_concurrent_subagents`` — count all active sub-agents for the project.

    Args:
        db: Open sqlite3 connection with row_factory already set or not.
        project_id: The project to check limits for.
        parent_task_id: The task that wants to spawn (None for no task).
        depth: Nesting depth of the proposed sub-agent (1 = direct child).

    Returns:
        ``(True, "")`` if the spawn is allowed, or ``(False, reason)`` if not.
    """
    db.row_factory = sqlite3.Row

    cursor = db.execute(
        "SELECT max_subagent_depth, max_subagents_per_task, max_concurrent_subagents"
        " FROM projects WHERE id = ?",
        (project_id,),
    )
    row = cursor.fetchone()
    if row is None:
        return (False, f"Project '{project_id}' not found")

    max_depth = row["max_subagent_depth"]
    max_per_task = row["max_subagents_per_task"]
    max_concurrent = row["max_concurrent_subagents"]

    if depth > max_depth:
        return (False, f"Depth {depth} exceeds project limit of {max_depth}")

    if parent_task_id is not None:
        cursor = db.execute(
            "SELECT COUNT(*) FROM sub_agents"
            " WHERE parent_task_id = ? AND status IN ('pending', 'running')",
            (parent_task_id,),
        )
        active_for_task = cursor.fetchone()[0]
        if active_for_task >= max_per_task:
            return (
                False,
                f"Task {parent_task_id} already has {active_for_task} active sub-agent(s)"
                f" (limit: {max_per_task})",
            )

    cursor = db.execute(
        "SELECT COUNT(*) FROM sub_agents"
        " WHERE project_id = ? AND status IN ('pending', 'running')",
        (project_id,),
    )
    active_total = cursor.fetchone()[0]
    if active_total >= max_concurrent:
        return (
            False,
            f"Project '{project_id}' already has {active_total} concurrent sub-agent(s)"
            f" (limit: {max_concurrent})",
        )

    return (True, "")


def check_idempotency(
    db: sqlite3.Connection,
    idempotency_key: str,
) -> SubAgent | None:
    """Return an existing sub-agent if this key was already spawned.

    - If a sub-agent with this key is pending/running/completed → return it.
    - If it failed/cancelled → delete the old record and return None
      so a fresh spawn can proceed.
    - Empty key → always return None (no idempotency check).

    Args:
        db: Open sqlite3 connection.
        idempotency_key: Hash or label identifying this specific spawn.

    Returns:
        Existing :class:`SubAgent` if one should be reused, otherwise ``None``.
    """
    if not idempotency_key:
        return None

    db.row_factory = sqlite3.Row
    cursor = db.execute(
        "SELECT * FROM sub_agents WHERE idempotency_key = ?"
        " ORDER BY created_at DESC LIMIT 1",
        (idempotency_key,),
    )
    row = cursor.fetchone()
    if row is None:
        return None

    agent = _row_to_sub_agent(row)

    if agent.status in ("pending", "running", "completed"):
        log.info(
            "Idempotency key '%s' matches sub-agent #%d (status=%s) — skipping spawn",
            idempotency_key,
            agent.id,
            agent.status,
        )
        return agent

    # Failed or cancelled — remove stale record and allow re-spawn
    log.info(
        "Idempotency key '%s' matched %s sub-agent #%d — deleting for re-spawn",
        idempotency_key,
        agent.status,
        agent.id,
    )
    db.execute("DELETE FROM sub_agents WHERE id = ?", (agent.id,))
    db.commit()
    return None


def register_spawn(
    db: sqlite3.Connection,
    spawn_request: SpawnRequest,
) -> SubAgent:
    """Validate limits, check idempotency, and insert a new sub-agent row.

    Args:
        db: Open sqlite3 connection.
        spawn_request: Parameters describing the desired sub-agent.

    Returns:
        The newly created (or previously existing idempotent) :class:`SubAgent`.

    Raises:
        ValueError: If spawn limits would be violated.
    """
    allowed, reason = validate_spawn(
        db,
        spawn_request.project_id,
        spawn_request.parent_task_id,
        spawn_request.depth,
    )
    if not allowed:
        raise ValueError(f"Spawn not allowed: {reason}")

    if spawn_request.idempotency_key:
        existing = check_idempotency(db, spawn_request.idempotency_key)
        if existing is not None:
            return existing

    db.row_factory = sqlite3.Row
    cursor = db.execute(
        """
        INSERT INTO sub_agents
            (parent_session_id, parent_task_id, project_id, label,
             idempotency_key, model, status, depth, timeout_seconds)
        VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?)
        """,
        (
            spawn_request.parent_session_id,
            spawn_request.parent_task_id,
            spawn_request.project_id,
            spawn_request.label,
            spawn_request.idempotency_key,
            spawn_request.model,
            spawn_request.depth,
            spawn_request.timeout_seconds,
        ),
    )
    db.commit()

    new_id = cursor.lastrowid
    cursor = db.execute("SELECT * FROM sub_agents WHERE id = ?", (new_id,))
    row = cursor.fetchone()
    log.info(
        "Registered sub-agent #%d '%s' (project=%s, depth=%d)",
        new_id,
        spawn_request.label,
        spawn_request.project_id,
        spawn_request.depth,
    )
    return _row_to_sub_agent(row)


def monitor_timeouts(db: sqlite3.Connection) -> int:
    """Kill processes and mark running sub-agents that have exceeded their timeout.

    For each timed-out running sub-agent:
    - Attempts to kill the process via PID (SIGTERM then SIGKILL after 5 s).
    - Updates ``status`` to ``'timeout'`` and sets ``ended_at``.

    Args:
        db: Open sqlite3 connection.

    Returns:
        Number of sub-agents that were timed out.
    """
    db.row_factory = sqlite3.Row
    cursor = db.execute(
        """
        SELECT * FROM sub_agents
        WHERE status = 'running'
          AND started_at IS NOT NULL
          AND CAST(strftime('%s', started_at) AS INTEGER) + timeout_seconds
              < CAST(strftime('%s', 'now') AS INTEGER)
        """
    )
    rows = cursor.fetchall()

    timed_out = 0
    for row in rows:
        agent = _row_to_sub_agent(row)
        log.warning(
            "Sub-agent #%d ('%s') timed out (timeout=%ds, pid=%s)",
            agent.id,
            agent.label,
            agent.timeout_seconds,
            agent.pid,
        )

        if agent.pid is not None:
            _try_kill_pid(agent.pid, sigterm_timeout=5)

        db.execute(
            "UPDATE sub_agents SET status = 'timeout', ended_at = CURRENT_TIMESTAMP"
            " WHERE id = ?",
            (agent.id,),
        )
        timed_out += 1

    if timed_out:
        db.commit()
        log.info("Timed out %d sub-agent(s)", timed_out)

    return timed_out


def cleanup_stale(db: sqlite3.Connection, max_age_hours: int = 24) -> int:
    """Mark long-stuck sub-agents as cancelled or timeout.

    Rules:
    - ``pending`` for more than 1 hour → ``cancelled``.
    - ``running`` for more than ``max_age_hours`` → ``timeout``.

    This is a safety net for sub-agents whose process exited without updating
    their own status (e.g., crashed before writing results).

    Args:
        db: Open sqlite3 connection.
        max_age_hours: Maximum allowed age (in hours) for a running sub-agent
            before it is forcibly timed out.

    Returns:
        Total number of sub-agents cleaned up.
    """
    cursor = db.execute(
        """
        UPDATE sub_agents
        SET status = 'cancelled', ended_at = CURRENT_TIMESTAMP
        WHERE status = 'pending'
          AND CAST(strftime('%s', 'now') AS INTEGER)
              - CAST(strftime('%s', started_at) AS INTEGER) > 3600
        """
    )
    cancelled = cursor.rowcount

    cursor = db.execute(
        """
        UPDATE sub_agents
        SET status = 'timeout', ended_at = CURRENT_TIMESTAMP
        WHERE status = 'running'
          AND CAST(strftime('%s', 'now') AS INTEGER)
              - CAST(strftime('%s', started_at) AS INTEGER) > ?
        """,
        (max_age_hours * 3600,),
    )
    timed_out = cursor.rowcount

    total = cancelled + timed_out
    if total:
        db.commit()
        if cancelled:
            log.info("Cancelled %d stale pending sub-agent(s) (stuck > 1 h)", cancelled)
        if timed_out:
            log.info(
                "Timed out %d stale running sub-agent(s) (age > %d h)",
                timed_out,
                max_age_hours,
            )

    return total


def mark_running(db: sqlite3.Connection, sub_agent_id: int, pid: int | None = None) -> None:
    """Mark a sub-agent as running, optionally recording its OS process PID.

    Args:
        db: Open sqlite3 connection.
        sub_agent_id: ID of the sub_agents row to update.
        pid: OS process ID of the spawned agent process (optional).
    """
    db.execute(
        "UPDATE sub_agents SET status = 'running', pid = ?, started_at = CURRENT_TIMESTAMP"
        " WHERE id = ?",
        (pid, sub_agent_id),
    )
    db.commit()
    log.debug("Sub-agent #%d marked running (pid=%s)", sub_agent_id, pid)


def mark_completed(
    db: sqlite3.Connection,
    sub_agent_id: int,
    result_summary: str = "",
    tokens_used: int = 0,
) -> None:
    """Mark a sub-agent as completed.

    Args:
        db: Open sqlite3 connection.
        sub_agent_id: ID of the sub_agents row to update.
        result_summary: Short summary of what the agent accomplished.
        tokens_used: Approximate token count consumed by the agent.
    """
    db.execute(
        "UPDATE sub_agents SET status = 'completed', ended_at = CURRENT_TIMESTAMP,"
        " result_summary = ?, tokens_used = ? WHERE id = ?",
        (result_summary, tokens_used, sub_agent_id),
    )
    db.commit()
    log.debug("Sub-agent #%d marked completed", sub_agent_id)


def mark_failed(
    db: sqlite3.Connection,
    sub_agent_id: int,
    result_summary: str = "",
) -> None:
    """Mark a sub-agent as failed.

    Args:
        db: Open sqlite3 connection.
        sub_agent_id: ID of the sub_agents row to update.
        result_summary: Error description or failure reason.
    """
    db.execute(
        "UPDATE sub_agents SET status = 'failed', ended_at = CURRENT_TIMESTAMP,"
        " result_summary = ? WHERE id = ?",
        (result_summary, sub_agent_id),
    )
    db.commit()
    log.debug("Sub-agent #%d marked failed", sub_agent_id)


def mark_timeout(db: sqlite3.Connection, sub_agent_id: int) -> None:
    """Mark a sub-agent as timed out.

    Args:
        db: Open sqlite3 connection.
        sub_agent_id: ID of the sub_agents row to update.
    """
    db.execute(
        "UPDATE sub_agents SET status = 'timeout', ended_at = CURRENT_TIMESTAMP"
        " WHERE id = ?",
        (sub_agent_id,),
    )
    db.commit()
    log.debug("Sub-agent #%d marked timeout", sub_agent_id)
