from fastapi import FastAPI, File, UploadFile
from fastapi.responses import FileResponse
from ml import inference
from PIL import Image
import io
import os
import uuid

app = FastAPI()


@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.post("/translate/image/")
async def process_image(file: UploadFile = File(...)):
    img_id = uuid.uuid4()
    temp_path = f"temp/temp_input_image_{img_id}.jpg"

    # Create an in-memory buffer of the uploaded file
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))

    # Convert the image to JPEG if it's not already in that format
    if image.format != 'JPEG':
        # Convert image to RGB if necessary (e.g., handling transparency in PNGs)
        if image.mode in ("RGBA", "P"):
            image = image.convert("RGB")

        image.save(temp_path, 'JPEG')
    else:
        # Save the original image if already a JPEG
        with open(temp_path, "wb") as buffer:
            buffer.write(contents)

    # Process the image (this part will depend on your specific processing logic)
    # For demonstration, assume the processed image is saved as 'processed_image.jpg'
    inf_res = inference(temp_path, img_id=img_id, show_bound=True)
    print('inf_res', inf_res)

    processed_image_path = os.path.join(os.getcwd(), f'result/translate_result_{img_id}.jpg')

    # Return the processed image
    return FileResponse(processed_image_path)