from fastapi import FastAPI, UploadFile, File,BackgroundTasks
from rembg import remove

from fastapi.responses import FileResponse
import os
from starlette.middleware.cors import CORSMiddleware
# from fastapi.staticfiles import StaticFiles


app = FastAPI()




origins = [
    "http://localhost:3000",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Directory to store processed images
output_dir = "output"
os.makedirs(output_dir, exist_ok=True)
# app.mount("/processed_images", StaticFiles(directory=output_dir), name="processed_images")

@app.post("/remove-background/")
async def remove_background(file: UploadFile):
    # Ensure the uploaded file is an image (you can add more validation)
    if not file.content_type.startswith("image/"):
        return {"error": "Invalid file format"}

    # Read the uploaded image file
    image_bytes = await file.read()

    # Process the image to remove the background using Rembg
    output_image_bytes = remove(image_bytes)

    # Save the processed image to a file
    output_filename = f"{output_dir}/{file.filename}"

    # output_filename  = f"processed_images/{file.filename.split('/')[-1]}"
    with open(output_filename, "wb") as output_file:
        output_file.write(output_image_bytes)

    return {"result": "Background removed successfully", "output_filename": output_filename}

    

@app.get("/download/{filename}")
async def download_processed_image(filename: str):
    file_path = os.path.join(output_dir, filename)
    # file_path = f"/output/{filename}"
    if not os.path.exists(file_path):
        return {"error": "File not found"}

    return FileResponse(file_path)


@app.get("/delete")
def deleteAll():
    os.rmdir("output")