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

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('ultralyticsplus/yolov8s')

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

    return results


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("/", response_class=HTMLResponse)
async def detect_yolo(request: Request, url: str = Form(...)):
    try:
        # Download the image from the specified URL
        async with httpx.AsyncClient() as client:
            response = await client.get(url)
            response.raise_for_status()  # Raise an exception if there is an error in the request
            content = response.content

        image = Image.open(io.BytesIO(content))

        results = predict_yolo(image)

        render = render_result(model=model, image=image, result=results[0])

        # Save the modified image to a byte sequence
        image_byte_array = io.BytesIO()
        image.save(image_byte_array, format="PNG")
        
        # Return the image as a response with content type "image/png"
        return templates.TemplateResponse("result.html", {"request": request, "image": base64.b64encode(image_byte_array.getvalue()).decode()})

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")


@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!"}