import json
from datetime import datetime, timezone

from fastapi import APIRouter, Depends, Query
from fastapi.responses import StreamingResponse
from sqlmodel import Session, select, func

from ..database import get_session, engine
from ..models import AIConfig, ChatMessage
from ..schemas import ChatRequest, ChatMessageRead, ChatMessageList
from ..services.ai_provider import stream_chat

router = APIRouter(prefix="/api/chat", tags=["chat"])


@router.post("/stream")
async def chat_stream(
    data: ChatRequest,
    session: Session = Depends(get_session),
):
    """SSE 流式对话端点"""
    config = session.exec(select(AIConfig)).first()
    if not config or not config.api_key:
        return StreamingResponse(
            _error_stream("请先配置AI服务"),
            media_type="text/event-stream",
        )

    # 保存用户消息
    user_msg = data.messages[-1] if data.messages else None
    if user_msg and user_msg.role == "user":
        db_msg = ChatMessage(
            novel_id=data.novel_id,
            role="user",
            content=user_msg.content,
        )
        session.add(db_msg)
        session.commit()

    # 构建消息列表
    messages = [{"role": m.role, "content": m.content} for m in data.messages]

    # 构建上下文摘要（前端展示用）
    context_preview = [{"role": m["role"], "content": m["content"][:200]} for m in messages]

    async def generate():
        # 先发送上下文预览
        yield f"data: {json.dumps({'context': context_preview})}\n\n"

        full_response = ""
        usage_info = {}
        try:
            async for event in stream_chat(config, messages):
                if event.text:
                    full_response += event.text
                    yield f"data: {json.dumps({'content': event.text})}\n\n"
                if event.usage:
                    # 累加usage
                    for k, v in event.usage.items():
                        if v:
                            usage_info[k] = usage_info.get(k, 0) + v
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"
        finally:
            # 保存AI回复
            if full_response:
                with Session(engine) as s:
                    ai_msg = ChatMessage(
                        novel_id=data.novel_id,
                        role="assistant",
                        content=full_response,
                    )
                    s.add(ai_msg)
                    s.commit()
            yield f"data: {json.dumps({'done': True, 'usage': usage_info or None})}\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


async def _error_stream(message: str):
    yield f"data: {json.dumps({'error': message})}\n\n"
    yield f"data: {json.dumps({'done': True})}\n\n"


@router.get("/messages", response_model=ChatMessageList)
def list_messages(
    novel_id: int | None = Query(default=None),
    limit: int = Query(default=50, ge=1, le=200),
    session: Session = Depends(get_session),
):
    query = select(ChatMessage)
    count_query = select(func.count(ChatMessage.id))
    if novel_id is not None:
        query = query.where(ChatMessage.novel_id == novel_id)
        count_query = count_query.where(ChatMessage.novel_id == novel_id)
    else:
        query = query.where(ChatMessage.novel_id.is_(None))  # type: ignore
        count_query = count_query.where(ChatMessage.novel_id.is_(None))  # type: ignore

    total = session.exec(count_query).one()
    items = session.exec(
        query.order_by(ChatMessage.created_at.desc()).limit(limit)
    ).all()
    # 返回时按时间正序
    items = list(reversed(items))
    return ChatMessageList(
        items=[ChatMessageRead.model_validate(m, from_attributes=True) for m in items],
        total=total,
    )


@router.delete("/messages", status_code=204)
def clear_messages(
    novel_id: int | None = Query(default=None),
    session: Session = Depends(get_session),
):
    query = select(ChatMessage)
    if novel_id is not None:
        query = query.where(ChatMessage.novel_id == novel_id)
    else:
        query = query.where(ChatMessage.novel_id.is_(None))  # type: ignore
    messages = session.exec(query).all()
    for msg in messages:
        session.delete(msg)
    session.commit()