import gradio as gr
import os
import base64
from PIL import Image
from io import BytesIO
from specklepy.api.client import SpeckleClient
from specklepy.transports.server import ServerTransport
from specklepy.api import operations
from specklepy.objects import Base

# Initialize Speckle client with your token
speckle_token = os.getenv('SPECKLE_TOKEN')
speckle_branch = os.getenv('SPECKLE_BRANCH')
speckle_stream_id = os.getenv('SPECKLE_STREAM_ID')

CLIENT = SpeckleClient(host="https://speckle.xyz/")
CLIENT.authenticate_with_token(token=speckle_token)

def send_to_speckle(stream_id, branch_name, client, data_object):
    try:
        # Get transport
        transport = ServerTransport(client=client, stream_id=stream_id)
        
        # Send the data object to the speckle stream
        object_id = operations.send(data_object, [transport])

        # Create a new commit with the new object
        commit_id = client.commit.create(
            stream_id=stream_id,
            object_id=object_id,
            branch_name=branch_name,
            message="Updated the data object from Gradio app",
        )
        
        return f"Success! Data sent to Speckle stream. Commit ID: {commit_id}"
    except Exception as e:
        return f"An error occurred: {e}"

def process_data(image, text):
    # Process image: Resize and convert to base64
    img_str = ""
    if image is not None:
        # Resize image, maintaining aspect ratio, and max width 1024 pixels
        base_width = 1024
        w_percent = (base_width / float(image.size[0]))
        h_size = int((float(image.size[1]) * float(w_percent)))
        image = image.resize((base_width, h_size), Image.Resampling.LANCZOS)
        
        # Convert to base64
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')


    # Prepare data object
    obj = Base()
    obj["image"] = img_str
    obj["text"] = text

    # Send data to Speckle
    return send_to_speckle(speckle_stream_id, speckle_branch, CLIENT, obj)
    
with gr.Blocks() as demo:
    gr.Markdown("### Upload Image and Enter Text")
    with gr.Row():
        image = gr.Image(type="pil", label="Add Image")
        text = gr.TextArea(label="Enter Text or Speak")
    submit_btn = gr.Button("Submit")
    output = gr.Textbox(label="Status", visible=False)  # Set visible to False to hide output

    submit_btn.click(fn=process_data, inputs=[image, text], outputs=output)

demo.launch()