# works with gradio file upload, not image upload
import base64
from fastapi import FastAPI #, UploadFile, File
from img_upload import upload_image_to_blob
from image_processing import process_image
from pydantic import BaseModel, validator
from PIL import Image
import io
import gradio as gr
import uuid

app = FastAPI()

class FileUpload(BaseModel):
    filename: str
    data: str

    # @validator('data')
    # def validate_image(cls, data: str):
    #     try:
    #         image_data = base64.b64decode(data)
    #         image = Image.open(BytesIO(image_data))
    #         if image.format not in ['JPEG', 'PNG']:
    #             raise ValueError('Invalid file type')
    #         if max(image.size) > 5000:
    #             raise ValueError('Image dimensions are too large')
    #         if len(data) > 5000 * 5000:  # adjust this value based on your needs
    #             raise ValueError('File size is too large')
    #         return data
    #     except Exception as e:
    #         raise ValueError('Invalid image') from e
    #

class Response(BaseModel):
    result: str


@app.post("/upload", response_model=Response)
async def create_upload_file(file: FileUpload):
    data = base64.b64decode(file.data)
    sas_url = upload_image_to_blob(data, file.filename)
    result = process_image(sas_url)
    return result

async def gradio_interface(image: Image.Image):
    # Convert PIL Image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format="JPEG")
    img_byte_arr = img_byte_arr.getvalue()

    # Encode bytes to base64
    data = base64.b64encode(img_byte_arr).decode()

    # Generate a unique ID for the image
    unique_id = str(uuid.uuid4())

    response = await create_upload_file(FileUpload(filename=unique_id + ".jpg", data=data))
    return response

iface = gr.Interface(fn=gradio_interface, inputs=gr.Image(type="pil"), outputs="text")
app = gr.mount_gradio_app(app, iface, "/gradio")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000, )