# main.py
from fastapi import FastAPI, Request, UploadFile, Form
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import shutil
import os
import app
import appImage

app = FastAPI()

# Mount static and templates
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/resources", StaticFiles(directory="resources"), name="resources")

templates = Jinja2Templates(directory="templates")

# Serve your main home page
@app.get("/", response_class=HTMLResponse)
async def serve_home(request: Request):
    return templates.TemplateResponse("home.html", {"request": request})

# Route to handle prediction requests (image or doc)
@app.post("/predict")
async def predict(file: UploadFile = Form(...), question: str = Form(...)):
    ext = file.filename.split(".")[-1].lower()
    with open(file.filename, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    if ext in ["pdf", "docx", "pptx", "xlsx"]:
        answer, audio = app.answer_question_from_doc(file=buffer, question=question)
    else:
        image = appImage.Image.open(file.filename)
        answer, audio = appImage.answer_question_from_image(image, question)

    os.remove(file.filename)
    return JSONResponse({"answer": answer, "audio": audio})