import asyncio
import os
import shutil
import subprocess
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image

app = FastAPI()

# Set up CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

async def delete_files_after_delay(delay, path):
    await asyncio.sleep(delay)
    # Check if path is a file or directory
    if os.path.isfile(path):
        os.remove(path)
    elif os.path.isdir(path):
        shutil.rmtree(path)

@app.post("/experiments/remove-background/")
async def remove_background(file: UploadFile = File(...)):
    input_path = f"temp/{file.filename}"
    compressed_path = f"temp/compressed-{file.filename}"
    output_path = f"temp/no-bg-{file.filename}"

    # Ensure the temp folder exists
    os.makedirs(os.path.dirname(input_path), exist_ok=True)

    # Save the uploaded file
    with open(input_path, "wb") as buffer:
        buffer.write(await file.read())

    # Load the image
    original_image = Image.open(input_path)

    # If the image has an alpha channel, convert it to RGB
    if original_image.mode == 'RGBA':
        original_image = original_image.convert('RGB')

    # Save the image with reduced quality (e.g., quality=85)
    original_image.save(compressed_path, 'JPEG', quality=85)

    # Run rembg command on the compressed image
    subprocess.run(["rembg", "i", compressed_path, output_path], check=True)

    # Schedule deletion of the input and output images after 30 seconds
    asyncio.create_task(delete_files_after_delay(60, input_path))
    asyncio.create_task(delete_files_after_delay(60, compressed_path))
    asyncio.create_task(delete_files_after_delay(60, output_path))

    # Return the processed file
    return FileResponse(output_path)