from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from ultralytics import YOLO
import requests
from PIL import Image
import numpy as np
import cv2
import io
import os

app = FastAPI()

# Mount a static directory for serving images
app.mount("/tmp", StaticFiles(directory="/tmp"), name="static")

# Templates for rendering HTML
templates = Jinja2Templates(directory="templates")


def predict_yolo(image_path):
    # Load a model
    model = YOLO('yolov8n.pt')  # pretrained YOLOv8n model

    # Run batched inference on a list of images
    results = model(image_path)  # return a list of Results objects

    # Process results list
    for result in results:
        boxes = result.boxes  # Boxes object for bbox outputs
        # masks = result.masks  # Masks object for segmentation masks outputs
        # keypoints = result.keypoints  # Keypoints object for pose outputs
        # probs = result.probs  # Probs object for classification outputs
    predictions = boxes.json()
    return predictions


def draw_boxes(image, boxes):
    for box in boxes:
        x, y, w, h = box["bbox"]
        cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 0), 2)
    return image


@app.post("/uploadfile/")
async def create_upload_file(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    
    # Save the image to a static directory
    save_path = f"/tmp/{file.filename}"
    image.save(save_path)
    
    # Perform YOLO prediction
    predictions = predict_yolo(save_path)
    
    # Draw bounding boxes on the image
    image_np = np.array(image)
    image_with_boxes = draw_boxes(image_np, predictions)
    
    # Save the image with bounding boxes
    image_with_boxes_path = f"/tmp/{file.filename.split('.')[0]}_with_boxes.jpg"
    cv2.imwrite(image_with_boxes_path, cv2.cvtColor(image_with_boxes, cv2.COLOR_RGB2BGR))
    
    # Render the HTML with the image and bounding boxes
    return templates.TemplateResponse("prediction.html", {"request": file, "image_path": image_with_boxes_path})

@app.get("/test")
async def read_root():
    return {"message": "TEST"}


@app.get("/")
async def read_root():
    return {"message": "Hello, this is a YOLO prediction API using FastAPI!"}