from __future__ import annotations

import json
from collections.abc import Iterator

from fastapi import APIRouter, Depends, 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

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:
    answer = ChatService.create_answer(db, current_user, payload.sessionId, payload.message)
    return StreamingResponse(_sse_chunks(answer), 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 _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 f"data: {json.dumps({'content': chunk}, ensure_ascii=False)}\n\n"
    yield "data: [DONE]\n\n"
