#!/lsiopy/bin/python3
"""Lazy TCP gate for the Headroom HTTP proxy.

The gate remains resident on the public Headroom port while the heavy Headroom
proxy (and its optional ONNX Kompress model) is started only for real traffic.
After an idle period the backend process is terminated, releasing its Python,
ONNX and model allocations. The next connection starts a fresh backend.

Only the Python standard library is imported here deliberately: the idle path
must not import Headroom, ONNX Runtime, transformers or the MCP SDK.
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import os
import signal
import socket
import sys
import time
from pathlib import Path


def _int_env(name: str, default: int, minimum: int = 1) -> int:
    try:
        value = int(os.environ.get(name, str(default)))
    except (TypeError, ValueError):
        value = default
    return max(minimum, value)


GATE_HOST = os.environ.get("HEADROOM_GATE_HOST", "127.0.0.1")
GATE_PORT = _int_env("HEADROOM_GATE_PORT", 8787)
BACKEND_HOST = os.environ.get("HEADROOM_BACKEND_HOST", "127.0.0.1")
BACKEND_PORT = _int_env("HEADROOM_BACKEND_PORT", 8789)
HEADROOM_BIN = os.environ.get("HEADROOM_REAL_BIN", "/usr/bin/headroom")
IDLE_TIMEOUT = _int_env("HEADROOM_IDLE_TIMEOUT_SECONDS", 900)
START_TIMEOUT = _int_env("HEADROOM_START_TIMEOUT_SECONDS", 60)
STOP_TIMEOUT = _int_env("HEADROOM_STOP_TIMEOUT_SECONDS", 10)
HF_HOME = os.environ.get("HF_HOME", str(Path.home() / ".headroom/hf"))
MAX_HEADER_BYTES = 128 * 1024
COPY_CHUNK = 64 * 1024


def log(message: str) -> None:
    stamp = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{stamp}] headroom-gate: {message}", file=sys.stderr, flush=True)


class BackendUnavailable(RuntimeError):
    pass


class ProxyGate:
    def __init__(self) -> None:
        self._lock = asyncio.Lock()
        self._process: asyncio.subprocess.Process | None = None
        self._last_activity = time.monotonic()
        self._active_connections = 0
        self._stopping = False
        self._watch_task: asyncio.Task[None] | None = None

    @property
    def state(self) -> str:
        process = self._process
        if process is None:
            return "dormant"
        if process.returncode is None:
            return "running"
        return "stopped"

    def touch(self) -> None:
        self._last_activity = time.monotonic()

    async def _backend_healthy(self) -> bool:
        for path in ("/livez", "/health"):
            try:
                reader, writer = await asyncio.wait_for(
                    asyncio.open_connection(BACKEND_HOST, BACKEND_PORT), timeout=1.0
                )
                request = (
                    f"GET {path} HTTP/1.1\r\n"
                    f"Host: {BACKEND_HOST}:{BACKEND_PORT}\r\n"
                    "Connection: close\r\n\r\n"
                ).encode()
                writer.write(request)
                await writer.drain()
                line = await asyncio.wait_for(reader.readline(), timeout=1.0)
                writer.close()
                with contextlib.suppress(Exception):
                    await writer.wait_closed()
                if line.startswith(b"HTTP/") and b" 2" in line[:16]:
                    return True
            except (OSError, asyncio.TimeoutError):
                continue
        return False

    async def ensure_backend(self) -> None:
        if await self._backend_healthy():
            self.touch()
            return

        async with self._lock:
            if await self._backend_healthy():
                self.touch()
                return

            process = self._process
            if process is not None and process.returncode is None:
                await self._wait_until_healthy()
                self.touch()
                return

            if not os.path.isfile(HEADROOM_BIN) or not os.access(HEADROOM_BIN, os.X_OK):
                raise BackendUnavailable(f"Headroom executable is unavailable: {HEADROOM_BIN}")

            Path(HF_HOME).mkdir(parents=True, exist_ok=True)

            environment = os.environ.copy()
            environment.update(
                {
                    "HF_HOME": HF_HOME,
                    "HEADROOM_PROXY_GATE": "1",
                }
            )
            command = [
                HEADROOM_BIN,
                "proxy",
                "--host",
                BACKEND_HOST,
                "--port",
                str(BACKEND_PORT),
                "--code-aware",
            ]
            log(f"starting backend on {BACKEND_HOST}:{BACKEND_PORT}")
            try:
                self._process = await asyncio.create_subprocess_exec(
                    *command,
                    env=environment,
                    start_new_session=True,
                )
            except Exception as exc:
                self._process = None
                raise BackendUnavailable(f"unable to start Headroom: {exc}") from exc

            self._watch_task = asyncio.create_task(self._watch_backend(self._process))
            try:
                await self._wait_until_healthy()
            except Exception:
                await self._terminate_backend_locked("startup failure")
                raise
            self.touch()

    async def _wait_until_healthy(self) -> None:
        deadline = time.monotonic() + START_TIMEOUT
        while time.monotonic() < deadline:
            process = self._process
            if process is not None and process.returncode is not None:
                raise BackendUnavailable(
                    f"Headroom exited during startup with status {process.returncode}; "
                    "see the add-on log"
                )
            if await self._backend_healthy():
                log("backend is ready")
                return
            await asyncio.sleep(0.25)
        raise BackendUnavailable(
            f"Headroom did not become ready within {START_TIMEOUT}s; see the add-on log"
        )

    async def _watch_backend(self, process: asyncio.subprocess.Process) -> None:
        returncode = await process.wait()
        async with self._lock:
            was_current = self._process is process
            if was_current:
                self._process = None
        if not self._stopping and was_current:
            log(f"backend exited with status {returncode}")

    async def stop_backend(self, reason: str) -> None:
        async with self._lock:
            await self._terminate_backend_locked(reason)

    async def _terminate_backend_locked(self, reason: str) -> None:
        process = self._process
        if process is None:
            return
        if process.returncode is not None:
            self._process = None
            return

        log(f"stopping backend ({reason})")
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
        try:
            await asyncio.wait_for(process.wait(), timeout=STOP_TIMEOUT)
        except asyncio.TimeoutError:
            log("backend did not stop after SIGTERM; sending SIGKILL")
            with contextlib.suppress(ProcessLookupError):
                os.killpg(process.pid, signal.SIGKILL)
            with contextlib.suppress(Exception):
                await process.wait()
        self._process = None

    async def idle_monitor(self) -> None:
        interval = max(5, min(30, IDLE_TIMEOUT // 4))
        while not self._stopping:
            await asyncio.sleep(interval)
            process = self._process
            idle_for = time.monotonic() - self._last_activity
            if (
                process is not None
                and process.returncode is None
                and idle_for >= IDLE_TIMEOUT
            ):
                await self.stop_backend(f"idle for {int(idle_for)}s")

    def status_payload(self) -> bytes:
        process = self._process
        payload = {
            "status": self.state,
            "backend_pid": process.pid if process is not None and process.returncode is None else None,
            "active_connections": self._active_connections,
            "idle_seconds": round(time.monotonic() - self._last_activity, 1),
            "idle_timeout_seconds": IDLE_TIMEOUT,
            "backend": f"{BACKEND_HOST}:{BACKEND_PORT}",
        }
        body = json.dumps(payload, separators=(",", ":")).encode()
        return (
            b"HTTP/1.1 200 OK\r\n"
            b"Content-Type: application/json\r\n"
            + f"Content-Length: {len(body)}\r\n".encode()
            + b"Connection: close\r\n\r\n"
            + body
        )

    async def handle_client(
        self,
        client_reader: asyncio.StreamReader,
        client_writer: asyncio.StreamWriter,
    ) -> None:
        self._active_connections += 1
        real_traffic = False
        peer = client_writer.get_extra_info("peername")
        try:
            try:
                initial = await asyncio.wait_for(
                    client_reader.readuntil(b"\r\n\r\n"), timeout=15.0
                )
            except (
                asyncio.IncompleteReadError,
                asyncio.LimitOverrunError,
                asyncio.TimeoutError,
            ):
                return
            if len(initial) > MAX_HEADER_BYTES:
                await self._send_error(client_writer, 431, "request headers too large")
                return

            first_line = initial.split(b"\r\n", 1)[0]
            parts = first_line.split(b" ")
            target = parts[1].split(b"?", 1)[0] if len(parts) >= 2 else b""
            if target == b"/gate/status":
                client_writer.write(self.status_payload())
                await client_writer.drain()
                return

            real_traffic = True
            self.touch()
            try:
                await self.ensure_backend()
                backend_reader, backend_writer = await asyncio.wait_for(
                    asyncio.open_connection(BACKEND_HOST, BACKEND_PORT), timeout=5.0
                )
            except (BackendUnavailable, OSError, asyncio.TimeoutError) as exc:
                log(f"backend unavailable for {peer}: {exc}")
                await self._send_error(client_writer, 503, str(exc))
                return

            backend_writer.write(initial)
            await backend_writer.drain()
            self.touch()

            upstream = asyncio.create_task(self._pipe(client_reader, backend_writer))
            downstream = asyncio.create_task(self._pipe(backend_reader, client_writer))
            done, pending = await asyncio.wait(
                {upstream, downstream}, return_when=asyncio.FIRST_COMPLETED
            )
            for task in pending:
                task.cancel()
            for task in done | pending:
                with contextlib.suppress(
                    asyncio.CancelledError,
                    ConnectionError,
                    OSError,
                ):
                    await task
            backend_writer.close()
            with contextlib.suppress(Exception):
                await backend_writer.wait_closed()
        finally:
            self._active_connections = max(0, self._active_connections - 1)
            if real_traffic:
                self.touch()
            client_writer.close()
            with contextlib.suppress(Exception):
                await client_writer.wait_closed()

    async def _pipe(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        while True:
            data = await reader.read(COPY_CHUNK)
            if not data:
                with contextlib.suppress(Exception):
                    writer.write_eof()
                return
            writer.write(data)
            await writer.drain()
            self.touch()

    @staticmethod
    async def _send_error(
        writer: asyncio.StreamWriter,
        status: int,
        detail: str,
    ) -> None:
        reason = (
            "Service Unavailable"
            if status == 503
            else "Request Header Fields Too Large"
        )
        body = json.dumps({"error": detail}).encode()
        response = (
            f"HTTP/1.1 {status} {reason}\r\n"
            "Content-Type: application/json\r\n"
            f"Content-Length: {len(body)}\r\n"
            "Connection: close\r\n\r\n"
        ).encode() + body
        writer.write(response)
        with contextlib.suppress(Exception):
            await writer.drain()

    async def shutdown(self) -> None:
        self._stopping = True
        await self.stop_backend("gate shutdown")


async def async_main() -> None:
    gate = ProxyGate()
    loop = asyncio.get_running_loop()
    stop_event = asyncio.Event()
    for sig in (signal.SIGTERM, signal.SIGINT):
        with contextlib.suppress(NotImplementedError):
            loop.add_signal_handler(sig, stop_event.set)

    server = await asyncio.start_server(
        gate.handle_client,
        GATE_HOST,
        GATE_PORT,
        limit=MAX_HEADER_BYTES + 1,
        family=socket.AF_INET,
    )
    addresses = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
    log(
        f"listening on {addresses}; backend is lazy and stops after {IDLE_TIMEOUT}s idle; "
        "status endpoint: /gate/status"
    )
    monitor = asyncio.create_task(gate.idle_monitor())
    async with server:
        await stop_event.wait()
    monitor.cancel()
    with contextlib.suppress(asyncio.CancelledError):
        await monitor
    await gate.shutdown()


if __name__ == "__main__":
    try:
        asyncio.run(async_main())
    except KeyboardInterrupt:
        pass