#!/usr/bin/env python3
"""
User-space allowlist TCP proxy for the 111 Ollama fallback.

Purpose:
  - Keep the real Ollama server bound to 127.0.0.1:11434.
  - Expose 192.168.0.111:11434 only to approved momo-pro clients.
  - Reject noisy LAN / VM clients without requiring sudo/pfctl.
"""

from __future__ import annotations

import asyncio
import ipaddress
import logging
import os
import signal
import sys
import time


LISTEN_HOST = os.getenv("OLLAMA111_PROXY_LISTEN_HOST", "192.168.0.111")
LISTEN_PORT = int(os.getenv("OLLAMA111_PROXY_LISTEN_PORT", "11434"))
TARGET_HOST = os.getenv("OLLAMA111_PROXY_TARGET_HOST", "127.0.0.1")
TARGET_PORT = int(os.getenv("OLLAMA111_PROXY_TARGET_PORT", "11434"))
ALLOWED_CIDRS = tuple(
    item.strip()
    for item in os.getenv(
        "OLLAMA111_PROXY_ALLOWED_CIDRS",
        "127.0.0.1/32,192.168.0.111/32,192.168.0.188/32",
    ).split(",")
    if item.strip()
)


def _allowed_networks() -> tuple[ipaddress._BaseNetwork, ...]:
    return tuple(ipaddress.ip_network(item, strict=False) for item in ALLOWED_CIDRS)


ALLOWED_NETWORKS = _allowed_networks()
REJECT_LOG_DEDUP_SEC = float(os.getenv("OLLAMA111_PROXY_REJECT_LOG_DEDUP_SEC", "60"))
_LAST_REJECT_LOG: dict[str, float] = {}


def _is_allowed(peer_ip: str) -> bool:
    try:
        ip = ipaddress.ip_address(peer_ip)
    except ValueError:
        return False
    return any(ip in network for network in ALLOWED_NETWORKS)


async def _pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    try:
        while True:
            data = await reader.read(65536)
            if not data:
                break
            writer.write(data)
            await writer.drain()
    except (ConnectionError, asyncio.CancelledError):
        pass
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass


async def _handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    peer = writer.get_extra_info("peername")
    peer_ip = peer[0] if peer else "unknown"
    if not _is_allowed(peer_ip):
        now = time.monotonic()
        last_log = _LAST_REJECT_LOG.get(peer_ip, 0.0)
        if now - last_log >= REJECT_LOG_DEDUP_SEC:
            logging.warning("reject peer=%s allowed=%s", peer_ip, ",".join(ALLOWED_CIDRS))
            _LAST_REJECT_LOG[peer_ip] = now
        writer.close()
        await writer.wait_closed()
        return

    try:
        target_reader, target_writer = await asyncio.open_connection(TARGET_HOST, TARGET_PORT)
    except Exception as exc:
        logging.error("target connect failed peer=%s target=%s:%s err=%r", peer_ip, TARGET_HOST, TARGET_PORT, exc)
        writer.close()
        await writer.wait_closed()
        return

    logging.info("proxy peer=%s -> %s:%s", peer_ip, TARGET_HOST, TARGET_PORT)
    await asyncio.gather(
        _pipe(reader, target_writer),
        _pipe(target_reader, writer),
    )


async def _main() -> None:
    logging.basicConfig(
        level=os.getenv("OLLAMA111_PROXY_LOG_LEVEL", "INFO"),
        format="%(asctime)s %(levelname)s %(message)s",
        stream=sys.stdout,
    )
    server = await asyncio.start_server(_handle_client, LISTEN_HOST, LISTEN_PORT)
    sockets = ", ".join(str(sock.getsockname()) for sock in (server.sockets or []))
    logging.info(
        "ollama111 allow proxy listening=%s target=%s:%s allowed=%s",
        sockets,
        TARGET_HOST,
        TARGET_PORT,
        ",".join(ALLOWED_CIDRS),
    )

    stop = asyncio.Event()
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, stop.set)

    async with server:
        await stop.wait()
        server.close()
        await server.wait_closed()


if __name__ == "__main__":
    asyncio.run(_main())