main.py

import io
from fastapi import FastAPI, File, UploadFile, Form, Request
from fastapi.responses import JSONResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image

from ai_router import route_model
from utils.logger import log_request
from utils.security import check_rate_limit, validate_image_size
from config import DEVICE, ENABLE_EXTERNAL_AI


# ================= CREATE APP FIRST =================
app = FastAPI(title="GenAI VQA System")


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


# ================= STATIC + TEMPLATE =================
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")


# ================= HOME =================
@app.get("/")
async def home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


# ================= API =================
@app.get("/api/logs")
async def get_logs_endpoint(user: str):
    import os
    logs = []
    if os.path.exists("logs/requests.log"):
        with open("logs/requests.log", "r") as f:
            for line in f:
                if " - User: " in line:
                    try:
                        parts = line.split(" - INFO - User: ")
                        timestamp = parts[0]
                        rest = parts[1].split(" | Model: ")
                        log_user = rest[0].strip()
                        model_q = rest[1].split(" | Q: ")
                        model = model_q[0].strip()
                        question = model_q[1].strip()
                        
                        # Admin sees all, User sees only their own
                        if user == "vignesh" or log_user == user:
                            logs.append({
                                "timestamp": timestamp, 
                                "user": log_user, 
                                "model": model, 
                                "question": question
                            })
                    except Exception as e:
                        pass
    logs.reverse()
    return {"logs": logs}


@app.post("/ask")
async def ask_question(
    file: UploadFile = File(...),
    question: str = Form(...),
    model_choice: str = Form("reasoning"),
    user: str = Form("guest"),
    lang: str = Form("en")
):

    check_rate_limit(user)

    image_bytes = await file.read()
    validate_image_size(len(image_bytes))

    try:
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except Exception as e:
        return JSONResponse({
            "device": "Error",
            "model_used": model_choice,
            "caption": "Image Load Error",
            "answer": f"The supplied image was empty or an invalid format. Size received: {len(image_bytes)} bytes.",
            "explanation": str(e),
            "external_enabled": False
        })

    caption, answer, explanation = route_model(
        model_choice,
        image,
        question,
        lang
    )

 

    log_request(user, model_choice, question)

    return JSONResponse({
        "device": DEVICE,
        "model_used": model_choice,
        "caption": caption,
        "answer": answer,
        "explanation": explanation,
        "external_enabled": ENABLE_EXTERNAL_AI
    })