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.security import check_rate_limit, validate_image_size
from config import DEVICE, ENABLE_EXTERNAL_AI
from database import init_db, login_user, signup_user, log_request_db, check_gemini_limit, increment_gemini_limit, get_user_logs
from fastapi import HTTPException


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

@app.on_event("startup")
async def startup_event():
    init_db()

# ================= 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):
    import os
    cid = os.getenv("GOOGLE_CLIENT_ID", "")
    return templates.TemplateResponse("index.html", {
        "request": request, 
        "google_client_id": cid
    })


# ================= API =================
@app.post("/api/login")
async def handle_login(request: Request):
    data = await request.json()
    username = data.get("username")
    password = data.get("password")
    
    if login_user(username, password):
        return {"status": "success", "user": username}
    return JSONResponse(status_code=401, content={"status": "error", "message": "Invalid password or user does not exist."})

@app.post("/api/signup")
async def handle_signup(request: Request):
    data = await request.json()
    username = data.get("username")
    password = data.get("password")
    
    if not username or not password:
        return JSONResponse(status_code=400, content={"status": "error", "message": "Username and password required."})
        
    if signup_user(username, password):
        return {"status": "success", "user": username}
    return JSONResponse(status_code=400, content={"status": "error", "message": "Username already taken."})

@app.post("/api/google_login")
async def handle_google_login(request: Request):
    data = await request.json()
    email = data.get("email")
    if not email or "@" not in email:
        return JSONResponse(status_code=400, content={"status": "error", "message": "Invalid Google Email."})
    
    username = email.split("@")[0].lower()
    from database import login_or_create_google_user
    login_or_create_google_user(username)
    
    return {"status": "success", "user": username}

@app.get("/api/logs")
async def get_logs_endpoint(user: str):
    logs = get_user_logs(user)
    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)

    if model_choice == "gemini":
        if not check_gemini_limit(user):
            return JSONResponse({
                "device": "Error",
                "model_used": "gemini",
                "caption": "Quota Exceeded",
                "answer": "You have reached your limit of 10 Gemini requests per day.",
                "explanation": "To prevent abusive cost spikes, users are limited to 10 Gemini requests per day. The Admin has unlimited access.",
                "external_enabled": True
            })

    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_db(user, model_choice, question)
    
    if model_choice == "gemini":
        increment_gemini_limit(user)

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