ForcePilot/backend/package/yuxi/channels/adapters/irc/probe.py

from __future__ import annotations

import asyncio
import ssl
import time
from collections.abc import Callable

from yuxi.channels.models import HealthStatus


async def ping_probe(
    send_fn: Callable[[str], None],
    read_fn,
    server: str,
    timeout: float = 8.0,
    metadata: dict | None = None,
) -> HealthStatus:
    start = time.monotonic()
    try:
        ping_token = f"fp-health-{int(time.time())}"
        send_fn(f"PING :{ping_token}")

        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            try:
                line = await asyncio.wait_for(read_fn(), timeout=deadline - time.monotonic())
            except TimeoutError:
                break
            if line and "PONG" in line:
                latency_ms = (time.monotonic() - start) * 1000
                return HealthStatus(
                    status="healthy",
                    latency_ms=latency_ms,
                    metadata=metadata or {},
                )

        return HealthStatus(
            status="degraded",
            last_error="PING no PONG response within timeout",
        )
    except (ConnectionError, OSError) as e:
        return HealthStatus(status="unhealthy", last_error=str(e))
    except Exception as e:
        return HealthStatus(status="unhealthy", last_error=str(e))


async def standalone_tcp_probe(
    host: str,
    port: int = 6697,
    use_tls: bool = True,
    nick: str = "ForcePilotProbe",
    username: str = "probe",
    realname: str = "ForcePilot Probe",
    timeout: float = 15.0,
) -> HealthStatus:
    """Independent TCP probe that connects, waits for 001, and disconnects."""
    start = time.monotonic()
    reader: asyncio.StreamReader | None = None
    writer: asyncio.StreamWriter | None = None

    try:
        if use_tls:
            ctx = ssl.create_default_context()
            ctx.check_hostname = True
            ctx.verify_mode = ssl.CERT_REQUIRED
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(
                    host=host,
                    port=port,
                    ssl=ctx,
                    server_hostname=host,
                ),
                timeout=timeout,
            )
        else:
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(host=host, port=port),
                timeout=timeout,
            )

        caps = ["server-time"]
        writer.write(b"CAP LS 302\r\n")
        writer.write(f"CAP REQ :{' '.join(caps)}\r\n".encode())
        writer.write(f"NICK {nick}\r\n".encode())
        writer.write(f"USER {username} 0 * :{realname}\r\n".encode())
        await writer.drain()

        registered = False
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            line = await asyncio.wait_for(reader.readline(), timeout=deadline - time.monotonic())
            if not line:
                break
            decoded = line.decode("utf-8", errors="replace").rstrip("\r\n")
            parts = decoded.split()
            if len(parts) >= 2 and parts[1] in ("001", "376", "422"):
                registered = True
                break
            if "ERROR" in decoded:
                return HealthStatus(
                    status="unhealthy",
                    latency_ms=(time.monotonic() - start) * 1000,
                    metadata={"host": host, "port": port, "tls": use_tls},
                    last_error=f"Server returned: {decoded}",
                )

        if not registered:
            return HealthStatus(
                status="degraded",
                latency_ms=(time.monotonic() - start) * 1000,
                metadata={"host": host, "port": port, "tls": use_tls},
                last_error="Registration not completed within timeout",
            )

        latency_ms = (time.monotonic() - start) * 1000
        try:
            writer.write(b"QUIT :ForcePilot probe done\r\n")
            await writer.drain()
        except Exception:
            pass

        return HealthStatus(
            status="healthy",
            latency_ms=latency_ms,
            metadata={"host": host, "port": port, "tls": use_tls, "probe_type": "standalone_tcp"},
        )

    except TimeoutError:
        return HealthStatus(
            status="unhealthy",
            latency_ms=(time.monotonic() - start) * 1000,
            metadata={"host": host, "port": port, "tls": use_tls},
            last_error="Connection timed out",
        )
    except (ConnectionError, OSError) as e:
        return HealthStatus(
            status="unhealthy",
            latency_ms=(time.monotonic() - start) * 1000,
            metadata={"host": host, "port": port, "tls": use_tls},
            last_error=str(e),
        )
    except Exception as e:
        return HealthStatus(
            status="unhealthy",
            latency_ms=(time.monotonic() - start) * 1000,
            metadata={"host": host, "port": port, "tls": use_tls},
            last_error=str(e),
        )
    finally:
        if writer:
            try:
                writer.close()
            except Exception:
                pass