"""
ttyd process lifecycle management for the tmux-web muxplex.

Constants:
    TTYD_PID_DIR  — directory for PID files (default: ~/.local/share/tmux-web/)
    TTYD_PID_PATH — full path to the legacy PID file (TTYD_PID_DIR / 'ttyd.pid')
    TTYD_PORT     — base port for the ttyd pool (7682)

Module state:
    _active_process — the currently running ttyd subprocess (or None) [backward compat]
    _pool           — session→PoolEntry mapping for the process pool

Pool API:
    get_or_spawn(session_name)  — Return port; spawn if needed, poll until listening
    pool_port(session_name)     — Look up port for session without spawning
    kill_session(session_name)  — Kill ttyd for one session, remove from pool
    kill_all()                  — Kill all pooled ttyd processes
    kill_orphans()              — Startup cleanup: PID files, port range scan

Backward-compat API:
    spawn_ttyd(session_name)    — spawn ttyd attached to a tmux session, write PID file
    kill_ttyd()                 — kill the running ttyd process, clean up PID file
    kill_orphan_ttyd()          — kill any orphaned ttyd from a previous coordinator run
"""

import asyncio
import os
import re as _re
import signal
import socket as _socket
import subprocess as _subprocess
import time
from dataclasses import dataclass
from pathlib import Path

# ---------------------------------------------------------------------------
# Paths and constants
# ---------------------------------------------------------------------------

_default_ttyd_pid_dir = Path.home() / ".local" / "share" / "tmux-web"
TTYD_PID_DIR: Path = Path(os.environ.get("TMUX_WEB_STATE_DIR", _default_ttyd_pid_dir))
TTYD_PID_PATH: Path = TTYD_PID_DIR / "ttyd.pid"

TTYD_PORT: int = 7682

_PORT_MIN: int = 7682
_PORT_MAX: int = 7701

# ---------------------------------------------------------------------------
# Module state
# ---------------------------------------------------------------------------

_active_process: asyncio.subprocess.Process | None = None


@dataclass
class PoolEntry:
    """A pooled ttyd process and its assigned port."""

    process: asyncio.subprocess.Process
    port: int


_pool: dict[str, PoolEntry] = {}

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


def _pid_path_for_port(port: int) -> Path:
    """Return the per-port PID file path: ``ttyd-{port}.pid`` in TTYD_PID_DIR."""
    return TTYD_PID_DIR / f"ttyd-{port}.pid"


def _is_alive(entry: PoolEntry) -> bool:
    """Return True if the pool entry's process has not exited."""
    return entry.process.returncode is None


def _ttyd_is_listening(port: int = TTYD_PORT) -> bool:
    """Return True if something is accepting TCP connections on *port*.

    Uses a raw socket connect (no WebSocket handshake, no PTY spawned).
    Takes < 1 ms on localhost when ttyd is running; fails immediately with
    ConnectionRefusedError when it's not.  OSError/TimeoutError are also
    caught so the caller always gets a bool.
    """
    try:
        with _socket.create_connection(("127.0.0.1", port), timeout=0.5):
            return True
    except (ConnectionRefusedError, OSError, TimeoutError):
        return False


def _allocate_port() -> int:
    """Return the next free port in the pool range [_PORT_MIN, _PORT_MAX]."""
    used_ports = {e.port for e in _pool.values()}
    for port in range(_PORT_MIN, _PORT_MAX + 1):
        if port not in used_ports:
            return port
    raise RuntimeError("ttyd pool exhausted: all ports in range are in use")


