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

# Initialize FastAPI
app = FastAPI()

# Load environment variables if needed
HF_TOKEN = os.getenv("HF_TOKEN")

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

@app.post("/upload/")
async def upload_image(file: UploadFile = File(...)):
    try:
        # Save the uploaded file to a temporary location
        file_location = f"temp_{file.filename}"
        with open(file_location, "wb") as f:
            f.write(await file.read())

        # Use the Gradio client to send the image to the Gradio app
        result = client.predict(
            img=handle_file(file_location),
            api_name="/predict"
        )

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

        # Read the sketch image from the result
        sketch_image_path = result[0]  # Path to the sketch image

        # Convert the sketch image to base64
        with open(sketch_image_path, "rb") as img_file:
            sketch_image_base64 = base64.b64encode(img_file.read()).decode('utf-8')

        # Prepare the response
        response = {
            "sketch_image_base64": f"data:image/jpeg;base64,{sketch_image_base64}"
        }

        return JSONResponse(content=response)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))