from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator

from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.core.database import get_db
from app.core.dependencies import get_current_user
from app.core.responses import api_success
from app.models.user import User
from app.schemas.chat import (
    ChatCompletionRequest,
    ChatMessageRead,
    ChatSessionRead,
    CreateSessionResponse,
    StopChatRequest,
    UpdateSessionTitleRequest,
)
from app.services.chat_service import ChatService
from app.services.chat_queue_runtime import (
    cancel_chat_waiter,
    release_chat_slot,
    request_chat_slot,
    wait_for_chat_slot,
)
from app.services.chat_queue_service import load_chat_queue_config
from app.services.chat_stream_service import ChatStreamService

router = APIRouter()


@router.post("/session")
def create_session(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)) -> dict:
    session = ChatService.create_session(db, current_user)
    return api_success(CreateSessionResponse(sessionId=session.id).model_dump())


@router.get("/session/list")
def list_sessions(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)) -> dict:
    sessions = ChatService.list_sessions(db, current_user)
    return api_success([ChatSessionRead.model_validate(session).model_dump() for session in sessions])


@router.get("/history")
def history(
    sessionId: int = Query(...),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
) -> dict:
    messages = ChatService.get_history(db, current_user, sessionId)
    return api_success([ChatMessageRead.model_validate(message).model_dump(mode="json") for message in messages])


@router.put("/session/title")
def update_title(
    payload: UpdateSessionTitleRequest,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
) -> dict:
    session = ChatService.update_title(db, current_user, payload.sessionId, payload.title)
    return api_success(ChatSessionRead.model_validate(session).model_dump())


@router.delete("/session/{session_id}")
def delete_session(
    session_id: int,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
) -> dict:
    ChatService.delete_session(db, current_user, session_id)
    return api_success()


@router.post("/completions")
def completions(
    payload: ChatCompletionRequest,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
) -> StreamingResponse:
    return StreamingResponse(_chat_stream(payload, db, current_user), media_type="text/event-stream")


@router.post("/stop")
def stop(
    payload: StopChatRequest,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
) -> dict:
    ChatService.stop_generation(db, current_user, payload.sessionId)
    return api_success()


async def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User) -> AsyncIterator[str]:
    config = load_chat_queue_config(db)
    queue_lease = await request_chat_slot(config)
    queue_request = queue_lease.request
    acquired = queue_request.status == "acquired"
    should_cancel_waiter = queue_request.status == "queued"

    if queue_request.status == "rejected":
        yield _sse_event(
            "error",
            message="当前请求过多，请稍后再试。",
            activeCount=queue_request.active_count,
            waitingCount=queue_request.waiting_count,
        )
        yield _sse_done()
        return

    if queue_request.status == "queued" and queued_token is not None:
        yield _sse_event(
            "queued",
            message=f"当前请求较多，正在排队中，前方约 {max(queue_request.position - 1, 0)} 个请求。",
            position=queue_request.position,
            activeCount=queue_request.active_count,
            waitingCount=queue_request.waiting_count,
        )
        queue_lease = await wait_for_chat_slot(queue_lease, config)
        queue_request = queue_lease.request
        acquired = queue_request.status == "acquired"
        should_cancel_waiter = queue_request.status == "queued"
        if queue_request.status == "rejected":
            yield _sse_event(
                "error",
                message="当前请求过多，请稍后再试。",
                activeCount=queue_request.active_count,
                waitingCount=queue_request.waiting_count,
            )
            yield _sse_done()
            return
        if queue_request.status == "timeout":
            yield _sse_event("error", message="排队等待超时，请稍后再试。")
            yield _sse_done()
            return
        if queue_request.status == "canceled":
            yield _sse_event("error", message="排队请求已取消，请稍后再试。")
            yield _sse_done()
            return

    try:
        yield _sse_event(
            "generating",
            message="已进入生成队列，正在生成回答。",
            activeCount=queue_request.active_count,
            waitingCount=queue_request.waiting_count,
        )
        async for chunk in ChatStreamService.stream_answer_async(db, current_user, payload.sessionId, payload.message):
            yield _sse_event("content", content=chunk)
    except HTTPException as exc:
        yield _sse_event("error", message=str(exc.detail))
    except asyncio.CancelledError:
        raise
    except Exception:
        yield _sse_event("error", message="AI 回复生成失败，请稍后再试。")
    finally:
        if acquired:
            await release_chat_slot(queue_lease)
        elif should_cancel_waiter:
            await cancel_chat_waiter(queue_lease)

    yield _sse_done()

def _sse_event(event_type: str, **payload: object) -> str:
    return f"data: {json.dumps({'type': event_type, **payload}, ensure_ascii=False)}\n\n"


def _sse_done() -> str:
    return "data: [DONE]\n\n"