import socketio
import uvicorn
from fastapi import FastAPI

from main import build_rag_chain, build_vectorstore, load_env

load_env()

app = FastAPI()
sio = socketio.AsyncServer(async_mode="asgi", cors_allowed_origins="*")
socket_app = socketio.ASGIApp(sio, other_asgi_app=app)

rag_chain = None


@app.on_event("startup")
async def startup():
    global rag_chain
    print("Loading vectorstore...")
    vectorstore = build_vectorstore()
    print("Building RAG chain...")
    rag_chain = build_rag_chain(vectorstore)
    print("Server ready!")


@sio.event
async def connect(sid, environ):
    print(f"Client connected: {sid}")


@sio.event
async def disconnect(sid):
    print(f"Client disconnected: {sid}")


@sio.event
async def ask(sid, data):
    question = data.get("question", "").strip()
    if not question:
        await sio.emit("answer", {"error": "Câu hỏi trống"}, to=sid)
        return

    try:
        answer = await rag_chain.ainvoke(question)
        await sio.emit("answer", {"answer": answer}, to=sid)
    except Exception as e:
        print(f"Error: {e}")
        await sio.emit("answer", {"error": str(e)}, to=sid)


if __name__ == "__main__":
    uvicorn.run(socket_app, host="0.0.0.0", port=8000)