def _pids_on_port(port: int) -> list[int]:
    """Return PIDs of all processes listening on *port*.

    Tries three tools in order, stopping at the first that returns results:

    1. ``lsof -ti :<port>``  — one PID per line, most widely available.
    2. ``fuser <port>/tcp``  — space-separated PIDs on stdout (psmisc).
    3. ``ss -Hnltp``         — parses ``pid=N`` from users field (iproute2).

    Returns an empty list if none of the tools are available or find anything.
    Silently swallows all errors so callers always get a list.
    """
    pids: list[int] = []

    # --- Tool 1: lsof ---
    try:
        result = _subprocess.run(
            ["lsof", "-ti", f":{port}"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0 and result.stdout.strip():
            for tok in result.stdout.split():
                try:
                    pids.append(int(tok))
                except ValueError:
                    pass
    except Exception:  # noqa: BLE001
        pass

    if pids:
        return pids

    # --- Tool 2: fuser (psmisc) ---
    try:
        result = _subprocess.run(
            ["fuser", f"{port}/tcp"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        # fuser writes the port label to stderr and PIDs to stdout.
        if result.stdout.strip():
            for tok in result.stdout.split():
                try:
                    pids.append(int(tok))
                except ValueError:
                    pass
    except Exception:  # noqa: BLE001
        pass

    if pids:
        return pids

    # --- Tool 3: ss (iproute2) ---
    try:
        result = _subprocess.run(
            ["ss", "-Hnltp", f"sport = :{port}"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0 and result.stdout.strip():
            for match in _re.finditer(r"pid=(\d+)", result.stdout):
                try:
                    pids.append(int(match.group(1)))
                except ValueError:
                    pass
    except Exception:  # noqa: BLE001
        pass

    return pids


def _kill_pids_on_port(port: int, sig: int) -> bool:
    """Find and signal all processes listening on *port*.

    Uses :func:`_pids_on_port` (lsof → fuser → ss) to locate PIDs, then
    signals each with *sig*.

    Returns True if at least one PID was found and signalled.
    Silently ignores unavailable tools and already-dead processes.
    """
    pids = _pids_on_port(port)
    if not pids:
        return False
    sent = False
    for pid in pids:
        try:
            os.kill(pid, sig)
            sent = True
        except (ProcessLookupError, PermissionError):
            pass
    return sent


# ---------------------------------------------------------------------------
# Pool API
# ---------------------------------------------------------------------------


async def get_or_spawn(session_name: str) -> int:
    """Return the port for *session_name*, spawning ttyd if needed.

    If the session already has a live ttyd in the pool (process alive AND
    port is listening), returns immediately.  Otherwise allocates the next
    free port from [_PORT_MIN, _PORT_MAX], spawns ttyd, polls until
    listening (up to ~1 s), and returns the port.
    """
    # Check existing pool entry — require both alive process AND listening port
    entry = _pool.get(session_name)
    if entry is not None:
        if _is_alive(entry) and _ttyd_is_listening(entry.port):
            return entry.port
        # Stale entry — clean up
        try:
            entry.process.terminate()
        except ProcessLookupError:
            pass
        _pid_path_for_port(entry.port).unlink(missing_ok=True)
        del _pool[session_name]

    # Allocate next free port
    port = _allocate_port()

    # Force-free the port (catches races)
    if _kill_pids_on_port(port, signal.SIGKILL):
        await asyncio.sleep(0.3)

    # Spawn ttyd
    proc = await asyncio.create_subprocess_exec(
        "ttyd",
        "-W",
        "-m",
        "3",
        "-p",
        str(port),
        "tmux",
        "attach",
        "-t",
        session_name,
        stdout=asyncio.subprocess.DEVNULL,
        stderr=asyncio.subprocess.DEVNULL,
        start_new_session=True,
    )

    # Write per-port PID file
    TTYD_PID_DIR.mkdir(parents=True, exist_ok=True)
    _pid_path_for_port(port).write_text(str(proc.pid))

    # Add to pool
    _pool[session_name] = PoolEntry(process=proc, port=port)

    # Poll until listening (up to ~1 s)
    for _ in range(20):
        if _ttyd_is_listening(port):
            break
        await asyncio.sleep(0.05)

    return port


def pool_port(session_name: str) -> int | None:
    """Look up the port for *session_name* without spawning.

    Returns the port if the session has a live pool entry, else None.
    """
    entry = _pool.get(session_name)
    if entry is not None and _is_alive(entry):
        return entry.port
    return None


async def kill_session(session_name: str) -> bool:
    """Kill the ttyd process for *session_name* and remove from pool.

    Returns True if a pool entry was found and cleaned up, False otherwise.
    """
    entry = _pool.pop(session_name, None)
    if entry is None:
        return False
    try:
        entry.process.terminate()
        try:
            await asyncio.wait_for(entry.process.wait(), timeout=2.0)
        except asyncio.TimeoutError:
            entry.process.kill()
    except ProcessLookupError:
        pass
    _pid_path_for_port(entry.port).unlink(missing_ok=True)
    return True


async def kill_all() -> None:
    """Kill all pooled ttyd processes and clear the pool."""
    for name in list(_pool):
        await kill_session(name)


async def kill_orphans() -> None:
    """Startup cleanup: read PID files, SIGTERM, scan port range for stragglers.

    Handles both legacy ``ttyd.pid`` and per-port ``ttyd-{port}.pid`` files.
    """
    await kill_orphan_ttyd()


# ---------------------------------------------------------------------------
# Backward-compat public API
# ---------------------------------------------------------------------------


async def kill_ttyd() -> bool:
    """Kill the running ttyd process and clean up the PID file.

    Belt-and-suspenders strategy:

    Strategy 1 — PID file:
        Reads the PID from TTYD_PID_PATH.  If no PID file exists, returns False.
        If the file content is not a valid integer, removes the file and returns
        False.  Checks whether the process is alive via ``os.kill(pid, 0)``.  If
        already gone (ProcessLookupError), cleans up and proceeds.  Otherwise
        sends SIGTERM and polls every 0.1 s for up to 2 s.

    Strategy 2 — port-based fallback:
        After the PID-file kill, finds and kills any process still listening on
        TTYD_PORT via ``_pids_on_port()`` (lsof → fuser → ss).  This catches
        orphaned ttyd processes whose PID was never recorded in the file (e.g.
        after a coordinator crash).  A brief 0.3 s wait is added to let the OS
        release the port.

    The PID file and ``_active_process`` are cleared in all cases before
    returning.

    Returns:
        True  — a process was killed (or was already dead) via either strategy.
        False — no PID file found and no process was listening on the port.
    """
    global _active_process

    killed = False

    # -------------------------------------------------------------------
    # Strategy 1: PID file
    # -------------------------------------------------------------------
    if TTYD_PID_PATH.exists():
        try:
            pid = int(TTYD_PID_PATH.read_text().strip())
        except ValueError:
            TTYD_PID_PATH.unlink(missing_ok=True)
            pid = None
        else:
            # Check whether the process is still alive.
            try:
                os.kill(pid, 0)
            except ProcessLookupError:
                # Already dead — clean up and note success.
                TTYD_PID_PATH.unlink(missing_ok=True)
                killed = True
                pid = None
            else:
                # Process is alive — ask it to terminate.
                os.kill(pid, signal.SIGTERM)

                # Poll up to 2 s for the process to exit.
                deadline = time.time() + 2.0
                while time.time() < deadline:
                    try:
                        os.kill(pid, 0)
                    except (ProcessLookupError, PermissionError):
                        break
                    await asyncio.sleep(0.1)

                TTYD_PID_PATH.unlink(missing_ok=True)
                killed = True
                pid = None  # noqa: F841 (intentional)

    # -------------------------------------------------------------------
    # Strategy 2: port-based fallback — catch orphans not in PID file
    # -------------------------------------------------------------------
    if _kill_pids_on_port(TTYD_PORT, signal.SIGTERM):
        killed = True
        # Brief pause so the OS can release the port before the next spawn.
        await asyncio.sleep(0.3)

    _active_process = None
    return killed


async def kill_orphan_ttyd() -> bool:
    """Kill any orphaned ttyd process left over from a previous coordinator run.

    Handles both legacy ``ttyd.pid`` and pool-era ``ttyd-{port}.pid`` files,
    plus scans the full port range for stragglers.

    Returns:
        True  — an orphan was found (process was dead or alive).
        False — no PID file found, or PID file contained invalid content.
    """
    # Pool-era cleanup: per-port PID files
    if TTYD_PID_DIR.exists():
        for pid_file in TTYD_PID_DIR.glob("ttyd-*.pid"):
            try:
                pid = int(pid_file.read_text().strip())
                os.kill(pid, signal.SIGTERM)
            except (ValueError, ProcessLookupError, PermissionError):
                pass
            pid_file.unlink(missing_ok=True)

    # Scan additional ports beyond TTYD_PORT
    for port in range(_PORT_MIN + 1, _PORT_MAX + 1):
        _kill_pids_on_port(port, signal.SIGTERM)

    # Original behavior: legacy PID file + TTYD_PORT cleanup
    return await kill_ttyd()


async def spawn_ttyd(session_name: str) -> asyncio.subprocess.Process:
    """Spawn a ttyd process attached to *session_name* via ``tmux attach``.

    Runs::

        ttyd -W -m 3 -p 7682 tmux attach -t <session_name>

    Before spawning, verifies that TTYD_PORT is free.  If any process is still
    listening on the port (e.g. a race between kill_ttyd() and spawn_ttyd()),
    it sends SIGKILL to force-free the port immediately.

    stdout and stderr are discarded (DEVNULL).  The PID is written to
    TTYD_PID_PATH.  The process handle is stored in ``_active_process``.

    Args:
        session_name: The tmux session name to attach to.

    Returns:
        The asyncio.subprocess.Process object for the spawned ttyd.
    """
    global _active_process

    # Final port-free guard — catches races where kill_ttyd() returned but
    # the old ttyd hasn't fully released the socket yet.
    if _kill_pids_on_port(TTYD_PORT, signal.SIGKILL):
        await asyncio.sleep(0.3)

    proc = await asyncio.create_subprocess_exec(
        "ttyd",
        "-W",
        "-m",
        "3",
        "-p",
        str(TTYD_PORT),
        "tmux",
        "attach",
        "-t",
        session_name,
        stdout=asyncio.subprocess.DEVNULL,
        stderr=asyncio.subprocess.DEVNULL,
        start_new_session=True,  # detach from parent process group so ttyd survives independently
    )

    # Write PID file (create parent dirs if needed)
    TTYD_PID_DIR.mkdir(parents=True, exist_ok=True)
    TTYD_PID_PATH.write_text(str(proc.pid))

    _active_process = proc
    return proc
