gui.py

import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from moto_vision import DEFAULT_OPENROUTER_API_KEY
import process_video
import os
from PIL import Image, ImageTk
import cv2
import json
import time

class HaneGUI:
    def __init__(self, master):
        self.master = master
        master.title("hane")

        # Create Tab Control
        self.tabControl = ttk.Notebook(master)
        self.processing_tab = ttk.Frame(self.tabControl)
        self.analysis_tab = ttk.Frame(self.tabControl)

        self.tabControl.add(self.processing_tab, text='Processing')
        self.tabControl.add(self.analysis_tab, text='Analysis')
        self.tabControl.pack(expand=1, fill="both")

        # --- Processing Tab ---
        self.setup_processing_tab(self.processing_tab)

        # --- Analysis Tab ---
        self.OBJECT_COLORS = {
            "traffic light": "red",
            "stop sign": "orange",
            "motorcycle": "deep sky blue", # Changed from "blue" for better visibility
            "car": "lightgreen",
            "truck": "lightgreen"
        }
        self.OBJECT_PRIORITY = ["traffic light", "stop sign", "motorcycle", "car", "truck"]
        self.setup_analysis_tab(self.analysis_tab)

        self.debug_photo_images = [] # List to store PhotoImage objects for processing tab preview
        self.handle_yolo_model_selection(self.yolo_model_path_var.get()) # Initial call for processing tab
        self.toggle_debug_features() # Initial state of debug features for processing tab

        # Video Player Attributes
        self.analysis_data = None
        self.video_cap = None
        self.video_fps = 0
        self.video_frame_count = 0
        self.total_video_duration_seconds = 0
        self.current_video_time_seconds = 0
        self.is_video_playing = False
        self.video_player_image = None # For PhotoImage reference

        # Attributes for dynamic color strip
        self.color_strip_segments = []
        self.last_known_color_strip_width = 0
        self.SLIDER_VISUAL_OFFSET_X = 15 # Offset to align color strip with slider track

    def setup_processing_tab(self, tab_control):
        # Input Video Path
        tk.Label(tab_control, text="Input Video Path:").grid(row=0, column=0, sticky=tk.W, padx=5, pady=2)
        self.video_path_var = tk.StringVar()
        self.video_path_entry = tk.Entry(tab_control, textvariable=self.video_path_var, width=50)
        self.video_path_entry.grid(row=0, column=1, padx=5, pady=2)
        self.browse_video_button = tk.Button(tab_control, text="Browse...", command=self.browse_video)
        self.browse_video_button.grid(row=0, column=2, padx=5, pady=2)

        # Output JSON Path
        tk.Label(tab_control, text="Output JSON Path:").grid(row=1, column=0, sticky=tk.W, padx=5, pady=2)
        self.json_path_var = tk.StringVar()
        self.json_path_entry = tk.Entry(tab_control, textvariable=self.json_path_var, width=50)
        self.json_path_entry.grid(row=1, column=1, padx=5, pady=2)
        self.browse_json_button = tk.Button(tab_control, text="Browse...", command=self.browse_json)
        self.browse_json_button.grid(row=1, column=2, padx=5, pady=2)

        # YOLO Model Selection
        tk.Label(tab_control, text="YOLO Model:").grid(row=2, column=0, sticky=tk.W, padx=5, pady=2)
        self.yolo_model_choices = ["yolo11n.pt", "yolo11s.pt", "yolo11m.pt", "yolo11l.pt", "yolo11x.pt", "Custom..."]
        self.yolo_model_path_var = tk.StringVar(tab_control)
        self.yolo_model_path_var.set(self.yolo_model_choices[0]) # Default selection
        self.yolo_model_option_menu = tk.OptionMenu(tab_control, self.yolo_model_path_var, *self.yolo_model_choices, command=self.handle_yolo_model_selection)
        self.yolo_model_option_menu.config(width=15)
        self.yolo_model_option_menu.grid(row=2, column=1, sticky=tk.W, padx=5, pady=2)

        # Entry for custom YOLO model path (initially hidden)
        self.custom_yolo_label = tk.Label(tab_control, text="Custom YOLO Path:")
        self.custom_yolo_path_var = tk.StringVar()
        self.custom_yolo_entry = tk.Entry(tab_control, textvariable=self.custom_yolo_path_var, width=30)
        self.custom_yolo_browse_button = tk.Button(tab_control, text="Browse...", command=self.browse_custom_yolo)

        # OpenRouter API Key
        tk.Label(tab_control, text="OpenRouter API Key:").grid(row=3, column=0, sticky=tk.W, padx=5, pady=2)
        self.api_key_var = tk.StringVar(value=DEFAULT_OPENROUTER_API_KEY)
        self.api_key_entry = tk.Entry(tab_control, textvariable=self.api_key_var, width=50, show="*")
        self.api_key_entry.grid(row=3, column=1, padx=5, pady=2)

        # OpenRouter API Base (Optional)
        tk.Label(tab_control, text="OpenRouter API Base (Optional):").grid(row=4, column=0, sticky=tk.W, padx=5, pady=2)
        self.api_base_var = tk.StringVar(value="https://openrouter.ai/api/v1")
        self.api_base_entry = tk.Entry(tab_control, textvariable=self.api_base_var, width=50)
        self.api_base_entry.grid(row=4, column=1, padx=5, pady=2)

        # Interval Seconds
        tk.Label(tab_control, text="Frame Interval (seconds):").grid(row=5, column=0, sticky=tk.W, padx=5, pady=2)
        self.interval_var = tk.DoubleVar(value=1.0)
        self.interval_entry = tk.Entry(tab_control, textvariable=self.interval_var, width=10)
        self.interval_entry.grid(row=5, column=1, sticky=tk.W, padx=5, pady=2)

        # Speed Units
        tk.Label(tab_control, text="Speed Units:").grid(row=6, column=0, sticky=tk.W, padx=5, pady=2)
        self.speed_unit_var = tk.StringVar(value="mph") # Default to mph
        self.speed_unit_options = ["mph", "km/h"]
        self.speed_unit_menu = tk.OptionMenu(tab_control, self.speed_unit_var, *self.speed_unit_options)
        self.speed_unit_menu.config(width=10)
        self.speed_unit_menu.grid(row=6, column=1, sticky=tk.W, padx=5, pady=2)

        # Debug Frames Checkbox
        self.debug_frames_var = tk.BooleanVar(value=True)
        self.debug_frames_check = tk.Checkbutton(tab_control, text="Preview", variable=self.debug_frames_var, command=self.toggle_debug_features)
        self.debug_frames_check.grid(row=7, column=0, columnspan=2, sticky=tk.W, padx=5, pady=2)

        # Status Label
        self.status_var = tk.StringVar()
        self.status_label = tk.Label(tab_control, textvariable=self.status_var)
        self.status_label.grid(row=8, column=0, columnspan=3, sticky=tk.W, padx=5, pady=5)

        # Run Button
        self.run_button = tk.Button(tab_control, text="Run Analysis", command=self.run_analysis, bg="lightblue")
        self.run_button.grid(row=9, column=0, columnspan=3, pady=10)

        # Debug Frame Display Area
        self.debug_frame_label = tk.Label(tab_control)
        self.debug_frame_label.grid(row=10, column=0, columnspan=5, pady=5, padx=5, sticky=tk.NSEW)

        # Scrubber Slider for Debug Frames
        self.scrubber_var = tk.IntVar(value=0)
        self.scrubber_slider = tk.Scale(tab_control, from_=0, to=0, orient=tk.HORIZONTAL,
                                        variable=self.scrubber_var, command=self.on_scrubber_update,
                                        state=tk.DISABLED, length=400)
        self.scrubber_slider.grid(row=11, column=0, columnspan=5, pady=5, padx=5, sticky=tk.EW)
        self.scrubber_slider.grid_remove() # Ensure slider starts hidden

        tab_control.grid_rowconfigure(10, weight=1) # Allow row with frame display to expand
        tab_control.grid_columnconfigure(1, weight=1) # Allow column with entry fields to expand

    def setup_analysis_tab(self, tab_control):
        # JSON File Path
        tk.Label(tab_control, text="Analysis JSON File:").grid(row=0, column=0, sticky=tk.W, padx=5, pady=2)
        self.analysis_json_path_var = tk.StringVar()
        self.analysis_json_entry = tk.Entry(tab_control, textvariable=self.analysis_json_path_var, width=50)
        self.analysis_json_entry.grid(row=0, column=1, padx=5, pady=2)
        self.browse_analysis_json_button = tk.Button(tab_control, text="Browse...", command=self.browse_analysis_json)
        self.browse_analysis_json_button.grid(row=0, column=2, padx=5, pady=2)

        # Load Plot Button
        self.load_plot_button = tk.Button(tab_control, text="Load Data & Plot", command=self.load_analysis_data, bg="lightgreen")
        self.load_plot_button.grid(row=1, column=0, columnspan=3, pady=10)

        # Video Player Container
        self.video_player_container = tk.Frame(tab_control, width=720, height=480, bg="black")
        self.video_player_container.grid(row=2, column=0, columnspan=3, pady=5, padx=5, sticky=tk.NSEW)
        self.video_player_container.grid_propagate(False) # Prevent container from resizing to content

        self.video_player_label = tk.Label(self.video_player_container, bg="black")
        self.video_player_label.pack(fill=tk.BOTH, expand=True) # Label fills the container
        
        # Play/Pause Button
        self.play_pause_button = tk.Button(tab_control, text="Play", command=self.toggle_play_pause, state=tk.DISABLED)
        self.play_pause_button.grid(row=3, column=0, sticky=tk.W, padx=5, pady=2)

        # Playhead Slider (now on its own conceptual row below button, but button is small)
        self.video_playhead_var = tk.DoubleVar(value=0)
        self.video_playhead_slider = tk.Scale(tab_control, from_=0, to=100, orient=tk.HORIZONTAL,
                                              variable=self.video_playhead_var, command=self.on_playhead_change,
                                              state=tk.DISABLED, length=400, resolution=0.1, showvalue=1)
        self.video_playhead_slider.grid(row=4, column=0, columnspan=3, sticky=tk.EW, padx=5, pady=2)
        
        # Color Strip Canvas
        self.color_strip_canvas = tk.Canvas(tab_control, height=10, bg="grey")
        self.color_strip_canvas.grid(row=5, column=0, columnspan=3, sticky=tk.EW, padx=5, pady=2)
        self.color_strip_canvas.bind("<Configure>", self.on_color_strip_resize)

        # Legend Frame
        self.legend_frame = tk.Frame(tab_control)
        self.legend_frame.grid(row=6, column=0, columnspan=3, sticky=tk.EW, padx=5, pady=2)
        self.populate_legend()
        
        # Plot Display Area
        self.plot_display_label = tk.Label(tab_control)
        self.plot_display_label.grid(row=7, column=0, columnspan=3, pady=5, padx=5, sticky=tk.NSEW)

        tab_control.grid_rowconfigure(2, weight=1) # Video player container expansion
        tab_control.grid_rowconfigure(7, weight=1) # Plot display expansion
        tab_control.grid_columnconfigure(0, weight=0) # Play button column does not expand
        tab_control.grid_columnconfigure(1, weight=1) # Slider/main content column expands
        tab_control.grid_columnconfigure(2, weight=0) # Potentially empty or small col, does not expand unless content pushes

    def populate_legend(self):
        # Clear existing legend items if any
        for widget in self.legend_frame.winfo_children():
            widget.destroy()
        
        tk.Label(self.legend_frame, text="Legend:").pack(side=tk.LEFT, padx=(0,5))
        for obj_type in self.OBJECT_PRIORITY:
            color = self.OBJECT_COLORS.get(obj_type)
            if color:
                color_patch = tk.Label(self.legend_frame, text="  ", bg=color, width=2)
                color_patch.pack(side=tk.LEFT, padx=(5,0))
                label = tk.Label(self.legend_frame, text=obj_type.replace("_", " ").title())
                label.pack(side=tk.LEFT, padx=(0,10))

    def reset_analysis_tab_ui(self):
        # Clear plot
        self.plot_display_label.config(image='')
        self.plot_display_label.image = None
        # Clear video player
        self.video_player_label.config(image='')
        self.video_player_label.image = None
        if self.video_cap:
            self.video_cap.release()
            self.video_cap = None
        self.is_video_playing = False
        self.play_pause_button.config(text="Play", state=tk.DISABLED)
        self.video_playhead_slider.config(state=tk.DISABLED, to=100)
        self.video_playhead_var.set(0)
        # Clear color strip
        if hasattr(self, 'color_strip_canvas'):
            self.color_strip_canvas.delete("all")
        self.color_strip_segments = [] # Reset stored segments
        self.last_known_color_strip_width = 0
        self.analysis_data = None

    def browse_analysis_json(self):
        filename = filedialog.askopenfilename(
            title="Select Analysis JSON File",
            filetypes=(("JSON files", "*.json"), ("All files", "*.*"))
        )
        if filename:
            self.analysis_json_path_var.set(filename)
            # Automatically attempt to load data if a file is selected
            self.load_analysis_data() 

    def load_analysis_data(self):
        self.reset_analysis_tab_ui() # Clear previous state
        json_path = self.analysis_json_path_var.get()
        if not json_path or not os.path.exists(json_path):
            messagebox.showerror("Error", "Valid analysis JSON file path is required.")
            return

        try:
            with open(json_path, 'r') as f:
                self.analysis_data = json.load(f)
        except Exception as e:
            messagebox.showerror("Error", f"Could not read or parse JSON file: {e}")
            self.analysis_data = None
            return

        # Load and display plot
        self._display_plot_from_data(json_path)

        # Load video
        video_file_path = self.analysis_data.get("video_file")
        if video_file_path and os.path.exists(video_file_path):
            self.load_video(video_file_path)
        elif video_file_path:
            messagebox.showwarning("Video Not Found", f"Video file specified in JSON not found: {video_file_path}")
        else:
            messagebox.showwarning("No Video Path", "No 'video_file' path found in JSON.")

    def _display_plot_from_data(self, json_path):
        if not self.analysis_data:
            return

        plot_dir = os.path.dirname(json_path)
        output_plot_path = os.path.join(plot_dir, "speed_plot.png")
        speed_unit_for_plot = self.analysis_data.get("processing_speed_unit", "mph")

        try:
            # Generate the plot, also calculate outliers and return the modified data.
            updated_analysis_data = process_video.generate_speed_plot(self.analysis_data, output_plot_path, speed_unit_for_plot)
            self.analysis_data = updated_analysis_data # Keep the updated data

            try: # Save potentially updated JSON (with new outlier flags)
                with open(json_path, 'w') as f_out:
                    json.dump(self.analysis_data, f_out, indent=4)
            except Exception as e_save:
                messagebox.showwarning("Save Warning", f"Could not save updated analysis data back to {json_path}: {e_save}")

            if os.path.exists(output_plot_path):
                img = Image.open(output_plot_path)
                max_width = 720 
                max_height = 480
                
                img_width, img_height = img.size
                if img_width == 0 or img_height == 0 : # Avoid division by zero for placeholder images
                    imgtk = ImageTk.PhotoImage(image=img)
                elif img_width > max_width or img_height > max_height:
                    ratio = min(max_width / img_width, max_height / img_height)
                    new_width = int(img_width * ratio)
                    new_height = int(img_height * ratio)
                    if new_width > 0 and new_height > 0:
                        img = img.resize((new_width, new_height), Image.LANCZOS)
                    imgtk = ImageTk.PhotoImage(image=img)
                else:
                    imgtk = ImageTk.PhotoImage(image=img)
                
                self.plot_display_label.image = imgtk
                self.plot_display_label.config(image=imgtk)
            else:
                messagebox.showwarning("Plot Not Found", f"Plot image was expected at {output_plot_path} but not found.")
        except Exception as e:
            messagebox.showerror("Plot Error", f"Error generating or displaying plot: {e}")
            self.plot_display_label.config(image='')
            self.plot_display_label.image = None

    def load_video(self, video_path):
        if self.video_cap:
            self.video_cap.release()
        
        self.video_cap = cv2.VideoCapture(video_path)
        if not self.video_cap.isOpened():
            messagebox.showerror("Video Error", f"Could not open video file: {video_path}")
            self.video_cap = None
            return

        self.video_fps = self.video_cap.get(cv2.CAP_PROP_FPS)
        self.video_frame_count = int(self.video_cap.get(cv2.CAP_PROP_FRAME_COUNT))

        if self.video_fps == 0 or self.video_frame_count == 0:
            messagebox.showerror("Video Error", "Video has 0 FPS or 0 frames. Cannot play.")
            self.video_cap.release()
            self.video_cap = None
            return
            
        self.total_video_duration_seconds = self.video_frame_count / self.video_fps
        self.video_playhead_slider.config(state=tk.NORMAL, to=self.total_video_duration_seconds)
        self.video_playhead_var.set(0)
        self.play_pause_button.config(state=tk.NORMAL)
        
        self.show_video_frame_at_time(0)
        self.prepare_and_draw_color_strip()

    def show_video_frame_at_time(self, time_seconds):
        if not self.video_cap or self.video_fps == 0:
            return

        frame_number = int(time_seconds * self.video_fps)
        if not (0 <= frame_number < self.video_frame_count):
            return

        self.video_cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        ret, frame_bgr = self.video_cap.read()

        if ret:
            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)

            # Resize to fit video player container (fixed size)
            player_width = 720 # Fixed width of video_player_container
            player_height = 480 # Fixed height of video_player_container

            img_width, img_height = img.size
            if img_width > 0 and img_height > 0: # Ensure valid image dimensions
                if img_width > player_width or img_height > player_height:
                    ratio = min(player_width / img_width, player_height / img_height)
                    new_width = int(img_width * ratio)
                    new_height = int(img_height * ratio)
                    if new_width > 0 and new_height > 0:
                        img = img.resize((new_width, new_height), Image.LANCZOS)
            
            self.video_player_image = ImageTk.PhotoImage(image=img)
            self.video_player_label.config(image=self.video_player_image)
            self.current_video_time_seconds = time_seconds
            self.video_playhead_var.set(time_seconds)
        else:
            if self.is_video_playing:
                self.pause_video()

    def on_playhead_change(self, value_str):
        if self.video_cap:
            # If video is playing, slider movement should pause it.
            if self.is_video_playing:
                self.pause_video()
            
            try:
                time_seconds = float(value_str)
                self.show_video_frame_at_time(time_seconds)
            except ValueError:
                pass # Should not happen with tk.DoubleVar

    def toggle_play_pause(self):
        if not self.video_cap:
            return
        if self.is_video_playing:
            self.pause_video()
        else:
            self.play_video()

    def play_video(self):
        if not self.video_cap or self.video_fps == 0:
            return
        self.is_video_playing = True
        self.play_pause_button.config(text="Pause")
        self.next_frame()

    def pause_video(self):
        self.is_video_playing = False
        if hasattr(self, 'play_pause_button'): # Check if widget exists
            self.play_pause_button.config(text="Play")

    def next_frame(self):
        if not self.is_video_playing or not self.video_cap or self.video_fps == 0:
            self.pause_video()
            return
        
        start_time_ns = time.perf_counter_ns()

        ret, frame_bgr = self.video_cap.read()

        if ret:
            self.current_video_time_seconds = self.video_cap.get(cv2.CAP_PROP_POS_FRAMES) / self.video_fps
            
            if self.current_video_time_seconds >= self.total_video_duration_seconds:
                self.current_video_time_seconds = self.total_video_duration_seconds
                self.show_video_frame_at_time(self.current_video_time_seconds)
                self.pause_video()
                self.video_playhead_var.set(self.current_video_time_seconds)
                return

            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)
            
            player_width = 720
            player_height = 480

            img_width, img_height = img.size
            if img_width > 0 and img_height > 0:
                if img_width > player_width or img_height > player_height:
                    ratio = min(player_width / img_width, player_height / img_height)
                    new_width = int(img_width * ratio)
                    new_height = int(img_height * ratio)
                    if new_width > 0 and new_height > 0:
                        img = img.resize((new_width, new_height), Image.LANCZOS)
            
            self.video_player_image = ImageTk.PhotoImage(image=img)
            self.video_player_label.config(image=self.video_player_image)
            self.video_playhead_var.set(self.current_video_time_seconds)

            end_time_ns = time.perf_counter_ns()
            processing_duration_ms = (end_time_ns - start_time_ns) / 1_000_000 # Convert ns to ms

            ideal_delay_ms = 1000 / self.video_fps
            actual_delay_ms = ideal_delay_ms - processing_duration_ms

            # Clamp delay to a minimum positive value (e.g., 1ms)
            # If processing takes longer than ideal_delay, this will be 1ms.
            delay_ms_to_use = max(1, int(actual_delay_ms))

            self.master.after(delay_ms_to_use, self.next_frame)
        else:
            self.pause_video()
            if self.total_video_duration_seconds > 0 and self.current_video_time_seconds < self.total_video_duration_seconds:
                 # If read failed before the expected end, try to set slider to current known time
                 self.video_playhead_var.set(self.current_video_time_seconds)
            elif self.total_video_duration_seconds > 0:
                 # If it seems like a natural end, set to total duration
                 self.current_video_time_seconds = self.total_video_duration_seconds
                 self.video_playhead_var.set(self.current_video_time_seconds)
    
    def prepare_and_draw_color_strip(self):
        if not self.analysis_data or not self.video_cap or self.total_video_duration_seconds == 0:
            if hasattr(self, 'color_strip_canvas'):
                self.color_strip_canvas.delete("all")
            self.color_strip_segments = []
            return

        frames_data = self.analysis_data.get("frames_processed", [])
        if not frames_data:
            if hasattr(self, 'color_strip_canvas'):
                self.color_strip_canvas.delete("all")
            self.color_strip_segments = []
            return
            
        proc_interval = self.analysis_data.get("processing_interval_seconds", 1.0)
        num_processed_frames = len(frames_data)
        
        current_segments = [] # Use a local variable for calculation
        for i in range(num_processed_frames):
            frame_entry = frames_data[i]
            timestamp = frame_entry.get("timestamp_seconds")
            if timestamp is None: continue

            secondary_objects = frame_entry.get("secondary_objects", [])
            current_color = None
            for obj_type in self.OBJECT_PRIORITY:
                if any(obj.get("label") == obj_type for obj in secondary_objects):
                    current_color = self.OBJECT_COLORS.get(obj_type)
                    break
            
            if current_color:
                start_time = float(timestamp)
                if i < num_processed_frames - 1 and frames_data[i+1].get("timestamp_seconds") is not None:
                    end_time = float(frames_data[i+1].get("timestamp_seconds"))
                else:
                    end_time = start_time + float(proc_interval)
                
                end_time = min(end_time, self.total_video_duration_seconds)

                if start_time < end_time:
                    current_segments.append((start_time, end_time, current_color))
        
        self.color_strip_segments = current_segments # Store the calculated segments
        self.draw_color_strip_on_canvas() # Perform initial draw

    def draw_color_strip_on_canvas(self):
        canvas = self.color_strip_canvas
        canvas.delete("all") 

        canvas_width = canvas.winfo_width()
        canvas_height = canvas.winfo_height()

        # Adjust for visual offset of the slider thumb/borders
        effective_track_start_x = self.SLIDER_VISUAL_OFFSET_X
        effective_track_width = canvas_width - (2 * self.SLIDER_VISUAL_OFFSET_X)

        if not self.color_strip_segments or self.total_video_duration_seconds == 0 or effective_track_width <= 0:
            self.last_known_color_strip_width = canvas_width # Still update this for resize logic
            return

        for start_time, end_time, color in self.color_strip_segments:
            # Scale positions based on the effective track width
            x1_on_track = (start_time / self.total_video_duration_seconds) * effective_track_width
            x2_on_track = (end_time / self.total_video_duration_seconds) * effective_track_width

            # Add the offset to get the canvas coordinates
            canvas_x1 = effective_track_start_x + x1_on_track
            canvas_x2 = effective_track_start_x + x2_on_track
            
            canvas.create_rectangle(canvas_x1, 0, canvas_x2, canvas_height, fill=color, outline="")
        
        self.last_known_color_strip_width = canvas_width

    def on_color_strip_resize(self, event):
        new_width = event.width
        if new_width != self.last_known_color_strip_width and new_width > 1:
            self.draw_color_strip_on_canvas()

    def toggle_debug_features(self):
        if self.debug_frames_var.get(): # Preview is ON
            self.debug_frame_label.grid() # Frame area is visible (might be empty)
            if self.debug_photo_images: # Frames exist (i.e., analysis has run)
                 new_max_idx = max(0, len(self.debug_photo_images)-1)
                 self.scrubber_slider.config(state=tk.NORMAL, to=new_max_idx)
                 self.scrubber_slider.grid() # Slider is visible and enabled
                 self.scrubber_slider.set(new_max_idx)
            else: # No frames yet (e.g., before first run or if run had no frames)
                 self.scrubber_slider.config(state=tk.DISABLED, to=0)
                 self.scrubber_slider.grid_remove() # Slider is hidden
        else: # Preview is OFF
            self.debug_frame_label.grid_remove()
            self.scrubber_slider.grid_remove()
            # Optionally clear the displayed image
            self.debug_frame_label.config(image='')
            self.debug_frame_label.imgtk = None 

    def update_debug_frame(self, frame_bgr):
        """Receives a BGR frame, converts, stores, and displays it."""
        if frame_bgr is None or not self.debug_frames_var.get():
            return
        try:
            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)
            imgtk = ImageTk.PhotoImage(image=img)
            
            self.debug_photo_images.append(imgtk) # Store the PhotoImage
            
            # Display the latest frame (live preview)
            self.debug_frame_label.imgtk = imgtk # Keep a reference!
            self.debug_frame_label.config(image=imgtk)
            
            if self.debug_photo_images:
                current_max_idx = len(self.debug_photo_images) -1
                if self.scrubber_slider.cget('to') < current_max_idx:
                    self.scrubber_slider.config(to=current_max_idx)
                self.scrubber_slider.set(current_max_idx)

            self.master.update_idletasks()
        except Exception as e:
            print(f"Error updating debug frame: {e}")

    def on_scrubber_update(self, value_str):
        """Called when the scrubber slider value changes."""
        if not self.debug_photo_images or not self.debug_frames_var.get():
            return
        try:
            frame_idx = int(value_str)
            if 0 <= frame_idx < len(self.debug_photo_images):
                imgtk = self.debug_photo_images[frame_idx]
                self.debug_frame_label.imgtk = imgtk # Keep a reference!
                self.debug_frame_label.config(image=imgtk)
        except ValueError:
            print(f"Invalid slider value: {value_str}")
        except Exception as e:
            print(f"Error in on_scrubber_update: {e}")

    def handle_yolo_model_selection(self, selection):
        if selection == "Custom...":
            self.custom_yolo_label.grid(row=2, column=2, sticky=tk.W, padx=5, pady=2) # Place next to dropdown
            self.custom_yolo_entry.grid(row=2, column=3, sticky=tk.W, padx=5, pady=2)
            self.custom_yolo_browse_button.grid(row=2, column=4, padx=5, pady=2)
        else:
            self.custom_yolo_label.grid_remove()
            self.custom_yolo_entry.grid_remove()
            self.custom_yolo_browse_button.grid_remove()
            self.custom_yolo_path_var.set("") # Clear custom path if not selected

    def browse_video(self):
        filename = filedialog.askopenfilename(
            title="Select Video File",
            filetypes=(("MP4 files", "*.mp4"), ("MOV files", "*.mov"), ("WEBM files", "*.webm"), ("AVI files", "*.avi"), ("All files", "*.*"))
        )
        if filename:
            self.video_path_var.set(filename)
            # Suggest an output JSON path based on video name
            base, _ = os.path.splitext(os.path.basename(filename))
            output_dir = f"output/{base}"
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            self.json_path_var.set(f"{output_dir}/analysis_results.json")

    def browse_json(self):
        filename = filedialog.asksaveasfilename(
            title="Save JSON Analysis As...",
            defaultextension=".json",
            filetypes=(("JSON files", "*.json"), ("All files", "*.*"))
        )
        if filename:
            self.json_path_var.set(filename)

    def browse_custom_yolo(self):
        filename = filedialog.askopenfilename(
            title="Select Custom YOLO Model File",
            filetypes=(("PyTorch Model files", "*.pt"), ("All files", "*.*"))
        )
        if filename:
            self.custom_yolo_path_var.set(filename)

    def run_analysis(self):
        video_path = self.video_path_var.get()
        json_path = self.json_path_var.get()
        
        selected_yolo_option = self.yolo_model_path_var.get()
        yolo_model_to_use = None
        if selected_yolo_option == "Custom...":
            yolo_model_to_use = self.custom_yolo_path_var.get() or None
            if yolo_model_to_use and not os.path.exists(yolo_model_to_use):
                messagebox.showwarning("Warning", f"Custom YOLO model path does not exist: {yolo_model_to_use}. Will attempt to use default.")
                yolo_model_to_use = None # Fallback
        elif selected_yolo_option: # It's one of the predefined names
            yolo_model_to_use = selected_yolo_option

        api_key = self.api_key_var.get() or None
        api_base = self.api_base_var.get() or None
        interval = self.interval_var.get()
        debug_frames = self.debug_frames_var.get()
        speed_unit = self.speed_unit_var.get()

        if not video_path or not os.path.exists(video_path):
            messagebox.showerror("Error", "Valid input video path is required.")
            return
        if not json_path:
            messagebox.showerror("Error", "Output JSON path is required.")
            return
        if not api_key:
            messagebox.showerror("Error", "OpenRouter API key is required.")
            return
        
        if interval <= 0:
            messagebox.showerror("Error", "Frame interval must be greater than 0.")
            return

        # Prepare for new analysis
        self.status_var.set("Processing... Please wait. This might take a while.")
        self.run_button.config(state=tk.DISABLED)
        self.debug_photo_images.clear() # Clear previously stored frames
        self.debug_frame_label.config(image='') # Clear display
        self.debug_frame_label.imgtk = None
        self.scrubber_slider.config(state=tk.DISABLED, from_=0, to=0, command=self.on_scrubber_update) # Reset slider
        self.scrubber_var.set(0)
        
        # Ensure slider is hidden before processing, frame label visibility depends on checkbox
        self.scrubber_slider.grid_remove()
        if self.debug_frames_var.get():
            self.debug_frame_label.grid()
        else:
            self.debug_frame_label.grid_remove()

        self.master.update_idletasks() # Ensure GUI updates before blocking call

        try:
            print(f"GUI calling analyze_video with:")
            print(f"  Video: {video_path}")
            print(f"  JSON Output: {json_path}")
            print(f"  Interval: {interval}s")
            print(f"  Debug: {debug_frames}")
            print(f"  YOLO Model: {yolo_model_to_use if yolo_model_to_use else 'Default'}")
            print(f"  API Key: {'Provided' if api_key else 'Not Provided (using env/default)'}")
            print(f"  API Base: {api_base if api_base else 'Default'}")
            print(f"  Speed Unit: {speed_unit}")

            processing_time = process_video.analyze_video(
                video_path,
                json_path,
                interval_seconds=interval,
                debug_frames=debug_frames,
                yolo_model_path=yolo_model_to_use,
                openrouter_api_key_override=api_key,
                openrouter_api_base_override=api_base,
                debug_frame_callback=self.update_debug_frame if debug_frames else None,
                speed_unit=speed_unit
            )
            success_message = f"Analysis complete in {processing_time:.1f}s! Output saved to {json_path}"
            self.status_var.set(success_message)
            messagebox.showinfo("Success", success_message)
        except Exception as e:
            self.status_var.set(f"Error during analysis: {e}")
            messagebox.showerror("Error", f"An error occurred: {e}")
            print(f"Error during analysis: {e}") 
        finally:
            self.run_button.config(state=tk.NORMAL)
            self.toggle_debug_features()

if __name__ == "__main__":
    root = tk.Tk()
    gui = HaneGUI(root)
    root.mainloop()