from typing import Annotated

from fastapi import APIRouter, Depends, Query, Request, Response, status
from fastapi.responses import StreamingResponse
from sqlalchemy.ext.asyncio import AsyncSession

from kbqa.api.errors import AppError
from kbqa.chat.repository import ChatRepository
from kbqa.chat.schemas import (
    MessagePage,
    MessageResponse,
    MessageSourceResponse,
    ChatStreamRequest,
    SessionCreate,
    SessionPage,
    SessionResponse,
)
from kbqa.database import get_session

router = APIRouter(prefix="/sessions", tags=["sessions"])
chat_router = APIRouter(prefix="/chat", tags=["chat"])
SessionDependency = Annotated[AsyncSession, Depends(get_session)]


@router.post("", response_model=SessionResponse, status_code=status.HTTP_201_CREATED)
async def create_session(payload: SessionCreate, session: SessionDependency) -> SessionResponse:
    created = await ChatRepository(session).create_session(payload.title)
    return SessionResponse.model_validate(created)


@router.get("", response_model=SessionPage)
async def list_sessions(
    session: SessionDependency,
    offset: Annotated[int, Query(ge=0)] = 0,
    limit: Annotated[int, Query(ge=1, le=100)] = 20,
) -> SessionPage:
    items, total = await ChatRepository(session).page_sessions(offset, limit)
    return SessionPage(
        items=[SessionResponse.model_validate(item) for item in items],
        total=total,
        offset=offset,
        limit=limit,
    )


@router.get("/{session_id}/messages", response_model=MessagePage)
async def list_messages(
    session_id: str,
    session: SessionDependency,
    offset: Annotated[int, Query(ge=0)] = 0,
    limit: Annotated[int, Query(ge=1, le=100)] = 50,
) -> MessagePage:
    repository = ChatRepository(session)
    if await repository.get_session(session_id) is None:
        raise AppError("SESSION_NOT_FOUND", "会话不存在", 404)
    messages, total = await repository.page_messages(session_id, offset, limit)
    sources_by_message = await repository.sources_for_messages([message.id for message in messages])
    responses: list[MessageResponse] = []
    for message in messages:
        sources = sources_by_message.get(message.id, [])
        source_labels = {source.citation_label for source in sources}
        deleted = sorted(repository.citation_labels(message.content) - source_labels)
        responses.append(
            MessageResponse(
                id=message.id,
                session_id=message.session_id,
                role=message.role,
                content=message.content,
                status=message.status,
                error_code=message.error_code,
                order_index=message.order_index,
                created_at=message.created_at,
                sources=[
                    MessageSourceResponse(
                        label=source.citation_label,
                        document_id=source.document_id,
                        chunk_id=source.chunk_id,
                        filename=source.filename,
                        order_index=source.order_index,
                        page_number=source.page_number,
                        content=source.content,
                        score=source.score,
                    )
                    for source in sources
                ],
                deleted_source_labels=deleted,
            )
        )
    return MessagePage(items=responses, total=total, offset=offset, limit=limit)


@router.delete("/{session_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_session(session_id: str, session: SessionDependency) -> Response:
    if not await ChatRepository(session).delete_session(session_id):
        raise AppError("SESSION_NOT_FOUND", "会话不存在", 404)
    return Response(status_code=status.HTTP_204_NO_CONTENT)


@chat_router.post("/stream")
async def stream_chat(
    payload: ChatStreamRequest,
    request: Request,
    session: SessionDependency,
) -> StreamingResponse:
    if await ChatRepository(session).get_session(payload.session_id) is None:
        raise AppError("SESSION_NOT_FOUND", "会话不存在", 404)
    query = payload.query.strip()
    if not query:
        raise AppError("VALIDATION_ERROR", "问题不能为空", 422)
    service = request.app.state.chat_service
    service.reserve(payload.session_id)
    return StreamingResponse(
        service.stream(payload.session_id, query, request.state.request_id),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "X-Request-ID": request.state.request_id,
        },
    )
