from __future__ import annotations

import json
from typing import AsyncIterator, Optional

from fastapi import APIRouter, HTTPException, status
from fastapi.responses import StreamingResponse

from app.deps import AuthDep
from app.services.lightrag_service import lightrag_service
from app.services.supabase_rest import supabase_rest

router = APIRouter(prefix="/chat")


@router.get("")
async def chat(
    query: str,
    auth: AuthDep,
    document_id: Optional[str] = None,
) -> StreamingResponse:
    """
    基于 LightRAG 的 SSE 流式回答。
    query: 必填问题
    document_id: 可选，指定所属 workspace
    """
    if not query or not query.strip():
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="query 不能为空")

    workspace_id: Optional[str] = None
    if document_id:
        document = supabase_rest.select_one(
            "documents", {"id": document_id, "user_id": auth.user_id}
        )
        if not document:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found")
        workspace_id = (
            str(document.get("workspace_id")) if document.get("workspace_id") else None
        )

    health = await lightrag_service.run_healthcheck()
    if not health.get("ok"):
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail=health.get("error", "LightRAG 未就绪"),
        )

    lightrag_result = await lightrag_service.stream_answer(
        query_text=query, user_id=auth.user_id, workspace_id=workspace_id
    )

    async def event_stream() -> AsyncIterator[str]:
        if lightrag_result.get("is_streaming") and lightrag_result.get("iterator"):
            iterator = lightrag_result["iterator"]
            async for chunk in iterator:
                payload = json.dumps({"type": "chunk", "content": chunk})
                yield f"data: {payload}\n\n"
        else:
            payload = json.dumps(
                {"type": "chunk", "content": lightrag_result.get("content", "")}
            )
            yield f"data: {payload}\n\n"
        references = lightrag_result.get("references", [])
        yield f"data: {json.dumps({'type': 'references', 'data': references})}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")