"""SSE endpoint: sources event first, then token events, then done. LLM is simulated."""
import asyncio, json, time
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.middleware.gzip import GZipMiddleware

RETRIEVAL_S = 0.050   # simulated vector search
LLM_TTFT_S = 0.400    # simulated time to first token
TOKEN_S = 0.020       # simulated inter-token delay
N_TOKENS = 60
SOURCES = [{"id": i, "title": f"doc-{i}", "score": round(0.8 - i * 0.05, 2)} for i in range(5)]

def sse(event: str, data) -> bytes:
    return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()

async def llm_tokens(blocking=False):
    if blocking:
        time.sleep(LLM_TTFT_S)          # bug: blocks the event loop
    else:
        await asyncio.sleep(LLM_TTFT_S)
    for i in range(N_TOKENS):
        yield f"tok{i} "
        if blocking:
            time.sleep(TOKEN_S)
        else:
            await asyncio.sleep(TOKEN_S)

def make_app(gzip=False):
    app = FastAPI()
    if gzip:
        app.add_middleware(GZipMiddleware, minimum_size=1)

    @app.get("/plain")
    async def plain():
        await asyncio.sleep(RETRIEVAL_S)
        text = "".join([t async for t in llm_tokens()])
        return {"sources": SOURCES, "answer": text}

    async def stream(blocking=False):
        await asyncio.sleep(RETRIEVAL_S)
        yield sse("sources", SOURCES)
        async for t in llm_tokens(blocking):
            yield sse("token", t)
        yield sse("done", {})

    @app.get("/stream")
    async def stream_ep():
        return StreamingResponse(stream(), media_type="text/event-stream",
                                 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})

    @app.get("/stream-blocking")
    async def stream_blocking():
        return StreamingResponse(stream(True), media_type="text/event-stream")
    return app