from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
from gradio_client import Client, handle_file
import shutil
import base64
import os

app = FastAPI()


HF_TOKEN = os.getenv("HF_TOKEN")

# Initialize the Gradio client with the token
client = Client("Makhinur/Image_Face_Upscale_Restoration-GFPGAN", hf_token=HF_TOKEN)

# Version mapping from HTML to Gradio API
version_map = {
    "M1": "v1.2",
    "M2": "v1.3",
    "M3": "v1.4"
}

@app.post("/upload/")
async def enhance_image(
    file: UploadFile = File(...),
    version: str = Form(...),
    scale: int = Form(...)
):
    # Map version from HTML to Gradio expected value
    gradio_version = version_map.get(version, "v1.4")

    # Save the uploaded image to a temporary file
    temp_file_path = "temp_image.png"
    with open(temp_file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    try:
        # Use the Gradio client to process the image
        result = client.predict(
            img=handle_file(temp_file_path),
            version=gradio_version,
            scale=scale,
            api_name="/predict"
        )

        # Read the result image and encode it in base64
        with open(result[0], "rb") as img_file:
            b64_string = base64.b64encode(img_file.read()).decode('utf-8')

        # Clean up the temporary file
        os.remove(temp_file_path)

        return JSONResponse(content={"sketch_image_base64": f"data:image/png;base64,{b64_string}"})

    except Exception as e:
        # Log the error message for debugging
        print(f"Error processing image: {e}")
        return JSONResponse(status_code=500, content={"message": "Internal Server Error"})