#!/usr/bin/env python3
"""SSE streaming server: demonstrates server-sent events for LLM token streaming.

Run: python sse_server.py
Then: curl -N -X POST http://localhost:8000/chat -H "Content-Type: application/json" -d '{"messages": "hello"}'
"""
from __future__ import annotations

import asyncio
import json
import sys
import time

try:
    from fastapi import FastAPI
    from fastapi.responses import StreamingResponse
    import uvicorn
except ImportError:
    print("Install: pip install fastapi uvicorn")
    sys.exit(1)

app = FastAPI()


# Simulate an LLM streaming provider (replace with real provider: openai, vllm, etc.)
async def mock_llm_stream(messages: str) -> str:
    """Simulated LLM that yields tokens one at a time with realistic delays."""
    response = f"Echo: {messages[:40]}. This is a streaming response with multiple tokens."
    words = response.split()
    for i, word in enumerate(words):
        await asyncio.sleep(0.05)  # simulate per-token latency
        yield word + (" " if i < len(words) - 1 else "")


async def event_stream(messages: str):
    """SSE event generator: streams tokens as delta events, finishes with done."""
    try:
        token_buffer = []
        async for token in mock_llm_stream(messages):
            token_buffer.append(token)
            # Batch small tokens to reduce event count without hurting perceived latency
            if len(token_buffer) >= 3 or token.endswith(" "):
                delta = "".join(token_buffer)
                token_buffer.clear()
                yield f"event: delta\ndata: {json.dumps({'delta': delta})}\n\n"

        # Flush remaining tokens
        if token_buffer:
            delta = "".join(token_buffer)
            yield f"event: delta\ndata: {json.dumps({'delta': delta})}\n\n"

        yield json.dumps({"event": "done", "data": {"status": "ok", "usage": {"completion_tokens": 24}}})
        # SSE format for done event
        yield f"event: done\ndata: {json.dumps({'status': 'ok'})}\n\n"

    except asyncio.CancelledError:
        # Client disconnected: cancel upstream LLM call
        print("[server] Client disconnected — cancelling upstream LLM call")
        yield f"event: error\ndata: {json.dumps({'error': 'cancelled'})}\n\n"


@app.post("/chat")
async def chat(req: dict):
    """POST endpoint returning SSE stream."""
    messages = req.get("messages", "")
    return StreamingResponse(
        event_stream(messages),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",   # bypass nginx buffering
            "Connection": "keep-alive",
        },
    )


@app.get("/health")
async def health():
    return {"status": "ok"}


def main() -> None:
    print("SSE streaming server on http://localhost:8000")
    print("Test: curl -N -X POST http://localhost:8000/chat \\")
    print('     -H "Content-Type: application/json" -d \'{"messages":"hello"}\'')
    uvicorn.run(app, host="127.0.0.1", port=8000, log_level="warning")


if __name__ == "__main__":
    main()