"""SSE streaming endpoint for token-by-token LLM responses."""

from __future__ import annotations

import asyncio
import json
import uuid
from typing import Any

from fastapi import APIRouter
from fastapi.responses import StreamingResponse

from fusionagi._logger import logger
from fusionagi.api.dependencies import get_orchestrator

router = APIRouter()


async def _sse_generator(session_id: str, prompt: str) -> Any:
    """Generate SSE events for a streaming prompt response."""
    event_id = str(uuid.uuid4())[:8]

    yield f"event: start\ndata: {json.dumps({'session_id': session_id, 'event_id': event_id})}\n\n"

    orch = get_orchestrator()
    if orch is None:
        yield f"event: error\ndata: {json.dumps({'error': 'Orchestrator not available'})}\n\n"
        return

    try:
        yield f"event: heads_running\ndata: {json.dumps({'heads': ['logic', 'creativity', 'research', 'safety']})}\n\n"

        from fusionagi.schemas.task import Task
        task = Task(task_id=f"stream_{event_id}", prompt=prompt)
        result = orch.run(task)

        if result and hasattr(result, "final_answer"):
            answer = result.final_answer or ""
            # Stream token-by-token (simulate chunked response)
            words = answer.split()
            for i, word in enumerate(words):
                chunk = word + (" " if i < len(words) - 1 else "")
                yield f"event: token\ndata: {json.dumps({'token': chunk, 'index': i})}\n\n"
                await asyncio.sleep(0.02)

            yield f"event: complete\ndata: {json.dumps({'session_id': session_id, 'full_text': answer})}\n\n"
        else:
            yield f"event: complete\ndata: {json.dumps({'session_id': session_id, 'full_text': ''})}\n\n"

    except Exception as e:
        logger.error("SSE streaming error", extra={"error": str(e), "session_id": session_id})
        yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"


@router.post("/sessions/{session_id}/stream/sse")
async def stream_sse(session_id: str, body: dict[str, Any]) -> StreamingResponse:
    """Stream a prompt response as Server-Sent Events.

    Events emitted:
    - ``start``: Stream began
    - ``heads_running``: Which heads are processing
    - ``token``: Individual response token
    - ``complete``: Final response with full text
    - ``error``: Error occurred
    """
    prompt = body.get("prompt", "")
    return StreamingResponse(
        _sse_generator(session_id, prompt),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )