from __future__ import annotations

import json
from collections.abc import Iterator

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_service import chat_queue_manager, load_chat_queue_config

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()


def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User) -> Iterator[str]:
    config = load_chat_queue_config(db)
    queue_request = chat_queue_manager.request_slot(config)
    acquired = queue_request.status == "acquired"
    queued_token = queue_request.token

    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_request = chat_queue_manager.wait_for_slot(queued_token, config)
        acquired = queue_request.status == "acquired"
        if queue_request.status == "timeout":
            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,
        )
        answer = ChatService.create_answer(db, current_user, payload.sessionId, payload.message)
        yield from _sse_chunks(answer)
    except HTTPException as exc:
        yield _sse_event("error", message=str(exc.detail))
    except Exception:
        yield _sse_event("error", message="AI 回复生成失败，请稍后再试。")
    finally:
        if acquired:
            chat_queue_manager.release_slot()
        elif queued_token is not None:
            chat_queue_manager.cancel_waiter(queued_token)

    yield _sse_done()


def _sse_chunks(answer: str) -> Iterator[str]:
    chunk_size = 12
    for index in range(0, len(answer), chunk_size):
        chunk = answer[index : index + chunk_size]
        yield _sse_event("content", content=chunk)


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"