iLSP/ilsp/server.py

"""
iLSP main entrypoint.

Starts all services in one container:
  - pylsp        on localhost:2087  (internal only — Python LSP + i80 completions)
  - Bicep proxy  on localhost:2088  (internal only — Bicep LS wrapper + LRU modules)
  - aiohttp HTTP on $HTTP_PORT      (exposed — health, reload, WebSocket proxies)

Only the HTTP port is exposed to Traefik/Nomad. Editors connect via:
  ws://ilsp.i80.dk/python   →  pylsp
  ws://ilsp.i80.dk/bicep    →  Bicep LS
"""

import asyncio
import logging
import os
import signal
import subprocess
import sys
import threading

from aiohttp import web

from .python_lsp.catalog import PypiCatalog
from .bicep_lsp.modules import BicepModuleCatalog
from .bicep_lsp.proxy import serve_bicep
from .yaml_lsp.catalog import PipelineTemplateCatalog
from .yaml_lsp.proxy import yaml_ws_handler

logger = logging.getLogger(__name__)

PYTHON_LSP_PORT = int(os.getenv("PYTHON_LSP_PORT", "2087"))
BICEP_LSP_PORT  = int(os.getenv("BICEP_LSP_PORT",  "2088"))
HTTP_PORT       = int(os.getenv("HTTP_PORT", "8000"))

_CHUNK = 65536


def _serve_python_lsp(port: int) -> None:
    """Start pylsp in TCP mode on localhost only; restart on unexpected exit."""
    while True:
        proc = subprocess.Popen(
            [sys.executable, "-m", "pylsp", "--tcp", "--host", "127.0.0.1", "--port", str(port)],
        )
        logger.info("pylsp listening on localhost:%d  PID=%d", port, proc.pid)
        proc.wait()
        logger.warning("pylsp exited (code %s) — restarting", proc.returncode)


async def _pipe(reader: asyncio.StreamReader, writer) -> None:
    """Pipe bytes from reader to writer until EOF."""
    try:
        while True:
            data = await reader.read(_CHUNK)
            if not data:
                break
            writer.write(data)
            if hasattr(writer, "drain"):
                await writer.drain()
            else:
                await writer.send_bytes(data)
    except Exception:
        pass


async def _ws_proxy(request: web.Request, host: str, port: int) -> web.WebSocketResponse:
    """Bridge an incoming WebSocket connection to an internal TCP LSP server."""
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    try:
        tcp_reader, tcp_writer = await asyncio.wait_for(
            asyncio.open_connection(host, port), timeout=3.0
        )
    except (OSError, asyncio.TimeoutError) as exc:
        logger.error("Cannot connect to LSP on %s:%d — %s", host, port, exc)
        await ws.close(code=1011, message=b"LSP backend unavailable", timeout=2.0)
        return ws

    async def ws_to_tcp():
        try:
            async for msg in ws:
                if msg.type == web.WSMsgType.BINARY:
                    tcp_writer.write(msg.data)
                    await tcp_writer.drain()
                elif msg.type == web.WSMsgType.TEXT:
                    tcp_writer.write(msg.data.encode())
                    await tcp_writer.drain()
        except Exception:
            pass
        finally:
            tcp_writer.close()

    async def tcp_to_ws():
        try:
            while True:
                data = await tcp_reader.read(_CHUNK)
                if not data:
                    break
                await ws.send_bytes(data)
        except Exception:
            pass
        finally:
            await ws.close()

    await asyncio.gather(ws_to_tcp(), tcp_to_ws(), return_exceptions=True)
    return ws


async def _build_app() -> web.Application:
    app = web.Application()

    async def health(_: web.Request) -> web.Response:
        import shutil
        return web.json_response({
            "status": "ok",
            "pypi_packages": len(PypiCatalog._packages),
            "bicep_modules": len(BicepModuleCatalog._modules),
            "iac_source_modules": len(BicepModuleCatalog._iac),
            "yaml_lsp": bool(shutil.which("yaml-language-server")),
            "pipeline_templates": PipelineTemplateCatalog.template_count(),
        })

    async def reload(_: web.Request) -> web.Response:
        before_bicep = len(BicepModuleCatalog._modules)
        before_tmpl = PipelineTemplateCatalog.template_count()
        BicepModuleCatalog.load()
        PipelineTemplateCatalog.load()
        after_bicep = len(BicepModuleCatalog._modules)
        after_tmpl = PipelineTemplateCatalog.template_count()
        logger.info(
            "Catalog reloaded: bicep %d→%d  templates %d→%d",
            before_bicep, after_bicep, before_tmpl, after_tmpl,
        )
        return web.json_response({
            "status": "reloaded",
            "bicep_modules_before": before_bicep,
            "bicep_modules_after": after_bicep,
            "pipeline_templates_before": before_tmpl,
            "pipeline_templates_after": after_tmpl,
        })

    async def python_ws(request: web.Request) -> web.WebSocketResponse:
        return await _ws_proxy(request, "127.0.0.1", PYTHON_LSP_PORT)

    async def bicep_ws(request: web.Request) -> web.WebSocketResponse:
        return await _ws_proxy(request, "127.0.0.1", BICEP_LSP_PORT)

    async def yaml_ws(request: web.Request) -> web.WebSocketResponse:
        return await yaml_ws_handler(request)

    app.router.add_get("/health", health)
    app.router.add_post("/reload", reload)
    app.router.add_get("/python", python_ws)
    app.router.add_get("/bicep", bicep_ws)
    app.router.add_get("/yaml", yaml_ws)
    return app


async def main_async() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    logger.info("Pre-warming catalogs…")
    BicepModuleCatalog.load()
    PipelineTemplateCatalog.load()
    await PypiCatalog.start_background_refresh()

    # LSP servers run internally on localhost — not exposed outside the container
    threading.Thread(target=_serve_python_lsp, args=(PYTHON_LSP_PORT,), daemon=True).start()
    threading.Thread(target=serve_bicep, args=(BICEP_LSP_PORT,), daemon=True).start()
    # yaml-language-server is spawned per-connection in yaml_ws_handler (--stdio mode)

    app = await _build_app()
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "0.0.0.0", HTTP_PORT)
    await site.start()
    logger.info(
        "iLSP HTTP on http://0.0.0.0:%d  (ws://.../python  ws://.../bicep  ws://.../yaml)",
        HTTP_PORT,
    )

    await asyncio.Event().wait()


def main() -> None:
    loop = asyncio.new_event_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, loop.stop)
    loop.run_until_complete(main_async())