from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse
from langchain_core.messages import HumanMessage
from chatbot import app as app_graph
from tools import update_retriever
from utils import STT, TTS
import os
import asyncio

app = FastAPI()

# ---------------- CORS ---------------- #
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

UPLOAD_DIR = "uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)


@app.get("/")
def health():
    return {"status": "API is running"}


# =============================
# Upload PDF
# =============================
@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    file_path = os.path.join(UPLOAD_DIR, file.filename)

    with open(file_path, "wb") as f:
        f.write(await file.read())

    update_retriever(file_path)

    return {
        "status": "success",
        "filename": file.filename
    }


# =============================
# Chat Endpoint
# =============================
@app.post("/chat")
async def chat(message: str, session_id: str = "default"):

    async def event_stream():
        async for chunk in app_graph.astream(
            {"messages": [HumanMessage(content=message)]},
            config={"configurable": {"thread_id": session_id}},
            stream_mode="messages",
        ):
            msg = chunk[0] if isinstance(chunk, tuple) else chunk
            if hasattr(msg, "content") and msg.content:
                yield msg.content + "\n"

    return StreamingResponse(event_stream(), media_type="text/plain")


# =============================
# STT
# =============================
@app.post("/stt")
async def stt(file: UploadFile = File(...)):
    return await STT(file)


# =============================
# TTS
# =============================
@app.post("/tts")
async def tts(text: str):
    audio_path = await TTS(text)
    return FileResponse(audio_path)