import gradio as gr
import cv2
import numpy as np

# --- Dummy SAM2 integration ---
def sam2_segment(frame, prompt_points):
    """
    Dummy segmentation function.
    In a real application, load and run your SAM2 model here with the given frame and prompt points.
    Returns a list of segmentation boundaries (e.g. bounding boxes or masks).
    """
    # For demonstration, we simply return a fixed bounding box.
    return [(50, 50, 100, 150)]  # (x, y, width, height)

# --- Global storage for annotations ---
annotations = []

# --- Helper functions ---
def parse_prompt_points(prompt_points_str):
    """
    Parse a string of prompt points (format: "x1,y1; x2,y2; ...")
    into a list of (x, y) tuples.
    """
    points = []
    for item in prompt_points_str.split(";"):
        item = item.strip()
        if item:
            try:
                x, y = item.split(",")
                points.append((int(x.strip()), int(y.strip())))
            except Exception as e:
                continue
    return points

def extract_frame(video_path, time_in_sec):
    """
    Extract a frame from the video at a given time (in seconds)
    using OpenCV.
    """
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_no = int(time_in_sec * fps)
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_no)
    ret, frame = cap.read()
    cap.release()
    if ret:
        # Convert BGR (OpenCV format) to RGB for display
        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    else:
        return None

def annotate_clip(video_file, start_time, end_time, clip_name, prompt_points_str):
    """
    Given an uploaded video file, a start and end time (in seconds) and a clip name,
    extract a representative frame (from the middle of the clip), parse the prompt points,
    run the SAM2 segmentation (dummy in this demo), and save the annotation.
    """
    prompt_points = parse_prompt_points(prompt_points_str)
    mid_time = (start_time + end_time) / 2
    frame = extract_frame(video_file, mid_time)
    if frame is None:
        return "Error extracting frame.", None

    # Run SAM2 segmentation on the extracted frame using the prompt points.
    boundaries = sam2_segment(frame, prompt_points)
    
    # Save the annotation information.
    annotation = {
        "clip_name": clip_name,
        "start_time": start_time,
        "end_time": end_time,
        "object_boundaries": boundaries
    }
    annotations.append(annotation)
    
    return annotation, frame

def refresh_annotations():
    return annotations

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## Video Annotation Demo with SAM2 Integration")
    
    # Video upload
    with gr.Row():
        video_input = gr.Video(label="Upload Video", source="upload", type="filepath")
    
    # Start and end time inputs (in seconds)
    with gr.Row():
        start_time_input = gr.Number(label="Start Time (sec)", value=0)
        end_time_input = gr.Number(label="End Time (sec)", value=10)
    
    # Clip name input
    clip_name_input = gr.Textbox(label="Clip Name", placeholder="Enter a name for this clip")
    
    # Input for prompt points for SAM2 segmentation
    prompt_points_input = gr.Textbox(
        label="Prompt Points (format: x1,y1; x2,y2; ...)",
        placeholder="e.g., 100,150; 200,250"
    )
    
    # Button to perform annotation and segmentation
    segment_button = gr.Button("Annotate Clip and Segment")
    
    # Outputs: JSON for annotation result and the extracted frame image
    output_text = gr.JSON(label="Annotation Result")
    output_image = gr.Image(label="Extracted Frame for Segmentation", type="numpy")
    
    segment_button.click(
        fn=annotate_clip, 
        inputs=[video_input, start_time_input, end_time_input, clip_name_input, prompt_points_input],
        outputs=[output_text, output_image]
    )
    
    # Accordion to display all saved annotations
    with gr.Accordion("Saved Annotations", open=False):
        annotation_history = gr.JSON(refresh_annotations())
        refresh_button = gr.Button("Refresh Annotations")
        refresh_button.click(fn=refresh_annotations, inputs=[], outputs=[annotation_history])

demo.launch()