"""AareLC ML Studio - Main GUI Application."""
import os
import time
import json
import tkinter as tk
import urllib.parse
from tkinter import ttk
from tkinter import filedialog, messagebox, simpledialog
import cv2
import requests
from collections import deque
import numpy as np
from pathlib import Path
from datetime import datetime
from PIL import Image, ImageTk

# Import modular components
from src.gui.canvas_panel import MainCanvas
from src.gui.control_panel import ControlPanel
from src.gui.class_bank_mixin import ClassBankMixin
from src.gui.dataset_prep_mixin import DatasetPrepMixin
from src.gui.canvas_interaction_mixin import CanvasInteractionMixin
from src.gui.review_mixin import ReviewMixin
from src.core.zmq_client import ZMQStreamClient
from src.core.inference_client import InferenceClient
from src.core.db_client import DatabaseClient
from src.core.image_processor import ImageProcessor

AAREDB_SHARED_PASSWORD = os.getenv('AAREDB_SHARED_PASSWORD')
ZMQ_CONTROL_TOKEN = os.getenv('AARELC_ZMQ_CONTROL_TOKEN')


class InferenceGUI(ClassBankMixin, DatasetPrepMixin, CanvasInteractionMixin, ReviewMixin):
    """Main GUI application for ML inference and annotation."""

    def __init__(self, root):
        self.root = root
        self.root.title("AareLC ML Studio")
        self.root.geometry("1400x900")

        # --- Clients ---
        self.zmq_client = None
        self.inference_client = None
        self.db_client = None
        self.image_processor = ImageProcessor()

        # --- Settings ---
        self.pref_inference_url = tk.StringVar(value="http://mx-ml.psi.ch:8002")
        self.pref_db_url = tk.StringVar(value="https://mx-db-01.psi.ch")
        self.pref_download_url = tk.StringVar(value="http://localhost:8002/backend")
        self.pref_zmq_control_url = tk.StringVar(value="http://localhost:8090")
        self.pref_zmq_control_token = tk.StringVar(value=ZMQ_CONTROL_TOKEN)
        self.pref_shared_pw = tk.StringVar(value=AAREDB_SHARED_PASSWORD)
        self.pref_username = tk.StringVar(value=os.getenv('USER', 'anonymous'))
        self.zmq_det_url = tk.StringVar(value="tcp://localhost:9091")
        self.focus_enabled_var = tk.BooleanVar(value=True)
        self.pref_zmq_stream_options = [
            "tcp://localhost:9091",
            "tcp://x10sa-spark-01:9091",
            "tcp://sls-gpu-003:9089",
        ]
        self.pref_zmq_control_options = [
            "http://localhost:8090",
            "http://x10sa-spark-01:8090",
            "http://sls-gpu-003:8090",
        ]
        self.pref_download_url_options = [
            "http://localhost:8002/backend",
            "http://mx-db-01:8002/backend",
        ]

        # --- Sync filters ---
        self.sync_since = tk.StringVar(value="")
        self.sync_until = tk.StringVar(value="")
        self.sync_limit = tk.IntVar(value=5000)
        self.sync_include_unannotated = tk.BooleanVar(value=False)
        self.sync_dry_run = tk.BooleanVar(value=True)
        self.review_mode = False
        self.review_items = []
        self.review_index = -1
        self.current_review_item = None
        self.review_format = None
        self.review_coco_path = None
        self.review_coco_data = None

        # --- State ---
        self.current_raw_frame = None
        self.current_image_id = None  # Track current image ID from database
        self.current_annotation_id = None  # Track annotation ID for this image
        self.last_results = []
        self._last_overlap_pick = None
        self.selected_island_idx = 0
        self.selected_island_kind = "outer"
        self.history = deque(maxlen=20)  # NEW: Store last 20 states
        self.focus_history = deque(maxlen=500)
        self.stream_out_timestamps = deque(maxlen=240)
        self.stream_fps_in = 0.0
        self.stream_fps_out = 0.0
        self.stream_fps_pred = 0.0
        self._last_server_frame_id = None
        self._last_server_frame_ts = None
        # UI render cap for live stream; 0 means uncapped while still dropping stale frames.
        # Override with env var AARELC_STREAM_UI_MAX_FPS (e.g. 8, 15, 0).
        try:
            self.stream_ui_max_fps = float(os.getenv("AARELC_STREAM_UI_MAX_FPS", "0"))
        except ValueError:
            self.stream_ui_max_fps = 0.0
        if self.stream_ui_max_fps > 0:
            print(f"INFO: Stream UI FPS cap enabled at {self.stream_ui_max_fps:.1f} FPS", flush=True)
        else:
            print("INFO: Stream UI FPS cap disabled (uncapped)", flush=True)
        self._last_ui_render_ts = 0.0
        self._pending_stream_frame = None
        self._pending_stream_detections = None
        self._pending_stream_metadata = None
        self._last_stream_metadata = {}
        self._stream_render_scheduled = False

        # Status Bar Variables
        self.status_var = tk.StringVar(value="Ready")
        # Setup status truncation early to prevent X11 errors
        self._setup_status_truncation()
        self.local_output_dir = Path("data/custom_annotations")
        self.current_local_image_path = None
        self.class_set_source_path = None
        self.element_bank_path = self.local_output_dir / "element_bank.json"
        self.element_bank = []
        self.annotation_mode = "yolo_single_polygon"
        self.annotation_mode_labels = {
            "yolo_single_polygon": "YOLO single polygon",
            "coco_multi_polygon": "COCO multi-polygon",
        }

        # YOLO segmentation classes
        self.class_params = {
            "loop_all": {"low": 50, "high": 150, "clahe": 3.0, "morph": 7, "eps": 0.01},
            "pin": {"low": 100, "high": 200, "clahe": 2.0, "morph": 5, "eps": 0.002},
            "crystal": {"low": 50, "high": 150, "clahe": 3.0, "morph": 11, "eps": 0.01},
            "loop_face": {"low": 30, "high": 100, "clahe": 4.0, "morph": 3, "eps": 0.005},
            "ice": {"low": 50, "high": 150, "clahe": 3.0, "morph": 7, "eps": 0.01},
            "needle": {"low": 50, "high": 150, "clahe": 3.0, "morph": 7, "eps": 0.01}
        }

        # Color scheme mapping (BGR for OpenCV) for YOLO classes
        self.class_colors = {
            "loop_all": (0, 255, 0),  # Blue
            "pin": (0, 0, 255),  # Green
            "crystal": (255, 0, 0),  # Red
            "loop_face": (0, 255, 255),  # Yellow
            "ice": (255, 255, 255),  # White/Cyan
            "needle": (128, 0, 255)
        }

        # Map class names to YOLO class IDs (0-indexed for YOLO format)
        self.class_to_id = {
            "loop_all": 0,
            "pin": 1,
            "crystal": 2,
            "loop_face": 3,
            "ice": 4,
            "needle": 5,
        }
        self.id_to_class = {v: k for k, v in self.class_to_id.items()}
        self._class_color_palette = [
            (0, 255, 0),
            (0, 0, 255),
            (255, 0, 0),
            (0, 255, 255),
            (255, 255, 255),
            (128, 0, 255),
            (255, 165, 0),
            (255, 0, 255),
            (0, 128, 255),
            (128, 255, 0),
        ]
        self._visibility_trace_bound_classes = set()

        # Tkinter variables for control panel
        self.tk_vars = {
            'active_class': tk.StringVar(value=self._default_class_name()),
            'active_tool': tk.StringVar(value="select"), # select, brush, sam, or measure
            'canny_high': tk.IntVar(value=150),
            'morph_kernel': tk.IntVar(value=7),
            'poly_epsilon': tk.DoubleVar(value=0.01),
            'show_debug_edges': tk.BooleanVar(value=False),
            'plot_window_size': tk.IntVar(value=100),
            'annotation_line_width': tk.IntVar(value=2),
            'show_boxes': tk.BooleanVar(value=True),
            'show_segments': tk.BooleanVar(value=True),
            'brush_mode': tk.BooleanVar(value=False),
            'brush_type': tk.StringVar(value='background'),
            'brush_size': tk.IntVar(value=15),
            'segmentation_method': tk.StringVar(value='smart')  # Default to smart!
        }

        # Brush strokes storage: detection_idx -> {'bg': [], 'fg': []}
        self.brush_strokes = {}
        self.analysis_state = {
            'show_grid': False,
            'show_crosshair': True,
            'cursor': (0, 0),
            'line_start': None,
            'line_end': None,
            'line_dragging': False,
            'roi_start': None,
            'roi_end': None,
            'roi_dragging': False,
        }
        self.depth_backend_var = tk.StringVar(value=os.getenv("AARELC_DEPTH_BACKEND", "depthpro"))
        self.depth_max_side_var = tk.IntVar(value=int(os.getenv("AARELC_DEPTH_MAX_SIDE", "0")))
        self.depth_preview_window = None
        self.depth_preview_label = None
        self.depth_preview_photo = None

        # Initialize clients
        self._init_clients()

        # Setup UI
        self._setup_menus()
        self._setup_ui()
        self._load_default_element_bank_if_exists()

        # Status bar
        self.status_bar = tk.Label(
            self.root,
            textvariable=self.status_var,
            bd=1,
            relief=tk.SUNKEN,
            anchor=tk.W,
            font="fixed"
        )
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)

        # Setup parameter traces
        self._setup_traces()
        print(f"AAREDB_SHARED_PASSWORD={AAREDB_SHARED_PASSWORD}")
        print(f"AARELC_ZMQ_CONTROL_TOKEN={ZMQ_CONTROL_TOKEN}")

    def _init_clients(self):
        """Initialize client objects."""
        self.inference_client = InferenceClient(self.pref_inference_url.get())
        self.db_client = DatabaseClient(
            self.pref_db_url.get(), 
            self.pref_shared_pw.get(),
            download_base=self.pref_download_url.get()
        )

    def _setup_menus(self):
        """Setup menu bar."""
        menubar = tk.Menu(self.root)

        # File menu
        file_menu = tk.Menu(menubar, tearoff=0)
        file_menu.add_command(label="Load Image from Disk...", command=self.load_local_image)
        file_menu.add_command(label="Fetch Next from DB", command=self.fetch_db_image)
        file_menu.add_separator()
        file_menu.add_command(label="Set Local Save Folder...", command=self.choose_local_output_dir)
        file_menu.add_separator()
        file_menu.add_command(label="Add Class...", command=self.add_single_class)
        file_menu.add_command(label="Create New Class Set...", command=self.create_new_class_set)
        file_menu.add_command(label="Load Class Set from Disk...", command=self.load_class_set_from_disk)
        file_menu.add_command(label="Save Class Set As...", command=self.save_class_set_as)
        file_menu.add_separator()
        file_menu.add_command(label="Exit", command=self.root.quit)
        menubar.add_cascade(label="File", menu=file_menu)

        # Preferences menu
        pref_menu = tk.Menu(menubar, tearoff=0)
        pref_menu.add_command(label="Server Settings...", command=self.show_preferences)
        pref_menu.add_command(label="ZMQ Inference Control...", command=self.show_zmq_inference_control_dialog)
        menubar.add_cascade(label="Preferences", menu=pref_menu)

        # Sync menu
        sync_menu = tk.Menu(menubar, tearoff=0)
        sync_menu.add_command(
            label="Retrieve Images + Labels Locally...",
            command=self.show_sync_dialog,
        )
        sync_menu.add_command(
            label="Generate YOLO Segmentation Prep Script...",
            command=self.show_dataset_prep_script_dialog,
        )
        menubar.add_cascade(label="Sync", menu=sync_menu)

        # Review menu
        review_menu = tk.Menu(menubar, tearoff=0)
        review_menu.add_command(
            label="Open Review Folder...",
            command=self.open_review_folder,
        )
        review_menu.add_command(label="Previous Review Image (P)", command=self.prev_review_image)
        review_menu.add_command(label="Next Review Image (N)", command=self.next_review_image)
        review_menu.add_separator()
        review_menu.add_command(label="Save Current Review Annotation", command=self.save_review_annotation)
        review_menu.add_command(label="Exit Review Mode", command=self.exit_review_mode)
        menubar.add_cascade(label="Review", menu=review_menu)

        # Elements menu
        elements_menu = tk.Menu(menubar, tearoff=0)
        elements_menu.add_command(label="Add Selected to Bank...", command=self.add_selected_to_element_bank)
        elements_menu.add_command(label="Place Element from Bank...", command=self.place_element_from_bank)
        elements_menu.add_separator()
        elements_menu.add_command(label="Load Element Bank...", command=self.load_element_bank_from_disk)
        elements_menu.add_command(label="Save Element Bank As...", command=self.save_element_bank_as)
        menubar.add_cascade(label="Elements", menu=elements_menu)

        self.root.config(menu=menubar)

    def _setup_ui(self):
        """Setup main UI layout."""
        # 1. Top toolbars (two rows to avoid clipping on smaller screens)
        self.top_bar = tk.Frame(self.root, bg="#eee", pady=5, relief=tk.RAISED, bd=1)
        self.top_bar.pack(side=tk.TOP, fill=tk.X)
        self.top_bar_aux = tk.Frame(self.root, bg="#f3f3f3", pady=4, relief=tk.RAISED, bd=1)
        self.top_bar_aux.pack(side=tk.TOP, fill=tk.X)

        # Removed ZMQ Entry from here

        # NEW: Tool Palette Frame (Left side of toolbar)
        self.tool_frame = tk.Frame(self.top_bar, bg="#ddd", padx=5)
        self.tool_frame.pack(side=tk.LEFT, padx=5)

        self.select_tool_btn = tk.Radiobutton(
            self.tool_frame, text="Select", variable=self.tk_vars['active_tool'],
            value="select", indicatoron=0, width=8, padx=5, pady=2
        )
        self.select_tool_btn.pack(side=tk.LEFT)

        self.brush_tool_btn = tk.Radiobutton(
            self.tool_frame, text="Brush", variable=self.tk_vars['active_tool'],
            value="brush", indicatoron=0, width=8, padx=5, pady=2
        )
        self.brush_tool_btn.pack(side=tk.LEFT)

        self.sam_tool_btn = tk.Radiobutton(
            self.tool_frame, text="SAM", variable=self.tk_vars['active_tool'],
            value="sam", indicatoron=0, width=8, padx=5, pady=2
        )
        self.sam_tool_btn.pack(side=tk.LEFT)

        self.measure_tool_btn = tk.Radiobutton(
            self.tool_frame, text="Measure", variable=self.tk_vars['active_tool'],
            value="measure", indicatoron=0, width=10, padx=5, pady=2
        )
        self.measure_tool_btn.pack(side=tk.LEFT)
        
        tk.Label(self.top_bar, text="|", bg="#eee").pack(side=tk.LEFT, padx=5)
        
        # Connect button
        self.stream_btn = tk.Button(self.top_bar, text="Connect Live", command=self.toggle_stream, width=15)
        self.stream_btn.pack(side=tk.LEFT, padx=5)

        self.focus_btn = tk.Button(
            self.top_bar,
            text="Focus ON",
            command=self.toggle_focus_stream,
            width=12,
            bg="#17a2b8",
            fg="white"
        )
        self.focus_btn.pack(side=tk.LEFT, padx=5)

        # Database Fetch button (Next Random Image)
        self.db_next_btn = tk.Button(self.top_bar, text="Next Image", command=self.fetch_db_image, bg="#28a745",
                                     fg="white")
        self.db_next_btn.pack(side=tk.LEFT, padx=5)

        # Skip Image button
        self.skip_btn = tk.Button(self.top_bar, text="Skip", command=self.skip_current_image, bg="#ffc107",
                                  fg="black")
        self.skip_btn.pack(side=tk.LEFT, padx=2)

        # Predict button
        self.predict_btn = tk.Button(self.top_bar, text="Predict (/predict)", command=self.request_inference,
                                     bg="#007bff", fg="white")
        self.predict_btn.pack(side=tk.LEFT, padx=20)
        self.depth_btn = tk.Button(
            self.top_bar,
            text="Depth Map",
            command=self.request_depth_map,
            bg="#5a3d00",
            fg="white"
        )
        self.depth_btn.pack(side=tk.LEFT, padx=5)
        self.sam_boxes_btn = tk.Button(
            self.top_bar_aux,
            text="SAM from Boxes",
            command=self.refine_boxes_with_sam,
            bg="#17a2b8",
            fg="white"
        )
        self.sam_boxes_btn.pack(side=tk.LEFT, padx=5)
        tk.Label(self.top_bar_aux, text="Depth:", bg="#f3f3f3").pack(side=tk.LEFT, padx=(8, 2))
        self.depth_backend_combo = ttk.Combobox(
            self.top_bar_aux,
            textvariable=self.depth_backend_var,
            values=[
                "depthpro",
                "depth_anything_v1",
                "depth_anything_v1_trt",
                "depth_anything_v2",
                "depth_anything_v2_tiny",
                "depth_anything_v2_trt",
                "midas_v21_small",
            ],
            width=20,
            state="readonly"
        )
        self.depth_backend_combo.pack(side=tk.LEFT, padx=2)
        tk.Label(self.top_bar_aux, text="Max side:", bg="#f3f3f3").pack(side=tk.LEFT, padx=(6, 2))
        self.depth_max_side_combo = ttk.Combobox(
            self.top_bar_aux,
            textvariable=self.depth_max_side_var,
            values=[0, 384, 512, 640, 768, 1024],
            width=6,
            state="readonly"
        )
        self.depth_max_side_combo.pack(side=tk.LEFT, padx=2)

        # NEW: Zoom Slider in Toolbar
        tk.Label(self.top_bar_aux, text="Zoom:", bg="#f3f3f3").pack(side=tk.LEFT, padx=(10, 2))
        self.zoom_var = tk.DoubleVar(value=1.0)
        self.zoom_slider = tk.Scale(
            self.top_bar_aux, from_=0.1, to=5.0, resolution=0.1,
            orient=tk.HORIZONTAL, variable=self.zoom_var,
            showvalue=False, width=10, length=120, bg="#f3f3f3",
            command=self._on_zoom_slider_change
        )
        self.zoom_slider.pack(side=tk.LEFT, padx=5)

        # NEW: SAM refinement toggle
        self.use_sam_refine = tk.BooleanVar(value=False)
        self.sam_refine_check = tk.Checkbutton(
            self.top_bar_aux, text="Use SAM Refinement", variable=self.use_sam_refine,
            bg="#f3f3f3", font=('Arial', 9)
        )
        self.sam_refine_check.pack(side=tk.LEFT, padx=5)

        self.show_boxes_check = tk.Checkbutton(
            self.top_bar_aux,
            text="Show Boxes",
            variable=self.tk_vars['show_boxes'],
            bg="#f3f3f3",
            font=('Arial', 9)
        )
        self.show_boxes_check.pack(side=tk.LEFT, padx=4)

        self.show_segments_check = tk.Checkbutton(
            self.top_bar_aux,
            text="Show Segments",
            variable=self.tk_vars['show_segments'],
            bg="#f3f3f3",
            font=('Arial', 9)
        )
        self.show_segments_check.pack(side=tk.LEFT, padx=4)

        # 2. Main content area
        self.content_frame = tk.Frame(self.root)
        self.content_frame.pack(fill=tk.BOTH, expand=True)

        # 3. Control panel (right side)
        self.ctrl_panel = ControlPanel(
            self.content_frame,
            self.class_params,
            self.tk_vars,
            callbacks={
                'on_train_click': self.send_to_training,
                'on_clear_brush': self.clear_brush_strokes,
                'on_regenerate_edge': self.regenerate_selected_edge
            }
        )
        self.ctrl_panel.pack(side=tk.RIGHT, fill=tk.Y)

        # Initialize class visibility checkboxes
        self.ctrl_panel.update_class_visibility(list(self.class_params.keys()), self.class_colors)
        self.ctrl_panel.update_class_selector(list(self.class_params.keys()))
        self._bind_class_visibility_traces()

        # Add trace to redraw when brush mode changes
        self.tk_vars['brush_mode'].trace_add("write", lambda *args: self.redraw_current_frame())

        # 4. Canvas (center)
        self.canvas = MainCanvas(self.content_frame)
        # FIX: Use _on_canvas_interaction, which handles click/drag arguments
        self.canvas.set_redraw_callback(self._on_canvas_interaction)
        self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        self.add_box_btn = tk.Button(
            self.top_bar,
            text="Add Box",
            command=self.add_new_box,
            bg="#6c757d",
            fg="white"
        )
        self.add_box_btn.pack(side=tk.LEFT, padx=5)
        self.add_island_btn = tk.Button(
            self.top_bar_aux,
            text="Island",
            command=self.add_island_to_selected,
            bg="#6c757d",
            fg="white"
        )
        self.add_island_btn.pack(side=tk.LEFT, padx=2)
        self.add_hole_btn = tk.Button(
            self.top_bar_aux,
            text="Add Hole",
            command=self.add_hole_to_selected,
            bg="#6c757d",
            fg="white"
        )
        self.add_hole_btn.pack(side=tk.LEFT, padx=2)
        self.next_island_btn = tk.Button(
            self.top_bar_aux,
            text="Cycle Island",
            command=self.cycle_selected_island,
            bg="#f8f9fa"
        )
        self.next_island_btn.pack(side=tk.LEFT, padx=2)
        self.remove_island_btn = tk.Button(
            self.top_bar_aux,
            text="Remove Island",
            command=self.remove_selected_island,
            bg="#f8f9fa"
        )
        self.remove_island_btn.pack(side=tk.LEFT, padx=2)

        # NEW: Undo Button
        self.undo_btn = tk.Button(
            self.top_bar,
            text="Undo",
            command=self.undo_action,
            bg="#f8f9fa"
        )
        self.undo_btn.pack(side=tk.LEFT, padx=5)

    def _on_zoom_slider_change(self, val):
        """Sync slider value to canvas zoom level."""
        self.canvas.zoom_level = float(val)
        self.redraw_current_frame()

    def _setup_status_truncation(self):
        """Setup status message truncation to prevent X11 BadLength errors."""
        def truncate_status(*args):
            val = self.status_var.get()
            if len(val) > 100:
                truncated = val[:97] + "..."
                if val != truncated:
                    self.status_var.set(truncated)

        self.status_var.trace_add("write", truncate_status)

    def _setup_traces(self):
        """Setup variable traces for parameter changes."""
        self._after_id = None  # NEW: Track the debouncing timer
        for var_name in ['canny_high', 'morph_kernel', 'poly_epsilon']:
            self.tk_vars[var_name].trace_add("write", self.on_param_change)

        self.tk_vars['active_class'].trace_add("write", self.on_class_switch)
        self.tk_vars['active_tool'].trace_add("write", self._on_tool_switch)
        self.tk_vars['show_debug_edges'].trace_add("write", lambda *args: self.redraw_current_frame())
        self.tk_vars['plot_window_size'].trace_add("write", lambda *args: self.redraw_focus_plot())
        self.tk_vars['annotation_line_width'].trace_add("write", lambda *args: self.redraw_current_frame())
        self.tk_vars['show_boxes'].trace_add("write", lambda *args: self.redraw_current_frame())
        self.tk_vars['show_segments'].trace_add("write", lambda *args: self.redraw_current_frame())
        self.root.bind("<a>", lambda e: self.add_new_box())
        self.root.bind("<Delete>", self.delete_selected)
        self.root.bind("<Control-z>", lambda e: self.undo_action())  # NEW: Ctrl+Z shortcut
        self.root.bind("<Command-z>", lambda e: self.undo_action())  # NEW: Cmd+Z for macOS
        self.root.bind("<i>", lambda e: self.cycle_selected_island())
        # Review navigation shortcuts
        self.root.bind("<Right>", self._on_review_next_shortcut)
        self.root.bind("<Left>", self._on_review_prev_shortcut)
        self.root.bind("<n>", self._on_review_next_shortcut)
        self.root.bind("<p>", self._on_review_prev_shortcut)
        self.root.bind("<N>", self._on_review_next_shortcut)
        self.root.bind("<P>", self._on_review_prev_shortcut)
        self.root.bind("<g>", self._toggle_analysis_grid)
        self.root.bind("<c>", self._toggle_analysis_crosshair)
        self.root.bind("<x>", self._clear_measurements)

        # Global safety: Truncate any status message to 100 chars to avoid X11 BadLength
        # Moved to _setup_status_truncation for early setup

    def _default_class_name(self):
        """Return first available class name."""
        return next(iter(self.class_to_id.keys()), "class_0")

    @staticmethod
    def _default_class_params():
        """Default contour extraction params for newly created classes."""
        return {"low": 50, "high": 150, "clahe": 3.0, "morph": 7, "eps": 0.01}

    def _is_coco_multi_mode(self):
        """Whether current class set uses COCO multi-polygon mode."""
        return self.annotation_mode == "coco_multi_polygon"

    def _get_annotation_mode_label(self):
        return self.annotation_mode_labels.get(self.annotation_mode, self.annotation_mode)

    def _prompt_annotation_mode(self):
        """Ask user which annotation mode should be used for a new class set."""
        choice = messagebox.askyesnocancel(
            "Annotation Mode",
            (
                "Enable COCO multi-polygon mode for this class set?\n\n"
                "Yes: COCO multi-polygon (supports multi-island instances, local COCO save)\n"
                "No: YOLO single-polygon (current/default behavior)\n"
                "Cancel: abort"
            ),
            parent=self.root,
        )
        if choice is None:
            return None
        return "coco_multi_polygon" if choice else "yolo_single_polygon"

    def _det_get_polygons_relative(self, det):
        """Return list of relative polygons in (N,1,2) float32 format."""
        polys = det.get("polygons")
        out = []
        if isinstance(polys, list):
            for poly in polys:
                arr = np.array(poly, dtype=np.float32).reshape(-1, 1, 2)
                if len(arr) >= 3:
                    out.append(arr)
        if out:
            return out
        poly = det.get("poly")
        if poly is not None and len(poly) >= 3:
            out.append(np.array(poly, dtype=np.float32).reshape(-1, 1, 2))
        return out

    def _det_get_holes_relative(self, det):
        """Return list of relative hole polygons in (N,1,2) float32 format."""
        holes = det.get("holes")
        out = []
        if isinstance(holes, list):
            for poly in holes:
                arr = np.array(poly, dtype=np.float32).reshape(-1, 1, 2)
                if len(arr) >= 3:
                    out.append(arr)
        return out

    def _det_has_polygon(self, det):
        """True if detection has at least one explicit polygon island."""
        return len(self._det_get_polygons_relative(det)) > 0

    def _set_det_single_polygon(self, det, poly_arr):
        """Set detection polygon in both legacy and multi-poly fields."""
        arr = np.array(poly_arr, dtype=np.float32).reshape(-1, 1, 2)
        det["poly"] = arr
        det["polygons"] = [arr]

    def _set_det_polygons(self, det, polygon_list):
        """Set detection polygons list, updating legacy primary poly too."""
        cleaned = []
        for poly in polygon_list:
            arr = np.array(poly, dtype=np.float32).reshape(-1, 1, 2)
            if len(arr) >= 3:
                cleaned.append(arr)
        if cleaned:
            det["polygons"] = cleaned
            det["poly"] = cleaned[0]
        else:
            det.pop("polygons", None)
            det.pop("poly", None)

    def _set_det_holes(self, det, hole_list):
        """Set detection hole polygons list."""
        cleaned = []
        for poly in hole_list:
            arr = np.array(poly, dtype=np.float32).reshape(-1, 1, 2)
            if len(arr) >= 3:
                cleaned.append(arr)
        if cleaned:
            det["holes"] = cleaned
        else:
            det.pop("holes", None)

    def _det_get_absolute_polygons(self, det):
        """Return polygon islands as absolute image coordinates."""
        box = det.get("box", det)
        bx1 = float(box.get("x1", 0.0))
        by1 = float(box.get("y1", 0.0))
        polygons_abs = []
        for rel in self._det_get_polygons_relative(det):
            pts = rel.reshape(-1, 2).astype(np.float32) + [bx1, by1]
            if len(pts) >= 3:
                polygons_abs.append(pts)
        return polygons_abs

    def _det_get_absolute_holes(self, det):
        """Return hole polygons as absolute image coordinates."""
        box = det.get("box", det)
        bx1 = float(box.get("x1", 0.0))
        by1 = float(box.get("y1", 0.0))
        holes_abs = []
        for rel in self._det_get_holes_relative(det):
            pts = rel.reshape(-1, 2).astype(np.float32) + [bx1, by1]
            if len(pts) >= 3:
                holes_abs.append(pts)
        return holes_abs

    def _get_active_polygon(self, det):
        """Get currently selected polygon island for a detection."""
        rings = self._det_get_polygons_relative(det) if self.selected_island_kind == "outer" else self._det_get_holes_relative(det)
        if not rings:
            self.selected_island_kind = "outer"
            rings = self._det_get_polygons_relative(det)
            if not rings:
                return None
        idx = int(max(0, min(self.selected_island_idx, len(rings) - 1)))
        self.selected_island_idx = idx
        return rings[idx]

    def add_island_to_selected(self):
        """Add a new polygon island inside selected detection box."""
        idx = self.canvas.selected_obj_idx
        if idx == -1 or idx >= len(self.last_results):
            self.status_var.set("Select an instance first.")
            return
        det = self.last_results[idx]
        box = det.get('box', det)
        x1, y1, x2, y2 = float(box['x1']), float(box['y1']), float(box['x2']), float(box['y2'])
        w = max(6.0, x2 - x1)
        h = max(6.0, y2 - y1)
        cx, cy = w * 0.5, h * 0.5
        rw, rh = max(4.0, w * 0.2), max(4.0, h * 0.2)
        poly = np.array(
            [[[cx - rw, cy - rh]], [[cx + rw, cy - rh]], [[cx + rw, cy + rh]], [[cx - rw, cy + rh]]],
            dtype=np.float32
        )
        polygons = self._det_get_polygons_relative(det)
        polygons.append(poly)
        self._save_history()
        self._set_det_polygons(det, polygons)
        self.selected_island_idx = len(polygons) - 1
        self.canvas.selected_point_idx = -1
        self.redraw_current_frame()
        self.status_var.set(f"Added island {self.selected_island_idx + 1}/{len(polygons)}")

    def add_hole_to_selected(self):
        """Add a new hole polygon inside selected detection box."""
        idx = self.canvas.selected_obj_idx
        if idx == -1 or idx >= len(self.last_results):
            self.status_var.set("Select an instance first.")
            return
        det = self.last_results[idx]
        box = det.get('box', det)
        x1, y1, x2, y2 = float(box['x1']), float(box['y1']), float(box['x2']), float(box['y2'])
        w = max(10.0, x2 - x1)
        h = max(10.0, y2 - y1)
        cx, cy = w * 0.5, h * 0.5
        rw, rh = max(3.0, w * 0.12), max(3.0, h * 0.12)
        hole = np.array(
            [[[cx - rw, cy - rh]], [[cx + rw, cy - rh]], [[cx + rw, cy + rh]], [[cx - rw, cy + rh]]],
            dtype=np.float32
        )
        holes = self._det_get_holes_relative(det)
        holes.append(hole)
        self._save_history()
        self._set_det_holes(det, holes)
        self.selected_island_kind = "hole"
        self.selected_island_idx = len(holes) - 1
        self.canvas.selected_point_idx = -1
        self.redraw_current_frame()
        self.status_var.set(f"Added hole {self.selected_island_idx + 1}/{len(holes)}")

    def cycle_selected_island(self):
        """Cycle active island on selected detection."""
        idx = self.canvas.selected_obj_idx
        if idx == -1 or idx >= len(self.last_results):
            self.status_var.set("Select an instance first.")
            return
        det = self.last_results[idx]
        outers = self._det_get_polygons_relative(det)
        holes = self._det_get_holes_relative(det)
        if not outers and not holes:
            self.status_var.set("No rings on selected instance.")
            return

        if self.selected_island_kind == "outer":
            if outers:
                self.selected_island_idx = (self.selected_island_idx + 1) % len(outers)
                if self.selected_island_idx == 0 and holes:
                    self.selected_island_kind = "hole"
                    self.selected_island_idx = 0
            else:
                self.selected_island_kind = "hole"
                self.selected_island_idx = (self.selected_island_idx + 1) % len(holes)
        else:
            if holes:
                self.selected_island_idx = (self.selected_island_idx + 1) % len(holes)
                if self.selected_island_idx == 0 and outers:
                    self.selected_island_kind = "outer"
                    self.selected_island_idx = 0
            else:
                self.selected_island_kind = "outer"
                self.selected_island_idx = (self.selected_island_idx + 1) % len(outers)

        self.canvas.selected_point_idx = -1
        self.redraw_current_frame()
        total = len(self._det_get_polygons_relative(det)) if self.selected_island_kind == "outer" else len(self._det_get_holes_relative(det))
        self.status_var.set(f"Active {self.selected_island_kind} {self.selected_island_idx + 1}/{max(1, total)}")

    def remove_selected_island(self):
        """Remove active ring from selected detection."""
        idx = self.canvas.selected_obj_idx
        if idx == -1 or idx >= len(self.last_results):
            self.status_var.set("Select an instance first.")
            return
        det = self.last_results[idx]
        if self.selected_island_kind == "hole":
            holes = self._det_get_holes_relative(det)
            if not holes:
                self.status_var.set("No holes to remove.")
                return
            self._save_history()
            del holes[self.selected_island_idx]
            self.selected_island_idx = max(0, min(self.selected_island_idx, len(holes) - 1))
            self._set_det_holes(det, holes)
            if not holes:
                self.selected_island_kind = "outer"
                self.selected_island_idx = 0
        else:
            polygons = self._det_get_polygons_relative(det)
            if not polygons:
                self.status_var.set("No outer islands to remove.")
                return
            if len(polygons) == 1:
                messagebox.showwarning("Remove Island", "Instance has a single outer island. Delete instance if needed.")
                return
            self._save_history()
            del polygons[self.selected_island_idx]
            self.selected_island_idx = max(0, min(self.selected_island_idx, len(polygons) - 1))
            self._set_det_polygons(det, polygons)
        self.canvas.selected_point_idx = -1
        self.redraw_current_frame()
        self.status_var.set(f"Removed ring. Active {self.selected_island_kind} {self.selected_island_idx + 1}")

    # --- Event Handlers ---

    def on_class_switch(self, *args):
        """Load stored parameters when the user switches classes in the dropdown."""
        cls = self.tk_vars['active_class'].get()
        p = self.class_params.get(cls, self._default_class_params())

        # We use a flag to prevent on_param_change from being triggered
        # while we are just updating the sliders to match the selected class
        self._updating_ui = True
        self.tk_vars['canny_high'].set(p["high"])
        self.tk_vars['morph_kernel'].set(p["morph"])
        self.tk_vars['poly_epsilon'].set(p["eps"])
        self._updating_ui = False

        self.redraw_current_frame()

    def _on_tool_switch(self, *args):
        """Sync canvas interaction mode with selected tool."""
        is_measure = self.tk_vars['active_tool'].get() == "measure"
        self.canvas.suspend_pan = is_measure
        if is_measure:
            self.status_var.set("Measure mode: left-drag line, right-drag ROI, g/c/x toggles")
        self.redraw_current_frame()

    def on_param_change(self, *args):
        """Save slider values and CLEAR polygon ONLY for selected object."""
        if getattr(self, '_updating_ui', False):
            return

        # Cancel previous scheduled redraw if user is still sliding
        if self._after_id:
            self.root.after_cancel(self._after_id)

        def perform_update():
            cls = self.tk_vars['active_class'].get()
            self.class_params[cls].update({
                "high": self.tk_vars['canny_high'].get(),
                "morph": self.tk_vars['morph_kernel'].get(),
                "eps": self.tk_vars['poly_epsilon'].get()
            })

            # ONLY clear polygon for the SELECTED object, not all objects of this class
            if self.canvas.selected_obj_idx != -1 and self.canvas.selected_obj_idx < len(self.last_results):
                det = self.last_results[self.canvas.selected_obj_idx]
                label = det.get('name', det.get('label', self._default_class_name()))

                # Only clear if the selected object matches the active class
                if label == cls:
                    if 'poly' in det or 'polygons' in det:
                        det.pop('poly', None)
                        det.pop('polygons', None)
                        self.status_var.set(f"Parameters updated for selected {cls} - click Regenerate Edge")
                else:
                    self.status_var.set(f"Note: Selected object is '{label}', not '{cls}'. Switch class to edit.")
            else:
                self.status_var.set(f"Parameters saved for class '{cls}' - will apply to new detections")

            self.redraw_current_frame()
            self._after_id = None

        # Delay execution by 150ms to wait for slider to settle
        self._after_id = self.root.after(150, perform_update)

    def _save_history(self):
        """Save a deep copy of current detections to the history stack."""
        import copy
        # We store a copy of the list and the dictionaries inside
        snapshot = copy.deepcopy(self.last_results)
        self.history.append(snapshot)

    def undo_action(self):
        """Restore the last saved state from history."""
        if not self.history:
            self.status_var.set("Nothing to undo")
            return

        self.last_results = self.history.pop()
        self.canvas.selected_obj_idx = -1
        self.selected_island_idx = 0
        self.selected_island_kind = "outer"
        self.redraw_current_frame()
        self.status_var.set("Action undone")

    def _run_sam_inference(self, x, y):
        """Send a point to the DGX SAM backend and create a new detection."""
        if self.current_raw_frame is None: return

        try:
            self.root.config(cursor="watch")
            self.status_var.set(f"SAM: Segmenting at [{int(x)}, {int(y)}]...")
            self.root.update_idletasks()

            result = self.inference_client.segment_anything(self.current_raw_frame, x, y)
            poly_pts = result.get('polygon', [])

            if not poly_pts:
                self.status_var.set("SAM: No object found at that location.")
                return

            # Convert list of [x, y] to numpy format (N, 1, 2)
            poly_arr = np.array(poly_pts, dtype=np.float32).reshape(-1, 1, 2)

            # Calculate a bounding box for the polygon
            x_coords = [p[0] for p in poly_pts]
            y_coords = [p[1] for p in poly_pts]
            x1, y1, x2, y2 = min(x_coords), min(y_coords), max(x_coords), max(y_coords)

            # Polygons in our system are stored relative to box origin
            poly_relative = poly_arr - [x1, y1]

            new_det = {
                "box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
                "name": self.tk_vars['active_class'].get(),
                "poly": poly_relative.astype(np.float32),
                "polygons": [poly_relative.astype(np.float32)],
                "conf": result.get('score', 1.0)
            }

            self._save_history()
            self.last_results.append(new_det)
            self.canvas.selected_obj_idx = len(self.last_results) - 1
            self.selected_island_kind = "outer"
            self.redraw_current_frame()
            self.status_var.set(f"SAM: Object segmented successfully ({len(poly_pts)} points).")

        except Exception as e:
            messagebox.showerror("SAM Error", str(e))
            self.status_var.set("SAM inference failed.")
        finally:
            self.root.config(cursor="")

    # --- Image Loading ---

    def load_local_image(self):
        """Load image from local file system."""
        path = filedialog.askopenfilename(
            filetypes=[("Images", "*.jpg *.jpeg *.png *.bmp")]
        )
        if path:
            img = cv2.imread(path)
            if img is not None:
                self.exit_review_mode(silent=True)
                self._stop_streaming()
                self.current_raw_frame = img
                self.current_image_id = None
                self.current_annotation_id = None
                self.current_local_image_path = Path(path)
                self.last_results = []
                self.selected_island_idx = 0
                self.selected_island_kind = "outer"
                self.redraw_current_frame()
                self.status_var.set(f"Loaded: {path}")

    def fetch_db_image(self):
        """Fetch image from database server."""
        try:
            self.exit_review_mode(silent=True)
            self.root.config(cursor="watch")
            # Fetch returns (img, metadata)
            result = self.db_client.fetch_next_raw_image()

            # Check if result is just an image or (image, metadata)
            if isinstance(result, tuple):
                img, metadata = result
                self.current_image_id = metadata.get('image', {}).get('id')
                self.current_annotation_id = metadata.get('annotation_id')
            else:
                # Backwards compatibility - just image
                img = result
                self.current_image_id = None
                self.current_annotation_id = None

            self._stop_streaming()
            self.current_raw_frame = img
            self.current_local_image_path = None
            self.last_results = []  # Clear previous annotations
            self.selected_island_idx = 0
            self.selected_island_kind = "outer"
            self.redraw_current_frame()

            status_msg = f"Fetched image from database"
            if self.current_image_id:
                status_msg += f" (ID: {self.current_image_id})"
            self.status_var.set(status_msg)
        except Exception as e:
            messagebox.showerror("DB Error", str(e))
            self.status_var.set("Database fetch failed")
            import traceback
            traceback.print_exc()
        finally:
            self.root.config(cursor="")

    # --- Inference ---

    def request_inference(self):
        """Request /predict inference for the currently displayed frame snapshot."""
        if self.current_raw_frame is None:
            messagebox.showwarning("Warning", "No image loaded to predict.")
            return

        frame_for_predict = self.current_raw_frame.copy()
        from_stream = isinstance(getattr(self, '_last_stream_metadata', None), dict) and bool(self._last_stream_metadata)
        source_label = "live snapshot" if from_stream else "current image"

        try:
            self.root.config(cursor="watch")
            self.root.update_idletasks()

            self.status_var.set(f"Running /predict inference on {source_label}...")
            self.root.update_idletasks()
            self.last_results = self.inference_client.predict(frame_for_predict)

            # Optional SAM refinement runs AFTER a fresh /predict call.
            if self.use_sam_refine.get() and self.last_results:
                refine_indices = [
                    i for i, det in enumerate(self.last_results)
                    if not self._det_has_polygon(det)
                ]
                if refine_indices:
                    self.status_var.set(f"SAM: Auto-refining {len(refine_indices)} /predict detections...")
                    self.root.update_idletasks()
                    boxes_to_refine = []
                    for i in refine_indices:
                        b = self.last_results[i].get('box', self.last_results[i])
                        boxes_to_refine.append([int(b['x1']), int(b['y1']), int(b['x2']), int(b['y2'])])

                    sam_results = self.inference_client.segment_multi_box(frame_for_predict, boxes_to_refine)
                    for local_idx, res in enumerate(sam_results):
                        poly_pts = res.get('polygon', [])
                        if not poly_pts:
                            continue
                        det_idx = refine_indices[local_idx]
                        box = self.last_results[det_idx].get('box', self.last_results[det_idx])
                        poly_arr = np.array(poly_pts, dtype=np.float32).reshape(-1, 1, 2)
                        self._set_det_single_polygon(self.last_results[det_idx], poly_arr - [box['x1'], box['y1']])
                        self.last_results[det_idx].pop('suppress_auto_polygon', None)

            self.redraw_current_frame()
            self.root.update_idletasks()
            self.status_var.set(f"Prediction complete via /predict/. Found {len(self.last_results)} objects.")

        except Exception as e:
            messagebox.showerror("Inference Error", str(e))
            self.status_var.set("Inference failed")
        finally:
            self.root.config(cursor="")

    def request_depth_map(self):
        """Request FastDepth map and show it in a dedicated preview window."""
        if self.current_raw_frame is None:
            messagebox.showwarning("Warning", "No image loaded for depth analysis.")
            return

        try:
            self.root.config(cursor="watch")
            backend = self.depth_backend_var.get().strip() or "depthpro"
            depth_max_side = int(self.depth_max_side_var.get())
            self.status_var.set(f"Running depth inference ({backend}, max_side={depth_max_side})...")
            self.root.update_idletasks()
            depth_img = self.inference_client.predict_depth(
                self.current_raw_frame,
                colorize=True,
                depth_model=backend,
                depth_max_side=depth_max_side
            )
            self._show_depth_preview(depth_img)
            self.status_var.set(f"Depth map ready ({backend}, max_side={depth_max_side}).")
        except Exception as e:
            messagebox.showerror("Depth Error", str(e))
            self.status_var.set("Depth inference failed")
        finally:
            self.root.config(cursor="")

    def _show_depth_preview(self, depth_img):
        """Display the depth map in a non-modal window."""
        if self.depth_preview_window is None or not self.depth_preview_window.winfo_exists():
            self.depth_preview_window = tk.Toplevel(self.root)
            self.depth_preview_window.title("FastDepth Preview")
            self.depth_preview_window.geometry("900x700")
            self.depth_preview_window.configure(bg="#111")
            self.depth_preview_label = tk.Label(self.depth_preview_window, bg="#111")
            self.depth_preview_label.pack(fill=tk.BOTH, expand=True, padx=8, pady=8)

        h, w = depth_img.shape[:2]
        max_w, max_h = 1200, 900
        scale = min(max_w / max(1, w), max_h / max(1, h), 1.0)
        if scale < 1.0:
            depth_img = cv2.resize(depth_img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)

        if depth_img.ndim == 2:
            rgb = cv2.cvtColor(depth_img, cv2.COLOR_GRAY2RGB)
        else:
            rgb = cv2.cvtColor(depth_img, cv2.COLOR_BGR2RGB)

        pil = Image.fromarray(rgb)
        self.depth_preview_photo = ImageTk.PhotoImage(pil)
        self.depth_preview_label.configure(image=self.depth_preview_photo)
        self.depth_preview_window.lift()

    def refine_boxes_with_sam(self):
        """Send current rectangle boxes to SAM and update polygons."""
        if self.current_raw_frame is None:
            messagebox.showwarning("Warning", "No image loaded.")
            return
        if not self.last_results:
            messagebox.showwarning("Warning", "No boxes available. Add or predict boxes first.")
            return

        try:
            self.root.config(cursor="watch")
            refine_indices = [
                i for i, det in enumerate(self.last_results)
                if not self._det_has_polygon(det)
            ]
            if not refine_indices:
                self.status_var.set("SAM: All boxes already have polygons; nothing to refine.")
                return

            self.status_var.set(f"SAM: Refining {len(refine_indices)} boxes...")
            self.root.update_idletasks()

            boxes_to_refine = []
            for idx in refine_indices:
                det = self.last_results[idx]
                box = det.get('box', det)
                boxes_to_refine.append([int(box['x1']), int(box['y1']), int(box['x2']), int(box['y2'])])

            sam_results = self.inference_client.segment_multi_box(self.current_raw_frame, boxes_to_refine)
            refined = 0
            for local_idx, res in enumerate(sam_results):
                poly_pts = res.get('polygon', [])
                if not poly_pts:
                    continue
                det_idx = refine_indices[local_idx]
                box = self.last_results[det_idx].get('box', self.last_results[det_idx])
                poly_arr = np.array(poly_pts, dtype=np.float32).reshape(-1, 1, 2)
                self._set_det_single_polygon(self.last_results[det_idx], poly_arr - [box['x1'], box['y1']])
                self.last_results[det_idx].pop('suppress_auto_polygon', None)
                self.last_results[det_idx]['conf'] = res.get('score', self.last_results[det_idx].get('conf', 1.0))
                refined += 1

            self.redraw_current_frame()
            self.status_var.set(f"SAM: Refined {refined}/{len(refine_indices)} boxes")
        except Exception as e:
            messagebox.showerror("SAM Error", str(e))
            self.status_var.set("SAM box refinement failed")
        finally:
            self.root.config(cursor="")

    # --- Streaming ---

    def toggle_stream(self):
        """Toggle live streaming on/off."""
        if self.zmq_client and self.zmq_client.is_streaming():
            self._stop_streaming()
        else:
            self._start_streaming()

    def _start_streaming(self):
        """Start ZMQ live streaming."""
        self.stream_out_timestamps.clear()
        self.stream_fps_in = 0.0
        self.stream_fps_out = 0.0
        self.stream_fps_pred = 0.0
        self._last_server_frame_id = None
        self._last_server_frame_ts = None
        self.zmq_client = ZMQStreamClient(
            self.zmq_det_url.get(),
            on_frame_callback=self._on_zmq_frame,
            on_focus_callback=self._on_focus_update
        )
        if not self.focus_enabled_var.get():
            self.zmq_client.disable_focus()
        self.zmq_client.start()
        self.stream_btn.config(text="Stop Live", bg="red")
        cuda_text = "ON" if getattr(self.zmq_client, 'cuda_enabled', False) else "OFF"
        print(f"INFO: GUI stream CUDA focus: {cuda_text}", flush=True)
        self.status_var.set(f"Streaming... CUDA focus: {cuda_text}")

    def _stop_streaming(self):
        """Stop ZMQ live streaming."""
        if self.zmq_client:
            self.zmq_client.stop()
            self.zmq_client = None
        self.stream_out_timestamps.clear()
        self.stream_fps_in = 0.0
        self.stream_fps_out = 0.0
        self.stream_fps_pred = 0.0
        self._last_server_frame_id = None
        self._last_server_frame_ts = None
        self._pending_stream_frame = None
        self._pending_stream_detections = None
        self._pending_stream_metadata = None
        self._last_stream_metadata = {}
        self._stream_render_scheduled = False
        self.stream_btn.config(text="Connect Live", bg="blue")
        self.status_var.set("Streaming stopped")

    def toggle_focus_stream(self):
        """Enable/disable focus-score computation for live stream."""
        new_state = not self.focus_enabled_var.get()
        self.focus_enabled_var.set(new_state)

        if self.zmq_client and self.zmq_client.is_streaming():
            if new_state:
                self.zmq_client.enable_focus()
            else:
                self.zmq_client.disable_focus()

        if new_state:
            self.focus_btn.config(text="🎯 Focus ON", bg="#17a2b8")
            self.status_var.set("Focus scoring enabled")
        else:
            self.focus_btn.config(text="🎯 Focus OFF", bg="#6c757d")
            self.status_var.set("Focus scoring disabled")

    def _on_zmq_frame(self, frame, detections, metadata=None):
        """Callback for new ZMQ frame.
        Keep only newest frame and render in UI at a capped rate to stay live."""
        now = time.time()

        if isinstance(metadata, dict):
            source_fps = metadata.get('source_fps')
            if isinstance(source_fps, (int, float)):
                self.stream_fps_in = float(source_fps)
            else:
                frame_id = metadata.get('frame_id')
                if isinstance(frame_id, int):
                    if self._last_server_frame_id is not None and self._last_server_frame_ts is not None:
                        dt = now - self._last_server_frame_ts
                        df = frame_id - self._last_server_frame_id
                        if dt > 0 and df >= 0:
                            self.stream_fps_in = float(df) / float(dt)
                    self._last_server_frame_id = frame_id
                    self._last_server_frame_ts = now

            prediction_fps = metadata.get('prediction_fps')
            if isinstance(prediction_fps, (int, float)):
                self.stream_fps_pred = float(prediction_fps)

        # Overwrite pending payload: this intentionally drops stale frames.
        self._pending_stream_frame = frame
        self._pending_stream_detections = detections
        self._pending_stream_metadata = metadata

        if not self._stream_render_scheduled:
            self._stream_render_scheduled = True
            self.root.after(0, self._flush_latest_stream_frame)

    def _flush_latest_stream_frame(self):
        """Render latest pending stream frame, dropping any intermediate backlog."""
        self._stream_render_scheduled = False
        if self._pending_stream_frame is None:
            return

        now = time.time()
        if self.stream_ui_max_fps > 0:
            min_interval = 1.0 / float(self.stream_ui_max_fps)
            dt = now - self._last_ui_render_ts
            if dt < min_interval:
                delay_ms = max(1, int((min_interval - dt) * 1000))
                self._stream_render_scheduled = True
                self.root.after(delay_ms, self._flush_latest_stream_frame)
                return

        frame = self._pending_stream_frame
        detections = self._pending_stream_detections if self._pending_stream_detections is not None else []
        metadata = self._pending_stream_metadata if isinstance(self._pending_stream_metadata, dict) else {}
        self._pending_stream_frame = None
        self._pending_stream_detections = None
        self._pending_stream_metadata = None

        self.current_raw_frame = frame
        self.last_results = detections
        self._last_stream_metadata = metadata
        self.update_canvas(frame, metadata=metadata)
        self._last_ui_render_ts = now

        # Render FPS is measured on UI flushes, not on received packets.
        self.stream_out_timestamps.append(now)
        while self.stream_out_timestamps and (now - self.stream_out_timestamps[0]) > 1.0:
            self.stream_out_timestamps.popleft()
        self.stream_fps_out = float(len(self.stream_out_timestamps))

    def _on_focus_update(self, score):
        """Callback for focus score update."""
        self.focus_history.append(score)
        self.root.after_idle(self.redraw_focus_plot)

    # --- Display & Canvas Interaction ---

    def regenerate_selected_edge(self):
        """Force regeneration of polygon for selected detection using brush strokes or SAM."""
        if self.canvas.selected_obj_idx == -1:
            self.status_var.set("No detection selected")
            return

        if self.canvas.selected_obj_idx >= len(self.last_results):
            return

        det = self.last_results[self.canvas.selected_obj_idx]
        box = det.get('box', det)

        # --- NEW: SAM Refinement Path ---
        if self.use_sam_refine.get():
            try:
                self.root.config(cursor="watch")
                self.status_var.set("SAM: Refining based on box...")
                self.root.update_idletasks()

                res = self.inference_client.segment_anything_box(
                    self.current_raw_frame,
                    box['x1'], box['y1'], box['x2'], box['y2']
                )
                poly_pts = res.get('polygon', [])
                if poly_pts:
                    poly_arr = np.array(poly_pts, dtype=np.float32).reshape(-1, 1, 2)
                    self._set_det_single_polygon(det, poly_arr - [box['x1'], box['y1']])
                    det.pop('suppress_auto_polygon', None)
                    self.status_var.set("SAM: Edge refined successfully")
                else:
                    self.status_var.set("SAM: Could not find object in box")
            except Exception as e:
                self.status_var.set(f"SAM Error: {str(e)}")
            finally:
                self.root.config(cursor="")
                self.redraw_current_frame()
            return

        # Check if we have manual mask
        has_manual_mask = 'manual_mask' in det
        print(f"Regenerating edge: has_manual_mask={has_manual_mask}")

        if has_manual_mask:
            manual_mask = det['manual_mask']
            has_bg = np.any(manual_mask == 0)
            has_fg = np.any(manual_mask == 1)
            print(f"  Manual mask: has_bg={has_bg}, has_fg={has_fg}")

        # Delete polygon to force regeneration
        if 'poly' in det or 'polygons' in det:
            det.pop('poly', None)
            det.pop('polygons', None)
            print("  Deleted existing polygon(s)")
        det.pop('suppress_auto_polygon', None)

        # Trigger redraw which will regenerate the polygon
        self.redraw_current_frame()

        # Check if manual mask still exists after redraw
        has_manual_mask_after = 'manual_mask' in det
        print(f"  After redraw: has_manual_mask={has_manual_mask_after}")

        if has_manual_mask_after:
            self.status_var.set("Edge regenerated - paint more or click again to refine")
        else:
            self.status_var.set("Edge regenerated")

    def clear_brush_strokes(self):
        """Clear brush strokes for selected detection."""
        if self.canvas.selected_obj_idx == -1:
            self.status_var.set("No detection selected")
            return

        if self.canvas.selected_obj_idx >= len(self.last_results):
            return

        self._save_history()
        det = self.last_results[self.canvas.selected_obj_idx]

        # Clear manual mask and polygon
        if 'manual_mask' in det:
            del det['manual_mask']
        det.pop('poly', None)
        det.pop('polygons', None)

        self.redraw_current_frame()
        self.status_var.set("Brush strokes cleared for selected detection")

    def redraw_focus_plot(self):
        """Redraw focus history plot."""
        window = self.tk_vars['plot_window_size'].get()
        self.ctrl_panel.update_focus_plot(list(self.focus_history), window)

    # --- Training & Annotation Submission ---

    def _convert_to_yolo_segmentation(self):
        """Convert detections to YOLO segmentation format.
        YOLO segmentation format per line:
        class_id x1 y1 x2 y2 x3 y3 ... (normalized coordinates 0-1)
        """
        if self.current_raw_frame is None or not self.last_results:
            return []

        img_h, img_w = self.current_raw_frame.shape[:2]
        yolo_lines = []

        for det in self.last_results:
            # Get class
            class_name = det.get('name', det.get('label', self._default_class_name()))
            class_id = self.class_to_id.get(class_name, 0)

            box = det.get('box', det)
            polygons = self._det_get_absolute_polygons(det)

            if polygons:
                for points in polygons:
                    normalized_points = []
                    for pt in points:
                        nx = pt[0] / img_w
                        ny = pt[1] / img_h
                        normalized_points.extend([nx, ny])
                    line = f"{class_id} " + " ".join(f"{coord:.6f}" for coord in normalized_points)
                    yolo_lines.append(line)
            else:
                x1, y1 = box.get('x1', 0), box.get('y1', 0)
                x2, y2 = box.get('x2', 0), box.get('y2', 0)
                cx = ((x1 + x2) / 2) / img_w
                cy = ((y1 + y2) / 2) / img_h
                w = (x2 - x1) / img_w
                h = (y2 - y1) / img_h
                line = f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}"
                yolo_lines.append(line)

        return yolo_lines

    def _save_local_annotation_files(self, file_stem, yolo_seg_lines, yolo_det_lines):
        """Persist image + YOLO labels under the configured local output folder."""
        root = Path(self.local_output_dir)
        targets = {
            "segmentation": yolo_seg_lines,
            "detection": yolo_det_lines,
        }
        for task_type, lines in targets.items():
            base_dir = root / task_type
            img_dir = base_dir / "images"
            lbl_dir = base_dir / "labels"
            img_dir.mkdir(parents=True, exist_ok=True)
            lbl_dir.mkdir(parents=True, exist_ok=True)
            cv2.imwrite(str(img_dir / f"{file_stem}.jpg"), self.current_raw_frame)
            with open(lbl_dir / f"{file_stem}.txt", "w", encoding="utf-8") as f:
                f.write("\n".join(lines) + ("\n" if lines else ""))

        classes_txt = root / "classes.txt"
        class_names = [name for _, name in sorted((idx, name) for name, idx in self.class_to_id.items())]
        classes_txt.write_text("\n".join(class_names) + ("\n" if class_names else ""), encoding="utf-8")
        (root / "annotation_mode.txt").write_text(f"{self.annotation_mode}\n", encoding="utf-8")

    @staticmethod
    def _polygon_area(points):
        """Compute polygon area for COCO annotations."""
        arr = np.array(points, dtype=np.float32).reshape(-1, 1, 2)
        return float(abs(cv2.contourArea(arr)))

    @staticmethod
    def _rle_encode_binary_mask(mask):
        """Encode uint8 binary mask into COCO uncompressed RLE."""
        pixels = np.asarray(mask, dtype=np.uint8).flatten(order="F")
        counts = []
        prev = 0
        run_len = 0
        for pix in pixels:
            val = 1 if pix else 0
            if val == prev:
                run_len += 1
            else:
                counts.append(run_len)
                run_len = 1
                prev = val
        counts.append(run_len)
        return {"size": [int(mask.shape[0]), int(mask.shape[1])], "counts": counts}

    @staticmethod
    def _rle_decode_binary_mask(seg):
        """Decode COCO uncompressed RLE dict into uint8 binary mask."""
        size = seg.get("size", [])
        counts = seg.get("counts", [])
        if not isinstance(size, list) or len(size) != 2 or not isinstance(counts, list):
            return None
        h, w = int(size[0]), int(size[1])
        flat = np.zeros(h * w, dtype=np.uint8)
        idx = 0
        val = 0
        for run in counts:
            run = int(run)
            if run <= 0:
                val = 1 - val
                continue
            end = min(idx + run, flat.size)
            if val == 1:
                flat[idx:end] = 1
            idx = end
            if idx >= flat.size:
                break
            val = 1 - val
        return flat.reshape((h, w), order="F")

    def _build_mask_from_rings(self, img_h, img_w, polygons_abs, holes_abs):
        """Rasterize outer/hole rings to binary mask (outer minus holes)."""
        mask = np.zeros((img_h, img_w), dtype=np.uint8)
        for poly in polygons_abs:
            pts = np.round(poly).astype(np.int32).reshape(-1, 1, 2)
            cv2.fillPoly(mask, [pts], 1)
        for hole in holes_abs:
            pts = np.round(hole).astype(np.int32).reshape(-1, 1, 2)
            cv2.fillPoly(mask, [pts], 0)
        return mask

    def _save_local_coco_annotation_files(self, file_stem):
        """Persist current image+annotations in COCO instance format."""
        root = Path(self.local_output_dir) / "coco"
        images_dir = root / "images"
        images_dir.mkdir(parents=True, exist_ok=True)

        image_filename = f"{file_stem}.jpg"
        cv2.imwrite(str(images_dir / image_filename), self.current_raw_frame)

        coco_path = root / "annotations.json"
        if coco_path.exists():
            try:
                coco = json.loads(coco_path.read_text(encoding="utf-8"))
            except Exception:
                coco = {}
        else:
            coco = {}

        coco.setdefault("info", {"description": "AareLC local COCO export"})
        coco.setdefault("licenses", [])
        coco.setdefault("images", [])
        coco.setdefault("annotations", [])
        coco["categories"] = [
            {"id": idx + 1, "name": name}
            for name, idx in sorted(self.class_to_id.items(), key=lambda kv: kv[1])
        ]
        coco["meta"] = {"annotation_mode": self.annotation_mode}

        img_h, img_w = self.current_raw_frame.shape[:2]
        existing_img = next((img for img in coco["images"] if img.get("file_name") == image_filename), None)
        if existing_img:
            image_id = existing_img["id"]
            existing_img.update({"width": int(img_w), "height": int(img_h)})
        else:
            image_id = (max((img.get("id", 0) for img in coco["images"]), default=0) + 1)
            coco["images"].append({
                "id": image_id,
                "file_name": image_filename,
                "width": int(img_w),
                "height": int(img_h),
            })

        coco["annotations"] = [ann for ann in coco["annotations"] if ann.get("image_id") != image_id]
        next_ann_id = max((ann.get("id", 0) for ann in coco["annotations"]), default=0) + 1

        for det in self.last_results:
            class_name = det.get('name', det.get('label', self._default_class_name()))
            category_id = self.class_to_id.get(class_name, 0) + 1
            box = det.get('box', det)
            x1, y1 = float(box['x1']), float(box['y1'])
            x2, y2 = float(box['x2']), float(box['y2'])
            bw, bh = max(0.0, x2 - x1), max(0.0, y2 - y1)

            polygons_abs = self._det_get_absolute_polygons(det)
            holes_abs = self._det_get_absolute_holes(det)
            if not polygons_abs:
                polygons_abs = [np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype=np.float32)]

            if holes_abs:
                mask = self._build_mask_from_rings(img_h, img_w, polygons_abs, holes_abs)
                ys, xs = np.where(mask > 0)
                if len(xs) == 0 or len(ys) == 0:
                    continue
                min_x, max_x = float(np.min(xs)), float(np.max(xs))
                min_y, max_y = float(np.min(ys)), float(np.max(ys))
                segmentation_payload = self._rle_encode_binary_mask(mask)
                area_total = float(np.sum(mask > 0))
                bbox_payload = [min_x, min_y, max_x - min_x + 1.0, max_y - min_y + 1.0]
            else:
                segmentation = []
                area_total = 0.0
                for poly in polygons_abs:
                    flat = [float(v) for point in poly for v in point]
                    if len(flat) >= 6:
                        segmentation.append(flat)
                        area_total += self._polygon_area(poly)
                if not segmentation:
                    continue
                segmentation_payload = segmentation
                bbox_payload = [x1, y1, bw, bh]

            coco["annotations"].append({
                "id": next_ann_id,
                "image_id": image_id,
                "category_id": category_id,
                "segmentation": segmentation_payload,
                "area": float(area_total),
                "bbox": bbox_payload,
                "iscrowd": 0,
            })
            next_ann_id += 1

        coco_path.write_text(json.dumps(coco, indent=2), encoding="utf-8")
        (root / "annotation_mode.txt").write_text(f"{self.annotation_mode}\n", encoding="utf-8")

    def send_to_training(self):
        """Save YOLO segmentation annotations to database and locally for training."""
        if self.review_mode:
            self.save_review_annotation()
            return

        if self.current_raw_frame is None:
            messagebox.showwarning("Warning", "No image loaded.")
            return

        if not self.last_results:
            messagebox.showwarning("Warning", "No annotations to save.")
            return

        try:
            img_h, img_w = self.current_raw_frame.shape[:2]

            if self._is_coco_multi_mode():
                if self.current_local_image_path is not None:
                    file_stem = self.current_local_image_path.stem
                elif self.current_image_id is not None:
                    file_stem = f"img_{self.current_image_id}"
                else:
                    file_stem = datetime.now().strftime("manual_%Y%m%d_%H%M%S")

                confirm_msg = (
                    f"Save {len(self.last_results)} annotations in COCO multi-polygon mode "
                    f"to local folder only?\n\nMode: {self._get_annotation_mode_label()}"
                )
                if not messagebox.askyesno("Confirm", confirm_msg):
                    return

                self.root.config(cursor="watch")
                self._save_local_coco_annotation_files(file_stem)
                self.status_var.set(f"Saved COCO annotation: {file_stem}")
                messagebox.showinfo(
                    "Success",
                    (
                        f"Saved '{file_stem}' in COCO format under:\n{self.local_output_dir / 'coco'}\n\n"
                        "File: annotations.json (multi-polygon capable)"
                    ),
                )
                return

            if any(self._det_get_holes_relative(det) for det in self.last_results):
                if not messagebox.askyesno(
                    "YOLO Limitation",
                    (
                        "Some instances contain hole polygons, but YOLO txt export cannot encode holes.\n"
                        "Holes will be ignored in YOLO export.\n\nContinue?"
                    ),
                ):
                    return

            # 1. Prepare YOLO data (Normalized)
            segments = []
            yolo_seg_lines = [] # For polygons
            yolo_det_lines = [] # For bounding boxes

            for det in self.last_results:
                class_name = det.get('name', det.get('label', self._default_class_name()))
                class_id = self.class_to_id.get(class_name, 0)
                box = det.get('box', det)

                # --- 1a. Generate Detection Line (cx, cy, w, h) ---
                x1, y1, x2, y2 = float(box['x1']), float(box['y1']), float(box['x2']), float(box['y2'])
                cx = ((x1 + x2) / 2) / img_w
                cy = ((y1 + y2) / 2) / img_h
                bw = (x2 - x1) / img_w
                bh = (y2 - y1) / img_h
                yolo_det_lines.append(f"{class_id} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")

                # --- 1b. Generate Segmentation Line (x1, y1, x2, y2...) ---
                polygons = self._det_get_absolute_polygons(det)
                if polygons:
                    for points in polygons:
                        normalized_points = []
                        for pt in points:
                            nx, ny = float(pt[0] / img_w), float(pt[1] / img_h)
                            normalized_points.extend([nx, ny])

                        segments.append({
                            "class_id": class_id,
                            "points": normalized_points,
                            "box": [cx, cy, bw, bh]
                        })
                        pts_str = " ".join(f"{p:.6f}" for p in normalized_points)
                        yolo_seg_lines.append(f"{class_id} {pts_str}")
                else:
                    # Fallback: Use box corners as a 4-point polygon
                    norm_poly = [x1 / img_w, y1 / img_h, x2 / img_w, y1 / img_h, x2 / img_w, y2 / img_h, x1 / img_w,
                                 y2 / img_h]
                    segments.append({
                        "class_id": class_id,
                        "points": norm_poly,
                        "box": [cx, cy, bw, bh]
                    })
                    pts_str = " ".join(f"{p:.6f}" for p in norm_poly)
                    yolo_seg_lines.append(f"{class_id} {pts_str}")

            save_mode = "database + local training sets" if self.current_image_id is not None else "local folder only"
            if not messagebox.askyesno("Confirm", f"Save {len(segments)} annotations to {save_mode}?"):
                return

            self.root.config(cursor="watch")

            if self.current_image_id is None:
                if self.current_local_image_path is not None:
                    file_stem = self.current_local_image_path.stem
                else:
                    file_stem = datetime.now().strftime("manual_%Y%m%d_%H%M%S")

                self._save_local_annotation_files(file_stem, yolo_seg_lines, yolo_det_lines)
                self.status_var.set(f"Saved local annotation: {file_stem}")
                messagebox.showinfo(
                    "Success",
                    (
                        f"Saved '{file_stem}' under:\n{self.local_output_dir}\n\n"
                        "Subfolders created: segmentation/, detection/, classes.txt"
                    ),
                )
                return

            # --- 3. LOCAL SAVING (DB workflow, legacy path layout) ---
            file_stem = f"img_{self.current_image_id}"

            for task_type in ["seg_training", "det_training"]:
                base_dir = Path(f"data/img_sets/{task_type}")
                img_dir, lbl_dir = base_dir / "images", base_dir / "labels"
                img_dir.mkdir(parents=True, exist_ok=True)
                lbl_dir.mkdir(parents=True, exist_ok=True)

                # Save the image in both places (required for YOLO training structure)
                cv2.imwrite(str(img_dir / f"{file_stem}.jpg"), self.current_raw_frame)

                # Save appropriate label format
                lines = yolo_seg_lines if "seg" in task_type else yolo_det_lines
                with open(lbl_dir / f"{file_stem}.txt", "w") as f:
                    f.write("\n".join(lines))

            # --- 4. DATABASE SAVING ---
            db_result = self.db_client.save_annotation(
                self.current_image_id,
                segments,
                username=self.pref_username.get()
            )
            status_msg = f"Saved locally ({file_stem}) and to DB (ID: {db_result.get('annotation_id')})"
            self.status_var.set((status_msg[:120] + '...') if len(status_msg) > 123 else status_msg)

            # FIX: Removed the undefined 'img_path' variable and updated the message
            messagebox.showinfo(
                "Success",
                f"Annotated image '{file_stem}' saved to local training sets (Detection & Segmentation) and synced to database."
            )

            # 5. Fetch next image automatically
            self.fetch_db_image()

        except Exception as e:
            messagebox.showerror("Error", f"Failed to save:\n{str(e)}")
            import traceback; traceback.print_exc()
        finally:
            self.root.config(cursor="")

    def skip_current_image(self):
        """Skip the current image and fetch the next one."""
        if self.current_image_id is None:
            messagebox.showwarning(
                "Warning",
                "No image to skip. Please use 'Next Image' button first."
            )
            return

        try:
            # Confirm skip
            response = messagebox.askyesno(
                "Confirm Skip",
                f"Skip image ID {self.current_image_id}?\n\n"
                "This image will be marked as skipped in the database."
            )

            if not response:
                return

            self.root.config(cursor="watch")
            result = self.db_client.skip_image(self.current_image_id)

            self.status_var.set(f"Skipped image ID {self.current_image_id}")

            # Fetch next image automatically
            self.fetch_db_image()

        except Exception as e:
            messagebox.showerror("Error", f"Failed to skip image:\n{str(e)}")
            self.status_var.set("Skip failed")
            import traceback
            traceback.print_exc()
        finally:
            self.root.config(cursor="")

    def show_sync_dialog(self):
        """Show dialog to sync missing local images and labels from DB."""
        dialog = tk.Toplevel(self.root)
        dialog.title("Sync DB -> Local")
        dialog.geometry("640x420")
        dialog.minsize(520, 340)
        dialog.resizable(True, True)
        dialog.transient(self.root)
        dialog.grab_set()

        frame = tk.Frame(dialog, padx=20, pady=20)
        frame.pack(fill=tk.BOTH, expand=True)

        tk.Label(
            frame,
            text="Copy only missing local image/annotation files from DB records.",
            font=('Arial', 10),
        ).grid(row=0, column=0, columnspan=2, sticky=tk.W, pady=(0, 12))

        tk.Label(frame, text="Since (YYYY-MM-DD or ISO):").grid(
            row=1, column=0, sticky=tk.W, pady=6
        )
        tk.Entry(frame, textvariable=self.sync_since, width=32).grid(
            row=1, column=1, sticky=tk.W, padx=8
        )

        tk.Label(frame, text="Until (YYYY-MM-DD or ISO):").grid(
            row=2, column=0, sticky=tk.W, pady=6
        )
        tk.Entry(frame, textvariable=self.sync_until, width=32).grid(
            row=2, column=1, sticky=tk.W, padx=8
        )

        tk.Label(frame, text="Limit:").grid(row=3, column=0, sticky=tk.W, pady=6)
        tk.Entry(frame, textvariable=self.sync_limit, width=12).grid(
            row=3, column=1, sticky=tk.W, padx=8
        )

        tk.Checkbutton(
            frame,
            text="Include unannotated images",
            variable=self.sync_include_unannotated,
        ).grid(row=4, column=0, columnspan=2, sticky=tk.W, pady=(10, 4))

        tk.Checkbutton(
            frame,
            text="Dry run (preview only, no files written)",
            variable=self.sync_dry_run,
        ).grid(row=5, column=0, columnspan=2, sticky=tk.W, pady=4)

        hint = (
            "Tip: run dry-run first, then uncheck dry-run to copy files."
        )
        tk.Label(frame, text=hint, fg="#555", font=('Arial', 9)).grid(
            row=6, column=0, columnspan=2, sticky=tk.W, pady=(10, 4)
        )

        def run_sync():
            try:
                limit = int(self.sync_limit.get())
                if limit <= 0:
                    raise ValueError("Limit must be > 0")

                self.root.config(cursor="watch")
                dialog.config(cursor="watch")
                self.status_var.set("Syncing local files from database...")
                self.root.update_idletasks()

                result = self.db_client.sync_local_annotations(
                    since=self.sync_since.get().strip() or None,
                    until=self.sync_until.get().strip() or None,
                    include_unannotated=self.sync_include_unannotated.get(),
                    dry_run=self.sync_dry_run.get(),
                    limit=limit,
                )

                stats = result.get("stats", {})
                failures = result.get("failures", [])

                summary = (
                    f"Scanned: {stats.get('scanned', 0)}\n"
                    f"Images copied: {stats.get('image_copied', 0)}\n"
                    f"Images existing: {stats.get('image_existing', 0)}\n"
                    f"Images failed: {stats.get('image_failed', 0)}\n"
                    f"Labels copied: {stats.get('label_copied', 0)}\n"
                    f"Labels existing: {stats.get('label_existing', 0)}\n"
                    f"Labels failed: {stats.get('label_failed', 0)}"
                )

                if failures:
                    sample_lines = [f[:80] + "..." if len(f) > 80 else f for f in failures[:10]]
                    sample = "\n".join(sample_lines)
                    messagebox.showwarning(
                        "Sync Completed with Warnings",
                        f"{summary}\n\nFirst failures:\n{sample}",
                    )
                else:
                    messagebox.showinfo("Sync Completed", summary)

                self.status_var.set("DB->local sync completed")
            except Exception as e:
                messagebox.showerror("Sync Error", str(e))
                self.status_var.set("DB->local sync failed")
            finally:
                self.root.config(cursor="")
                dialog.config(cursor="")

        button_row = tk.Frame(frame)
        button_row.grid(row=7, column=0, columnspan=2, sticky=tk.EW, pady=(18, 0))

        tk.Button(button_row, text="Close", command=dialog.destroy, width=12).pack(
            side=tk.LEFT, padx=8
        )
        tk.Button(
            button_row,
            text="Run Sync",
            command=run_sync,
            bg="#007bff",
            fg="white",
            width=12,
            font=('Arial', 10, 'bold'),
        ).pack(side=tk.LEFT, padx=8)

    def show_zmq_inference_control_dialog(self):
        """Control runtime settings of remote zmq_inference server via HTTP API."""
        dialog = tk.Toplevel(self.root)
        dialog.title("ZMQ Inference Control")
        dialog.geometry("900x780")
        dialog.minsize(760, 700)
        dialog.resizable(True, True)
        dialog.transient(self.root)
        dialog.grab_set()

        frame = tk.Frame(dialog, padx=16, pady=14)
        frame.pack(fill=tk.BOTH, expand=True)
        frame.grid_columnconfigure(1, weight=1)

        control_url_var = tk.StringVar(value=self.pref_zmq_control_url.get().strip() or "http://localhost:8090")
        control_token_var = tk.StringVar(value=self.pref_zmq_control_token.get().strip())
        conf_var = tk.DoubleVar(value=0.25)
        infer_scale_var = tk.DoubleVar(value=1.0)
        skip_var = tk.IntVar(value=0)
        max_fps_var = tk.DoubleVar(value=0.0)
        device_var = tk.StringVar(value="0")
        shard_count_var = tk.IntVar(value=1)
        shard_index_var = tk.IntVar(value=0)
        model_choice_var = tk.StringVar(value="")
        model_seg_hint_var = tk.StringVar(value="Segmentation capable: unknown")
        cuda_hint_var = tk.StringVar(value="Server CUDA: unknown")
        task_var = tk.StringVar(value="auto")
        tracker_var = tk.StringVar(value="none")
        bayer_flip_var = tk.BooleanVar(value=False)
        publish_enabled_var = tk.BooleanVar(value=True)
        compute_target_point_var = tk.BooleanVar(value=False)
        focus_enabled_var = tk.BooleanVar(value=False)
        focus_epics_enabled_var = tk.BooleanVar(value=False)
        focus_pv_var = tk.StringVar(value="X10SA-ES-MS:cam1:FocusScore")
        focus_every_var = tk.IntVar(value=3)
        focus_scale_var = tk.DoubleVar(value=0.5)
        focus_pv_min_period_var = tk.DoubleVar(value=100.0)
        model_items_by_label = {}
        server_cuda_available = {"value": False}
        model_refresh_job = {"id": None}
        runtime_model_paths = {"engine": "", "pt": ""}

        tk.Label(frame, text="Control URL:", font=('Arial', 10, 'bold')).grid(row=0, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=control_url_var, width=52).grid(row=0, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, text="Control Token:").grid(row=1, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=control_token_var, width=30, show="*").grid(row=1, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, text="Confidence:").grid(row=2, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=conf_var, width=16).grid(row=2, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, text="Infer Scale:").grid(row=3, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=infer_scale_var, width=16).grid(row=3, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, text="Skip Frames:").grid(row=4, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=skip_var, width=16).grid(row=4, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, text="Max Predict FPS (0=off):").grid(row=5, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=max_fps_var, width=16).grid(row=5, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, text="Model (models/):").grid(row=6, column=0, sticky=tk.W, pady=4)
        model_combo = ttk.Combobox(frame, textvariable=model_choice_var, width=58, state="readonly")
        model_combo.grid(row=6, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, textvariable=model_seg_hint_var, fg="#126b2c", font=('Arial', 9, 'bold')).grid(
            row=7, column=1, sticky=tk.W, padx=8, pady=(0, 2)
        )
        tk.Label(frame, textvariable=cuda_hint_var, fg="#1c4f8a", font=('Arial', 9, 'bold')).grid(
            row=8, column=1, sticky=tk.W, padx=8, pady=(0, 2)
        )

        tk.Label(frame, text="Task:").grid(row=9, column=0, sticky=tk.W, pady=4)
        task_combo = ttk.Combobox(
            frame,
            textvariable=task_var,
            values=("auto", "detect", "segment"),
            width=14,
            state="readonly",
        )
        task_combo.grid(row=9, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, text="Tracker:").grid(row=10, column=0, sticky=tk.W, pady=4)
        tracker_combo = ttk.Combobox(
            frame,
            textvariable=tracker_var,
            values=("none", "bytetrack", "botsort"),
            width=14,
            state="readonly",
        )
        tracker_combo.grid(row=10, column=1, sticky=tk.W, padx=8)

        tk.Label(frame, text="Device (cpu or GPU idx):").grid(row=11, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=device_var, width=16).grid(row=11, column=1, sticky=tk.W, padx=8)
        tk.Label(frame, text="Shard Count (1=off, 0=auto):").grid(row=12, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=shard_count_var, width=16).grid(row=12, column=1, sticky=tk.W, padx=8)
        tk.Label(frame, text="Shard Index:").grid(row=13, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=shard_index_var, width=16).grid(row=13, column=1, sticky=tk.W, padx=8)

        tk.Checkbutton(frame, text="Bayer Flip", variable=bayer_flip_var).grid(
            row=14, column=1, sticky=tk.W, padx=8, pady=(6, 2)
        )
        tk.Checkbutton(frame, text="Publish Output", variable=publish_enabled_var).grid(
            row=15, column=1, sticky=tk.W, padx=8, pady=(2, 2)
        )
        tk.Checkbutton(frame, text="Compute Target Point", variable=compute_target_point_var).grid(
            row=16, column=1, sticky=tk.W, padx=8, pady=(2, 2)
        )
        tk.Checkbutton(frame, text="Compute Focus Score", variable=focus_enabled_var).grid(
            row=17, column=1, sticky=tk.W, padx=8, pady=(2, 2)
        )
        tk.Checkbutton(frame, text="Write Focus to EPICS", variable=focus_epics_enabled_var).grid(
            row=18, column=1, sticky=tk.W, padx=8, pady=(2, 2)
        )
        tk.Label(frame, text="Focus EPICS PV:").grid(row=19, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=focus_pv_var, width=52).grid(row=19, column=1, sticky=tk.W, padx=8)
        tk.Label(frame, text="Focus Every N Frames:").grid(row=20, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=focus_every_var, width=16).grid(row=20, column=1, sticky=tk.W, padx=8)
        tk.Label(frame, text="Focus Scale (0-1]:").grid(row=21, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=focus_scale_var, width=16).grid(row=21, column=1, sticky=tk.W, padx=8)
        tk.Label(frame, text="EPICS Min Period (ms):").grid(row=22, column=0, sticky=tk.W, pady=4)
        tk.Entry(frame, textvariable=focus_pv_min_period_var, width=16).grid(row=22, column=1, sticky=tk.W, padx=8)

        hint = (
            "Apply changes to the running zmq_inference service.\n"
            "Model/task changes trigger hot reload for all downstream ZMQ consumers."
        )
        tk.Label(frame, text=hint, fg="#555", justify=tk.LEFT, font=('Arial', 9)).grid(
            row=23, column=0, columnspan=2, sticky=tk.W, pady=(10, 6)
        )

        def _base_url():
            base = control_url_var.get().strip().rstrip("/")
            if not base:
                raise ValueError("Control URL is empty")
            return base

        def _norm_path(value):
            if not value:
                return ""
            return os.path.normpath(os.path.expanduser(str(value).strip()))

        def _canonical_model_ref(value):
            """Normalize model path refs to improve matching across server/client cwd differences."""
            raw = str(value or "").strip()
            if not raw:
                return ""
            normalized = raw.replace("\\", "/")
            if "/models/" in normalized:
                return "models/" + normalized.split("/models/", 1)[1]
            if normalized.startswith("../models/"):
                return "models/" + normalized[len("../models/"):]
            if normalized.startswith("./models/"):
                return "models/" + normalized[len("./models/"):]
            return normalized

        def _set_segmentation_hint(item):
            if not item:
                model_seg_hint_var.set("Segmentation capable: unknown")
                return
            yes_no = "yes" if bool(item.get("segmentation_capable")) else "no"
            model_seg_hint_var.set(f"Segmentation capable: {yes_no}")

        def _model_label(item):
            name = str(item.get("name") or item.get("project_path") or item.get("path") or "model")
            kind = str(item.get("kind", "")).upper()
            seg = "SEG" if bool(item.get("segmentation_capable")) else "DET"
            kind_lower = str(item.get("kind", "")).lower()
            cuda_mark = ""
            if server_cuda_available["value"] and kind_lower in ("engine", "engines"):
                cuda_mark = " | CUDA"
            if kind:
                return f"{name} [{kind} | {seg}{cuda_mark}]"
            return f"{name} [{seg}]"

        def _match_model_item(model_path):
            if not model_path:
                return None
            normalized = _norm_path(_canonical_model_ref(model_path))
            for item in model_items_by_label.values():
                candidates = {
                    _norm_path(_canonical_model_ref(item.get("path"))),
                    _norm_path(_canonical_model_ref(item.get("project_path"))),
                    _norm_path(_canonical_model_ref(item.get("name"))),
                }
                if normalized in candidates:
                    return item
            return None

        def _apply_model_item(item, update_task=True):
            if not item:
                return
            if update_task:
                task_hint = str(item.get("task_hint") or "").lower()
                if task_hint in ("detect", "segment"):
                    task_var.set(task_hint)
            _set_segmentation_hint(item)

        def refresh_models(quiet=False):
            try:
                base = _base_url()
                token = control_token_var.get().strip()
                headers = {"X-API-Key": token} if token else None

                resp = requests.get(f"{base}/models", headers=headers, timeout=3.0)
                resp.raise_for_status()
                payload = resp.json()
                models = payload.get("models", [])
                caps = payload.get("capabilities", {})
                server_cuda_available["value"] = bool(caps.get("cuda_available", False))
                cuda_hint_var.set(f"Server CUDA: {'yes' if server_cuda_available['value'] else 'no'}")
            except Exception:
                # Fallback for older servers without /models endpoint.
                local_models = []
                models_root = Path(__file__).resolve().parent / "models"
                if models_root.exists():
                    for p in sorted(models_root.rglob("*")):
                        if not p.is_file() or p.suffix.lower() not in {".pt", ".engine", ".engines", ".yaml", ".yml", ".onnx"}:
                            continue
                        rel = str(p.relative_to(models_root))
                        seg = any(tok in p.name.lower() for tok in ("seg", "segment", "mask"))
                        local_models.append(
                            {
                                "name": rel,
                                "path": str(p),
                                "project_path": str(Path("models") / rel),
                                "kind": p.suffix.lower().lstrip("."),
                                "segmentation_capable": seg,
                                "task_hint": "segment" if seg else "detect",
                            }
                        )
                models = local_models
                if not quiet and not models:
                    messagebox.showwarning(
                        "Models Unavailable",
                        "Could not fetch models from server and no local models folder entries were found.",
                    )

            # Ensure current runtime model is selectable even if it is outside scanned models/.
            current_model_path = runtime_model_paths.get("engine", "").strip() or runtime_model_paths.get("pt", "").strip()
            if current_model_path:
                current_norm = _norm_path(_canonical_model_ref(current_model_path))
                has_match = False
                for item in models:
                    candidates = {
                        _norm_path(_canonical_model_ref(item.get("path"))),
                        _norm_path(_canonical_model_ref(item.get("project_path"))),
                        _norm_path(_canonical_model_ref(item.get("name"))),
                    }
                    if current_norm in candidates:
                        has_match = True
                        break
                if not has_match:
                    normalized_current = _canonical_model_ref(current_model_path)
                    current_name = normalized_current.split("/")[-1] if "/" in normalized_current else normalized_current
                    current_kind = "engine" if normalized_current.endswith((".engine", ".engines")) else "pt"
                    seg_cap = (
                        str(task_var.get()).strip().lower() == "segment"
                        or any(tok in current_name.lower() for tok in ("seg", "segment", "mask"))
                    )
                    models.insert(
                        0,
                        {
                            "name": current_name,
                            "path": current_model_path,
                            "project_path": normalized_current,
                            "kind": current_kind,
                            "segmentation_capable": seg_cap,
                            "task_hint": "segment" if seg_cap else "detect",
                        },
                    )

            labels = []
            model_items_by_label.clear()
            for item in models:
                label = _model_label(item)
                if len(label) > 60:
                    label = label[:57] + "..."
                labels.append(label)
                model_items_by_label[label] = item

            model_combo["values"] = labels
            if not labels:
                model_choice_var.set("")
                _set_segmentation_hint(None)
                return

            current_model_path = runtime_model_paths.get("engine", "").strip() or runtime_model_paths.get("pt", "").strip()
            selected_item = _match_model_item(current_model_path)
            if selected_item:
                selected_label = _model_label(selected_item)
                model_choice_var.set(selected_label)
                _set_segmentation_hint(selected_item)
            else:
                model_choice_var.set(labels[0])
                _set_segmentation_hint(model_items_by_label.get(labels[0]))

        def _on_model_selected(_event=None):
            item = model_items_by_label.get(model_choice_var.get())
            _apply_model_item(item, update_task=True)

        model_combo.bind("<<ComboboxSelected>>", _on_model_selected)

        def load_current(quiet=False):
            try:
                base = _base_url()
                token = control_token_var.get().strip()
                headers = {"X-API-Key": token} if token else None
                self.root.config(cursor="watch")
                dialog.config(cursor="watch")
                resp = requests.get(f"{base}/config", headers=headers, timeout=3.0)
                resp.raise_for_status()
                data = resp.json()
                cfg = data.get("config", data)
                caps = data.get("capabilities", {})
                if "cuda_available" in caps:
                    server_cuda_available["value"] = bool(caps.get("cuda_available"))
                    cuda_hint_var.set(f"Server CUDA: {'yes' if server_cuda_available['value'] else 'no'}")
                conf_var.set(float(cfg.get("conf", conf_var.get())))
                infer_scale_var.set(float(cfg.get("infer_scale", infer_scale_var.get())))
                skip_var.set(int(cfg.get("skip", skip_var.get())))
                max_fps_var.set(float(cfg.get("max_fps", max_fps_var.get())))
                device_var.set(str(cfg.get("device", device_var.get())))
                shard_count_var.set(int(cfg.get("shard_count", shard_count_var.get())))
                shard_index_var.set(int(cfg.get("shard_index", shard_index_var.get())))
                task_var.set(str(cfg.get("task", task_var.get())))
                tracker_var.set(str(cfg.get("tracker", tracker_var.get())))
                runtime_model_paths["engine"] = str(cfg.get("engine", "")).strip()
                runtime_model_paths["pt"] = str(cfg.get("pt", "")).strip()
                bayer_flip_var.set(bool(cfg.get("bayer_flip", bayer_flip_var.get())))
                publish_enabled_var.set(bool(cfg.get("publish_enabled", publish_enabled_var.get())))
                compute_target_point_var.set(bool(cfg.get("compute_target_point", compute_target_point_var.get())))
                focus_enabled_var.set(bool(cfg.get("focus_enabled", focus_enabled_var.get())))
                focus_epics_enabled_var.set(bool(cfg.get("focus_epics_enabled", focus_epics_enabled_var.get())))
                focus_pv_var.set(str(cfg.get("focus_pv", focus_pv_var.get())))
                focus_every_var.set(int(cfg.get("focus_every", focus_every_var.get())))
                focus_scale_var.set(float(cfg.get("focus_scale", focus_scale_var.get())))
                focus_pv_min_period_var.set(float(cfg.get("focus_pv_min_period_ms", focus_pv_min_period_var.get())))
                refresh_models(quiet=True)
                current_item = _match_model_item(runtime_model_paths.get("engine") or runtime_model_paths.get("pt"))
                if current_item:
                    model_choice_var.set(_model_label(current_item))
                    _set_segmentation_hint(current_item)
                else:
                    _set_segmentation_hint(None)
                self.pref_zmq_control_url.set(base)
                self.pref_zmq_control_token.set(token)
                self.status_var.set("Loaded ZMQ inference runtime config")
            except Exception as e:
                if not quiet:
                    messagebox.showerror("Load Error", f"Failed loading config:\n{e}")
                self.status_var.set("Failed loading ZMQ inference config")
            finally:
                self.root.config(cursor="")
                dialog.config(cursor="")

        def _schedule_refresh_on_url_change(*_args):
            if model_refresh_job["id"] is not None:
                try:
                    dialog.after_cancel(model_refresh_job["id"])
                except Exception:
                    pass
            def _run():
                model_refresh_job["id"] = None
                try:
                    load_current(quiet=True)
                except Exception:
                    pass
            model_refresh_job["id"] = dialog.after(500, _run)

        control_url_var.trace_add("write", _schedule_refresh_on_url_change)

        def apply_changes():
            try:
                base = _base_url()
                token = control_token_var.get().strip()
                headers = {"X-API-Key": token} if token else None
                selected_item = model_items_by_label.get(model_choice_var.get())
                if not selected_item:
                    raise ValueError("Select a model from models/ before applying.")
                selected_kind = str(selected_item.get("kind") or "").lower()
                selected_project_path = str(
                    selected_item.get("project_path") or selected_item.get("name") or ""
                ).strip()
                if not selected_project_path:
                    raise ValueError("Selected model has no usable models/ path.")
                selected_engine = selected_project_path if selected_kind in ("engine", "engines") else ""
                selected_pt = selected_project_path if selected_kind not in ("engine", "engines") else ""
                device_value = str(device_var.get()).strip()
                if not device_value:
                    raise ValueError("Device is required (cpu or CUDA index).")
                shard_count_value = int(shard_count_var.get())
                shard_index_value = int(shard_index_var.get())
                payload = {
                    "conf": float(conf_var.get()),
                    "infer_scale": float(infer_scale_var.get()),
                    "skip": int(skip_var.get()),
                    "max_fps": float(max_fps_var.get()),
                    "device": device_value,
                    "shard_count": shard_count_value,
                    "shard_index": shard_index_value,
                    "task": str(task_var.get()),
                    "tracker": str(tracker_var.get()),
                    "bayer_flip": bool(bayer_flip_var.get()),
                    "publish_enabled": bool(publish_enabled_var.get()),
                    "compute_target_point": bool(compute_target_point_var.get()),
                    "focus_enabled": bool(focus_enabled_var.get()),
                    "focus_epics_enabled": bool(focus_epics_enabled_var.get()),
                    "focus_pv": str(focus_pv_var.get()).strip(),
                    "focus_every": int(focus_every_var.get()),
                    "focus_scale": float(focus_scale_var.get()),
                    "focus_pv_min_period_ms": float(focus_pv_min_period_var.get()),
                    "engine": selected_engine,
                    "pt": selected_pt,
                }
                self.root.config(cursor="watch")
                dialog.config(cursor="watch")
                resp = requests.post(f"{base}/config", json=payload, headers=headers, timeout=6.0)
                if resp.status_code >= 400:
                    raise RuntimeError(f"HTTP {resp.status_code}: {resp.text}")
                out = resp.json()
                self.pref_zmq_control_url.set(base)
                self.pref_zmq_control_token.set(token)
                messagebox.showinfo(
                    "Config Applied",
                    f"Updated runtime config.\nModel reloaded: {bool(out.get('model_reloaded', False))}",
                )
                self.status_var.set("Applied ZMQ inference runtime config")
            except Exception as e:
                messagebox.showerror("Apply Error", f"Failed applying config:\n{e}")
                self.status_var.set("Failed applying ZMQ inference config")
            finally:
                self.root.config(cursor="")
                dialog.config(cursor="")

        button_row = tk.Frame(frame)
        button_row.grid(row=24, column=0, columnspan=2, sticky=tk.EW, pady=(12, 0))
        tk.Button(button_row, text="Close", command=dialog.destroy, width=12).pack(side=tk.LEFT, padx=6)
        tk.Button(button_row, text="Refresh Models", command=refresh_models, width=14).pack(side=tk.LEFT, padx=6)
        tk.Button(button_row, text="Load Current", command=load_current, width=14).pack(side=tk.LEFT, padx=6)
        tk.Button(
            button_row,
            text="Apply",
            command=apply_changes,
            bg="#28a745",
            fg="white",
            width=12,
            font=('Arial', 10, 'bold'),
        ).pack(side=tk.LEFT, padx=6)

        # Initialize from the currently selected server.
        load_current(quiet=True)

    # --- Preferences ---

    def show_preferences(self):
        """Show a single consolidated preferences dialog."""
        dialog = tk.Toplevel(self.root)
        dialog.title("Server Settings")
        dialog.geometry("640x500")
        dialog.minsize(500, 390)
        dialog.resizable(True, True)
        dialog.transient(self.root)
        dialog.grab_set()

        main_frame = tk.Frame(dialog, padx=20, pady=20)
        main_frame.pack(fill=tk.BOTH, expand=True)
        original_stream_url = self.zmq_det_url.get().strip()

        # Define fields to display
        config_fields = [
            ("Inference API URL:", self.pref_inference_url, False),
            ("Database API URL:", self.pref_db_url, False),
            ("Image Download URL:", self.pref_download_url, False),
            ("ZMQ Control URL:", self.pref_zmq_control_url, False),
            ("ZMQ Control Token:", self.pref_zmq_control_token, True),
            ("ZMQ Unified Stream:", self.zmq_det_url, False),
            ("Shared Password:", self.pref_shared_pw, True),
            ("Annotator Name:", self.pref_username, False)
        ]
        combo_values_by_label = {
            "Image Download URL:": self.pref_download_url_options,
            "ZMQ Control URL:": self.pref_zmq_control_options,
            "ZMQ Unified Stream:": self.pref_zmq_stream_options,
        }

        for i, (label_text, var, is_pwd) in enumerate(config_fields):
            tk.Label(main_frame, text=label_text, font=('Arial', 10)).grid(row=i, column=0, sticky=tk.W, pady=8)
            if label_text in combo_values_by_label and not is_pwd:
                entry = ttk.Combobox(main_frame, textvariable=var, values=combo_values_by_label[label_text], width=38)
            else:
                entry = tk.Entry(main_frame, textvariable=var, width=40, show="*" if is_pwd else "")
            entry.grid(row=i, column=1, sticky=tk.EW, padx=10)

        def apply_and_close():
            # Update the actual clients with the new values
            self.inference_client.set_server_url(self.pref_inference_url.get())
            self.db_client.set_server_url(self.pref_db_url.get())
            self.db_client.set_shared_password(self.pref_shared_pw.get())

            if hasattr(self.db_client, 'set_download_base'):
                self.db_client.set_download_base(self.pref_download_url.get())

            new_stream_url = self.zmq_det_url.get().strip()
            if new_stream_url and new_stream_url != original_stream_url:
                try:
                    parsed = urllib.parse.urlparse(new_stream_url)
                    host = parsed.hostname
                    if host:
                        inferred_control = f"http://{host}:8090"
                        self.pref_zmq_control_url.set(inferred_control)
                except Exception:
                    pass

            self.status_var.set("Settings updated successfully.")
            dialog.destroy()

        btn_frame = tk.Frame(main_frame)
        btn_frame.grid(row=len(config_fields), column=0, columnspan=2, pady=20)

        tk.Button(btn_frame, text="Cancel", command=dialog.destroy, width=12).pack(side=tk.LEFT, padx=10)
        tk.Button(
            btn_frame,
            text="Save Changes",
            command=apply_and_close,
            bg="#28a745",
            fg="white",
            width=12,
            font=('Arial', 10, 'bold')
        ).pack(side=tk.LEFT, padx=10)


if __name__ == "__main__":
    # --- X11 CRASH PREVENTATIVE MEASURES ---
    # Force legacy rendering to avoid RenderAddGlyphs BadLength errors on virtual displays
    os.environ["GDK_RENDERING"] = "image"
    os.environ["XFT_CONFIG"] = "/dev/null"
    os.environ["XLIB_SKIP_ARGB_VISUALS"] = "1"

    try:
        root = tk.Tk()

        # Disable X Input Methods and Anti-aliasing which trigger the BadLength error
        root.tk.call('tk', 'useinputmethods', '0')

        # Force standard core bitmap fonts for ALL widgets globally
        # This bypasses the RENDER extension entirely
        root.option_add("*Font", "fixed")
        root.option_add("*TLabel*Font", "fixed")
        root.option_add("*TButton*Font", "fixed")
        root.option_add("*Menu*Font", "fixed")

        gui = InferenceGUI(root)
        root.mainloop()
    except tk.TclError as e:
        print(f"CRITICAL: Could not connect to X11 display. Check your SSH -X forwarding.")
        print(f"Error: {e}")