# import db.postgres as pg
import uuid
from uuid import UUID

from fastapi import APIRouter, Depends

import db.postgres as pg
from config.security import get_user_id_from_token
from models.BaseResponse import BaseResponse
from models.ChatRequest import ChatRequest

chatRouter = APIRouter()
from agent.dataAgent import get_graph_output
from llm.chatLLM import get_chat_response
from llm.titleChain import get_title


# 对话历史记录
@chatRouter.get("/history")
def getHistory(sessionId: str):
    return BaseResponse(data=pg.get_history(sessionId))


# 对话列表
@chatRouter.get("/sessionsForBot")
def getSessions(user_id: UUID = Depends(get_user_id_from_token)):
    if not user_id:
        return {"error": "userId is required"}
    return BaseResponse(data=pg.get_sessions(user_id, "bot"))


@chatRouter.post("/chatForBot")
def chat(req: ChatRequest, user_id: UUID = Depends(get_user_id_from_token)):
    if not user_id:
        return {"error": "userId is required"}
    if not req.aiId:
        return {"error": "aiId is required"}
    session_name = get_title(req.userInput)
    # 如果没有 sessionId 就新建
    if not req.sessionId:
        is_new_session = True
        req.sessionId = str(uuid.uuid4())
        pg.insert_session(user_id, req.aiId, req.sessionId, session_name, "bot")
    else:
        is_new_session = False
        pg.update_session_updated_at(req.sessionId)

    # 插入用户消息
    pg.insert_message(req.sessionId, False, req.userInput)

    # 调用 LLM
    if req.aiId == "9d157dd1-921b-c768-5b90-3e903b50f6f9":
        # 数据专家AI
        answer = get_graph_output(
            aiRole=pg.get_ai_personality(req.aiId),
            history=pg.get_history(req.sessionId),
            userInput=req.userInput,
        )
    else:
        answer = get_chat_response(
            aiRole=pg.get_ai_personality(req.aiId),
            history=pg.get_history(req.sessionId),
            userInput=req.userInput,
        ).content
    # 插入 AI 回复
    pg.insert_message(req.sessionId, True, answer)

    return BaseResponse(
        data={
            "session_name": session_name,
            "isNewSession": is_new_session,
            "content": answer,
            "sessionId": req.sessionId,
        }
    )