import asyncio
from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()

# async generator — can use await inside
async def async_text_generator():
    words = ["Hello", " ", "async", " ", "world", "!"]
    for word in words:
        yield word
        await asyncio.sleep(0.2)  # non-blocking sleep

@app.get("/stream/async")
async def stream_async():
    return StreamingResponse(
        async_text_generator(),
        media_type="text/plain"
    )

# Real-world: streaming rows from a database
async def stream_db_rows(db_session):
    # SQLAlchemy async streaming query
    result = await db_session.stream(
        select(User).order_by(User.id)
    )
    async for row in result:
        # serialize each row to JSON line
        yield json.dumps(row._asdict()) + "\n"

@app.get("/stream/users")
async def stream_users(db: AsyncSession = Depends(get_db)):
    return StreamingResponse(
        stream_db_rows(db),
        media_type="application/x-ndjson"  # newline-delimited JSON
    )

# Real-world: streaming an HTTP response from another service
import httpx

async def proxy_stream(url: str):
    async with httpx.AsyncClient() as client:
        async with client.stream("GET", url) as response:
            async for chunk in response.aiter_bytes():
                yield chunk

@app.get("/proxy/download")
async def proxy_download(url: str):
    return StreamingResponse(
        proxy_stream(url),
        media_type="application/octet-stream"
    )

async def stream_with_retry():
    # Tell the browser to wait 3 seconds before reconnecting
    yield "retry: 3000\n\n"

    for i in range(5):
        yield f"id: {i}\ndata: Message {i}\n\n"
        await asyncio.sleep(1)

# Default retry is 3 seconds if not specified

import asyncio
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse

app = FastAPI()

async def stream_with_keepalive(request: Request, data_source):
    async def generate():
        # Interleave data events and keep-alive pings
        ping_task = asyncio.create_task(send_pings())
        data_task = asyncio.create_task(send_data(data_source))

        queue = asyncio.Queue()

        async def send_pings():
            while True:
                await asyncio.sleep(15)                   # every 15s
                await queue.put(": ping\n\n")           # comment line

        async def send_data(source):
            async for item in source:
                await queue.put(f"data: {item}\n\n")
            await queue.put(None)                         # sentinel: done

        while True:
            if await request.is_disconnected():           # client gone?
                ping_task.cancel()
                break
            item = await queue.get()
            if item is None:                            # done signal
                ping_task.cancel()
                break
            yield item

    return StreamingResponse(generate(), media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})

# Simple version: just ping periodically
async def simple_keepalive_stream(request: Request):
    counter = 0
    while not await request.is_disconnected():
        counter += 1
        yield f"data: update {counter}\n\n"
        try:
            await asyncio.wait_for(asyncio.sleep(10), timeout=10)
        except asyncio.TimeoutError:
            yield ": ping\n\n"   # keep alive comment

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

# Lines starting with ":" are comments — browser ignores them
# but they keep the TCP connection alive through proxies
: ping

: heartbeat

: keep-alive

# Send one every 15-30 seconds

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from openai import AsyncOpenAI
from pydantic import BaseModel
import json

app = FastAPI()
client = AsyncOpenAI()

class ChatRequest(BaseModel):
    messages: list[dict]   # full conversation history
    model: str = "gpt-4o"

async def llm_sse_stream(request: Request, chat_req: ChatRequest):
    stream = await client.chat.completions.create(
        model=chat_req.model,
        messages=chat_req.messages,
        stream=True
    )
    async for chunk in stream:
        # Check if client disconnected
        if await request.is_disconnected():
            await stream.close()      # cancel LLM request
            break

        delta = chunk.choices[0].delta
        finish = chunk.choices[0].finish_reason

        if delta.content:
            # Send token as SSE event
            payload = json.dumps({"token": delta.content})
            yield f"event: token\ndata: {payload}\n\n"

        if finish == "stop":
            # Send done event with usage info
            usage = chunk.usage
            done_data = json.dumps({
                "finish_reason": "stop",
                "usage": usage.model_dump() if usage else {}
            })
            yield f"event: done\ndata: {done_data}\n\n"
        elif finish == "length":
            yield 'event: error\ndata: {"error":"max_tokens reached"}\n\n'

@app.post("/chat/stream")
async def chat_stream(request: Request, chat_req: ChatRequest):
    return StreamingResponse(
        llm_sse_stream(request, chat_req),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
    )

# JavaScript client
"""
// Using fetch with readable stream (more control than EventSource)
const response = await fetch("/chat/stream", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ messages: [{ role: "user", content: "Hello!" }] })
});

const reader = response.body.getReader();
const decoder = new TextDecoder();
let buffer = "";

while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    
    buffer += decoder.decode(value, { stream: true });
    const lines = buffer.split("\n\n");
    buffer = lines.pop();  // keep incomplete last line
    
    for (const chunk of lines) {
        if (chunk.startsWith("event: token")) {
            const dataLine = chunk.split("\n").find(l => l.startsWith("data:"));
            const { token } = JSON.parse(dataLine.slice(5));
            document.getElementById("output").textContent += token;
        }
    }
}
"""

let abortController = null;

async function startGeneration(prompt) {
    abortController = new AbortController();  // create abort signal

    const response = await fetch("/ai/generate", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ prompt }),
        signal: abortController.signal  // attach to fetch
    });

    // When user clicks Stop, abortController.abort() cancels the fetch
    // → browser closes the HTTP connection
    // → request.is_disconnected() returns True on server
    // → server cancels LLM stream
    const reader = response.body.getReader();
    try {
        while (true) {
            const { done, value } = await reader.read();
            if (done) break;
            // handle chunk...
        }
    } catch (e) {
        if (e.name === "AbortError") console.log("User stopped");
    }
}

// Stop button
document.getElementById("stop-btn").onclick = () => {
    if (abortController) abortController.abort();
};