import gradio as gr
import cv2
import numpy as np
from pathlib import Path

# Import the existing homography and warping function from your existing script
from billboard import apply_homography_and_warp


# Define a function to select points on the image (similar to your `get_point_interface`)
def get_point_interface(original_frame, points, evt: gr.SelectData):
    x, y = evt.index
    if points is None:
        points = []
    points = points.copy()  # Make a copy to avoid modifying in-place
    points.append((x, y))

    # Draw the points and lines on the image
    image = original_frame.copy()
    for pt in points:
        cv2.circle(image, pt, 5, (255, 0, 0), -1)
    if len(points) > 1:
        for i in range(len(points) - 1):
            cv2.line(image, points[i], points[i + 1], (255, 0, 0), 2)

    return image, points


# Function to process and apply homography (this will call your existing function)
def process_images(source_image, dest_image, roi_points):
    # Ensure that exactly four points are selected
    if len(roi_points) != 4:
        return "Error: Please select exactly four points."

    # Convert images to correct format (if not already in uint8)
    img_src = np.array(source_image)
    img_dst = np.array(dest_image)

    # Apply homography and get the warped image
    result_img = apply_homography_and_warp(img_src, img_dst, np.array(roi_points))

    # Return the processed image
    return result_img


# Gradio interface setup
with gr.Blocks(title="Homography Warping App") as demo:
    gr.Markdown("# Homography Warping App")
    gr.Markdown(
        "Upload two images (source and destination), select the four points for ROI on the destination image, and click 'Process Images' to warp the source image onto the destination."
    )

    with gr.Row(equal_height=True):
        source_image_input = gr.Image(label="Upload Source Image")
        dest_image_input = gr.Image(label="Upload Destination Image")

    with gr.Column():
        dest_image_roi = gr.Image(label="Click to select ROI points on destination image")
        original_frame_state = gr.State(None)
        selected_points = gr.State([])
        clear_points_button = gr.Button("Clear Points")
        process_button = gr.Button("Process Images")

    with gr.Row():
        output_image = gr.Image(label="Wrapped Image")

    examples = [
        ["./images/Apollo-8-Launch.png", "./images/times_square.jpg"],
    ]

    with gr.Row():
        gr.Examples(
            examples=examples,
            inputs=[source_image_input, dest_image_roi],
            label="Load Example Images",
        )

    dest_image_input.upload(
        lambda img: img, inputs=dest_image_input, outputs=dest_image_roi  # Simply return the uploaded image
    )

    # Adding the ROI point selection functionality
    dest_image_roi.select(
        get_point_interface, inputs=[dest_image_roi, selected_points], outputs=[dest_image_roi, selected_points]
    )

    # Clear the selected points
    clear_points_button.click(
        fn=lambda original_frame: (original_frame, []),
        inputs=dest_image_input,
        outputs=[dest_image_roi, selected_points],
    )

    # Callback for processing the images
    process_button.click(
        process_images, inputs=[source_image_input, dest_image_roi, selected_points], outputs=output_image
    )


# Launch the Gradio app
demo.launch()