#!/usr/bin/env python3
"""
LSP WebSocket bridge: stdin/stdout (LSP Content-Length framed) <-> WebSocket.

websocat reads stdin line-by-line and holds the JSON body until a trailing \\n
arrives -- causing a deadlock on the first LSP message. This script solves that
by reading complete Content-Length-framed messages from stdin before sending.

Usage:
    lsp_bridge.py <wss://host/path>

The bridge is transparent: it reads exactly one LSP message from stdin (header +
body), sends it as a single binary WebSocket frame, and writes any incoming
frames straight to stdout.
"""

import asyncio
import ssl
import sys
import traceback

import websockets


def _parse_header(buf: bytes) -> tuple[int, int]:
    """Return (content_length, body_start_offset) or (-1, -1) if header incomplete."""
    end = buf.find(b"\r\n\r\n")
    if end == -1:
        return -1, -1
    content_length = 0
    for line in buf[:end].split(b"\r\n"):
        if line.lower().startswith(b"content-length:"):
            content_length = int(line.split(b":")[1].strip())
    return content_length, end + 4


async def _read_stdin_lsp(reader: asyncio.StreamReader) -> bytes | None:
    """Read one complete LSP-framed message from the async stdin reader."""
    buf = b""
    content_length = -1
    body_start = -1

    while True:
        chunk = await reader.read(4096)
        if not chunk:
            return None
        buf += chunk

        if content_length == -1:
            content_length, body_start = _parse_header(buf)

        if content_length != -1 and len(buf) >= body_start + content_length:
            return buf[: body_start + content_length]


async def main(uri: str) -> None:
    ssl_ctx = ssl.create_default_context() if uri.startswith("wss://") else None

    loop = asyncio.get_event_loop()

    stdin_reader = asyncio.StreamReader()
    await loop.connect_read_pipe(
        lambda: asyncio.StreamReaderProtocol(stdin_reader), sys.stdin.buffer
    )

    async with websockets.connect(uri, ssl=ssl_ctx) as ws:

        async def stdin_to_ws() -> None:
            try:
                while True:
                    msg = await _read_stdin_lsp(stdin_reader)
                    if msg is None:
                        break
                    await ws.send(msg)
            except Exception as e:
                print(f"stdin_to_ws error: {e}", file=sys.stderr)
                traceback.print_exc(file=sys.stderr)

        async def ws_to_stdout() -> None:
            try:
                async for frame in ws:
                    data = frame if isinstance(frame, bytes) else frame.encode()
                    sys.stdout.buffer.write(data)
                    sys.stdout.buffer.flush()
            except Exception as e:
                print(f"ws_to_stdout error: {e}", file=sys.stderr)
                traceback.print_exc(file=sys.stderr)

        await asyncio.gather(stdin_to_ws(), ws_to_stdout())


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: lsp_bridge.py <wss://...>", file=sys.stderr)
        sys.exit(1)
    try:
        asyncio.run(main(sys.argv[1]))
    except Exception as e:
        print(f"Fatal error: {e}", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)