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

from __future__ import annotations

import asyncio
import ssl
from collections.abc import Callable

from yuxi.utils.logging_config import logger

from .splits import calc_effective_limit, split_markdown, split_utf8

_FLOOD_BURST = 4
_FLOOD_DELAY = 0.8


def send_raw_line(
    send_fn: Callable[[str], None],
    line: str,
) -> None:
    send_fn(line)


async def send_privmsg(
    send_fn: Callable[[str], None],
    target: str,
    text: str,
    nick: str = "",
    username: str = "",
    server: str = "",
    reply_to_id: str = "",
    reply_to_nick: str = "",
    chunker_mode: str = "length",
) -> None:
    if reply_to_id:
        text = f"[reply:{reply_to_id}] {text}"
    elif reply_to_nick:
        text = f"<{reply_to_nick}> {text}"
    limit = calc_effective_limit(target, nick, username, server)
    if chunker_mode == "markdown":
        chunks = split_markdown(text, limit)
    else:
        chunks = split_utf8(text, limit)

    for i, chunk in enumerate(chunks):
        if i >= _FLOOD_BURST:
            await asyncio.sleep(_FLOOD_DELAY)
        send_raw_line(send_fn, f"PRIVMSG {target} :{chunk}")


async def send_media_text(
    send_fn: Callable[[str], None],
    target: str,
    text: str,
    nick: str = "",
    username: str = "",
    server: str = "",
) -> None:
    await send_privmsg(send_fn, target, text, nick, username, server)


def send_notice(
    send_fn: Callable[[str], None],
    target: str,
    text: str,
) -> None:
    send_raw_line(send_fn, f"NOTICE {target} :{text}")


def send_whois(
    send_fn: Callable[[str], None],
    nick: str,
) -> None:
    send_raw_line(send_fn, f"WHOIS {nick}")


class DeliveryTracker:
    def __init__(self, max_pending: int = 100):
        self._pending: dict[str, asyncio.Future] = {}
        self._max_pending = max_pending

    def track(self, message_id: str) -> asyncio.Future:
        if len(self._pending) >= self._max_pending:
            oldest = next(iter(self._pending))
            self._pending.pop(oldest, None)
        future: asyncio.Future = asyncio.get_running_loop().create_future()
        self._pending[message_id] = future
        return future

    def confirm(self, message_id: str) -> bool:
        future = self._pending.pop(message_id, None)
        if future and not future.done():
            future.set_result(True)
            return True
        return False

    def fail(self, message_id: str, error: str = "") -> bool:
        future = self._pending.pop(message_id, None)
        if future and not future.done():
            future.set_exception(ConnectionError(error or "delivery failed"))
            return True
        return False

    def pending_count(self) -> int:
        return len(self._pending)

    def clear(self) -> None:
        for future in self._pending.values():
            if not future.done():
                future.cancel()
        self._pending.clear()


async def send_message_irc(
    server: str,
    port: int,
    use_tls: bool,
    nick: str,
    target: str,
    text: str,
    username: str = "",
    realname: str = "",
    password: str = "",
    timeout: float = 12.0,
) -> bool:

    async def _read_line(reader: asyncio.StreamReader) -> str:
        try:
            line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            return line.decode("utf-8", errors="replace").rstrip("\r\n")
        except TimeoutError:
            return ""

    try:
        if use_tls:
            ssl_context = ssl.create_default_context()
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(server, port, ssl=ssl_context),
                timeout=timeout,
            )
        else:
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(server, port),
                timeout=timeout,
            )

        def send_fn(line: str) -> None:
            writer.write((line + "\r\n").encode("utf-8"))

        if password:
            send_fn(f"PASS {password}")

        effective_user = username or nick
        effective_realname = realname or nick
        send_fn(f"NICK {nick}")
        send_fn(f"USER {effective_user} 0 * :{effective_realname}")

        deadline = asyncio.get_event_loop().time() + min(timeout, 10.0)
        registered = False
        while asyncio.get_event_loop().time() < deadline:
            line = await _read_line(reader)
            if not line:
                continue
            parts = line.split()
            if len(parts) >= 2 and parts[1] in ("376", "422", "001", "002", "003", "004"):
                registered = True
                break
            if len(parts) >= 2 and parts[0] == "PING":
                send_fn(f"PONG {parts[1] if len(parts) > 1 else 'standalone'}")
            if len(parts) >= 2 and parts[1] == "433":
                logger.warning(f"send_message_irc: nick {nick} in use on {server}")
                writer.close()
                return False

        if not registered:
            writer.close()
            return False

        await asyncio.sleep(0.5)

        effective_nick = nick
        effective_user_name = username or nick
        limit = calc_effective_limit(target, effective_nick, effective_user_name, server)
        chunks = split_utf8(text, limit)

        for i, chunk in enumerate(chunks):
            if i >= _FLOOD_BURST:
                await asyncio.sleep(_FLOOD_DELAY)
            writer.write(f"PRIVMSG {target} :{chunk}\r\n".encode())

        await writer.drain()

        await asyncio.sleep(0.5)
        writer.write(b"QUIT :send_message_irc done\r\n")
        await writer.drain()

        await asyncio.sleep(0.3)
        writer.close()
        return True

    except TimeoutError:
        logger.warning(f"send_message_irc: timed out connecting to {server}:{port}")
        return False
    except Exception as e:
        logger.warning(f"send_message_irc: error connecting to {server}:{port}: {e}")
        return False