#!/usr/bin/env python3
"""
oMLX compatibility proxy.

Routes reasoning_content → content in streaming deltas so consumers
that lack a native reasoning_content rendering path still see thinking output.

Gemma4 emits thinking tokens as reasoning_content in stream deltas.
This proxy moves them into content when no content field is present.
Tool-call deltas are forwarded unchanged — oMLX buffers the entire tool call
internally (via ToolCallStreamFilter) and emits one structured chunk at the
end; we must not touch it.

Streaming: reads line-by-line from upstream and awaits each write() so
tokens reach the client one-at-a-time without aiohttp's internal batching.
In aiohttp ≥ 3.9, await write() performs an implicit drain to TCP.

Runs on :8090, forwards to :8080.
"""
import json
from aiohttp import web, ClientSession, ClientTimeout, ClientPayloadError

UPSTREAM = "http://localhost:8080"
PROXY_PORT = 8090
# No total timeout — long streaming sessions must not be cut short.
# connect=10 is fine; sock_read=None lets the stream run indefinitely.
TIMEOUT = ClientTimeout(total=None, connect=10, sock_read=None)


def _route_reasoning(line: str) -> str:
    """
    For a single SSE data line, route reasoning_content → content when the
    delta carries no content and no tool_calls. Returns the (possibly rewritten)
    line. Non-data lines and [DONE] are returned unchanged.
    """
    if not line.startswith("data:") or line.strip() == "data: [DONE]":
        return line
    try:
        payload = json.loads(line[5:].strip())
        for choice in payload.get("choices", []):
            delta = choice.get("delta", {})
            # Never touch deltas that carry tool_calls — forward them intact.
            if not delta.get("tool_calls"):
                rc = delta.pop("reasoning_content", None)
                if rc and not delta.get("content"):
                    delta["content"] = rc
        return f"data: {json.dumps(payload)}"
    except (json.JSONDecodeError, KeyError):
        return line


async def _stream_response(response: web.StreamResponse, upstream_resp) -> None:
    """
    Read the upstream SSE stream line-by-line and forward each line immediately
    with an explicit drain() so the client sees tokens as they are generated.

    OSError (including ConnectionResetError) is caught and suppressed — it means
    the client disconnected, which is normal and should not crash the handler.
    """
    try:
        async for raw_line in upstream_resp.content:
            line = raw_line.decode(errors="replace").rstrip("\r\n")
            out = _route_reasoning(line) + "\n"
            # await write() is itself a flush in aiohttp ≥ 3.9 — each call
            # writes and drains to TCP immediately, ensuring per-token delivery.
            await response.write(out.encode())
    except OSError:
        # Client disconnected mid-stream — abort gracefully.
        pass
    except ClientPayloadError:
        # oMLX dropped the upstream connection mid-stream (e.g. hard memory
        # pressure burst). Send a clean [DONE] so the client gets a valid
        # SSE terminator rather than a hanging connection.
        try:
            await response.write(b"data: [DONE]\n\n")
        except OSError:
            pass


async def proxy_handler(request: web.Request) -> web.StreamResponse:
    path = request.path_qs
    headers = {k: v for k, v in request.headers.items()
               if k.lower() not in ("host", "content-length")}

    body = await request.read()

    # Detect streaming by parsing JSON — byte matching misses "stream": true (space)
    is_streaming_chat = False
    if request.method == "POST" and "/chat/completions" in request.path:
        try:
            is_streaming_chat = json.loads(body).get("stream", False) is True
        except (json.JSONDecodeError, AttributeError):
            pass

    if is_streaming_chat:
        async with ClientSession(timeout=TIMEOUT) as session:
            async with session.post(
                f"{UPSTREAM}{path}", headers=headers, data=body
            ) as upstream:
                sse_resp = web.StreamResponse(
                    status=upstream.status,
                    headers={
                        **{k: v for k, v in upstream.headers.items()
                           if k.lower() not in ("transfer-encoding",)},
                        "X-Accel-Buffering": "no",
                        "Cache-Control": "no-cache",
                    },
                )
                await sse_resp.prepare(request)
                await _stream_response(sse_resp, upstream)
                try:
                    await sse_resp.write_eof()
                except OSError:
                    pass
                return sse_resp

    # Generic passthrough
    async with ClientSession(timeout=TIMEOUT) as session:
        async with session.request(
            request.method, f"{UPSTREAM}{path}", headers=headers, data=body
        ) as upstream:
            data = await upstream.read()
            return web.Response(
                status=upstream.status,
                headers={k: v for k, v in upstream.headers.items()
                         if k.lower() not in ("transfer-encoding", "content-encoding")},
                body=data,
            )


app = web.Application()
app.router.add_route("*", "/{path_info:.*}", proxy_handler)

if __name__ == "__main__":
    print(
        f"oMLX proxy :8090 → {UPSTREAM}\n"
        f"  • reasoning_content routed to content (skips tool_calls deltas)\n"
        f"  • line-by-line streaming with await write() — no batching\n"
        f"  • OSError on client disconnect handled gracefully",
        flush=True,
    )
    web.run_app(app, host="127.0.0.1", port=PROXY_PORT, access_log=None)