from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from gradio_client import Client, handle_file
from PIL import Image
import shutil
import os

app = FastAPI()

client = Client("zhengchong/CatVTON")


def convert_to_rgb(input_path, output_path):
    img = Image.open(input_path)
    img = img.convert("RGB")
    img.save(output_path, "JPEG")


@app.get("/")
def home():
    return {"message": "FitVision API Running"}


@app.post("/tryon")
async def tryon(
    person: UploadFile = File(...),
    cloth: UploadFile = File(...)
):
    try:
        os.makedirs("temp", exist_ok=True)

        person_path = "temp/person.jpg"
        cloth_path = "temp/cloth.jpg"

        with open(person_path, "wb") as buffer:
            shutil.copyfileobj(person.file, buffer)

        with open(cloth_path, "wb") as buffer:
            shutil.copyfileobj(cloth.file, buffer)

        convert_to_rgb(person_path, person_path)
        convert_to_rgb(cloth_path, cloth_path)

        person_image = client.predict(
            image_path=handle_file(person_path),
            api_name="/person_example_fn"
        )

        result = client.predict(
            person_image=person_image,
            cloth_image=handle_file(cloth_path),
            cloth_type="upper",
            num_inference_steps=30,
            guidance_scale=2.5,
            seed=42,
            show_type="result only",
            api_name="/submit_function"
        )

        return JSONResponse({
            "success": True,
            "result": result
        })

    except Exception as e:
        return JSONResponse({
            "success": False,
            "error": str(e)
        })