import os
import random
import torch
from diffusers import StableDiffusionInpaintPipeline
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import FileResponse
from PIL import Image
from io import BytesIO
from dotenv import load_dotenv
import uvicorn

# Load environment variables
load_dotenv()

app = FastAPI()

# Load the inpainting model
model_name = os.getenv("MODEL_NAME", "stabilityai/stable-diffusion-2-inpainting")
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionInpaintPipeline.from_pretrained(model_name, torch_dtype=torch.float16 if device == "cuda" else torch.float32)

pipe.to(device)


@app.post("/inpaint/")
async def inpaint_image(file: UploadFile = File(...), target_width: int = 512, target_height: int = 512):
    # Load the input image
    input_image = Image.open(BytesIO(await file.read())).convert("RGB")

    random_numm = random.randint(0, 10000)

    # Resize the image to the target dimensions
    input_image = input_image.resize((target_width, target_height))
    input_image.save(f"./temp/inpainted_{random_numm}_{file.filename}")

    full_mask = Image.new('RGB', (target_width, target_height), color='black')
    full_mask.save(f"./temp/inpainted_mask_{file.filename}")

    # Generate the inpainted image
    inpainted_image = pipe(
        prompt=[
            "Enhance this image to improve quality.",
            "Better image quality",
            "Higher image quality",
            "High Resolution"
        ],
        negative_prompt=[
            "blur",
            "low resolution",
            "pixelation",
            "noise"
        ],
        image=input_image,
        mask_image=full_mask,
    ).images[0]

    # Save the image to the static directory
    output_path = f"static/images/inpainted_{random_numm}_{file.filename}"
    inpainted_image.save(output_path)

    print(output_path)

    return FileResponse(output_path)

@app.get("/")
async def root():
    return {"message": "Image Inpainting API is running!"}