import cv2
import gradio as gr
import numpy as np

# Global variable to store the list of arrows (start and end points)
arrows = []
start_point = None


# Function to draw all arrows (including zero-length arrows) on the image
def draw_arrow(image, click_coords):
    global start_point, arrows

    # Convert the image to a numpy array if it's not already
    img_array = np.array(image, dtype=np.uint8)

    # Get the current point from the user (click coordinates as (x, y))
    current_point = (int(click_coords[0]), int(click_coords[1]))  # Convert float coords to int

    # If start point is not set, set it as the current click position
    if start_point is None:
        start_point = current_point
        return img_array  # No arrow yet, just return the original image

    # If start point is already set, add an arrow (including zero-length ones)
    end_point = current_point
    arrows.append((start_point, end_point))  # Save the arrow

    # Reset start_point for the next arrow
    start_point = None

    # Draw all arrows from the saved list
    for arrow in arrows:
        start, end = arrow
        color = (0, 255, 0)  # Green arrow
        thickness = 3

        # Draw the arrow (even if it's zero-length, i.e., start == end)
        cv2.arrowedLine(img_array, start, end, color, thickness)

    return img_array


# Function to reset the canvas (clearing all arrows)
def reset_canvas():
    global arrows, start_point
    arrows = []
    start_point = None
    return load_image()  # Return a fresh image


# Load an image for the user to interact with
def load_image():
    img = np.ones((400, 400, 3), dtype=np.uint8) * 255  # White background image
    return img


# Define Gradio interface using Blocks
def interactive_arrow_interface():
    with gr.Blocks() as demo:
        image_input = gr.Image(value=load_image(), interactive=True,
                               label="Click to specify the arrow's start and end points")
        output_image = gr.Image(label="Image with Arrows")
        reset_button = gr.Button("Reset")

        # Set up interaction: Handle click events with 'handle_click'
        def handle_click(image, evt: gr.SelectData):
            print(f"Click coordinates: {evt.index}")
            # Pass click coordinates to draw_arrow function and update the image
            updated_image = draw_arrow(image, (evt.index[0], evt.index[1]))
            return updated_image

        image_input.select(handle_click, [image_input], output_image)

        # Set up the reset button to clear all arrows
        reset_button.click(fn=reset_canvas, inputs=None, outputs=image_input)

        return demo


# Launch the interactive demo
interactive_arrow_interface().launch(inbrowser=True)