# main.py

from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
import uvicorn
import shutil
import os
import uuid
import cv2
import numpy as np
import base64

from yolo_numbering import predict_yolo as predict_yolo_numbering
from yolo_anomaly import predict_yolo_anomaly
app = FastAPI()

UPLOAD_DIR = "/tmp/uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)
@app.get("/")
async def root():
    return {"message": "API is up and running!"}

@app.post("/predict/")
async def predict_endpoint(
    file: UploadFile = File(...),
    model: str = Form(...),  # either "detectron" or "yolo"
):
    # Save uploaded image
    file_ext = file.filename.split('.')[-1]
    filename = f"{uuid.uuid4()}.{file_ext}"
    file_path = os.path.join(UPLOAD_DIR, filename)

    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    try:
        if model == "Anomaly":
            result_img, predictions = predict_yolo_anomaly(file_path)
        elif model == "Numbering":
            result_img, predictions = predict_yolo_numbering(file_path)
        else:
            return JSONResponse({"error": "Invalid model choice"}, status_code=400)

        # Encode image to bytes (optional)
        _, img_encoded = cv2.imencode(".jpg", result_img)
        img_bytes = img_encoded.tobytes()

        img_b64 = base64.b64encode(img_bytes).decode('utf-8')

        return {
          "predictions": predictions,
          "image_base64": img_b64
              }
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)
    finally:
        os.remove(file_path)  # Clean up uploaded file


# Only for testing locally
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)