from fastapi import FastAPI, File, UploadFile
from fastapi import FastAPI, File, UploadFile, Form, Request
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from fastapi.responses import StreamingResponse
from gradio_client import Client
import os
import io
app = FastAPI()

hf_token = os.environ.get('HF_TOKEN')
client = Client("https://ashrafb-image-to-sketch.hf.space/", hf_token=hf_token)

import tempfile

import base64

@app.post("/upload/")
async def upload_file(file: UploadFile = File(...)):
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
        temp_file.write(await file.read())
        temp_file_path = temp_file.name

    try:
        result = client.predict(temp_file_path, api_name="/predict")
        
        # Convert the image data to base64 string
        with open(result[0], "rb") as image_file:
            image_data = base64.b64encode(image_file.read()).decode("utf-8")

        return {"sketch_image_base64": image_data}
    finally:
        os.unlink(temp_file_path)


app.mount("/", StaticFiles(directory="static", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/static/index.html", media_type="text/html")