"""
Simple customizable annotation GUI with auto-labeling support.

Built with Gradio - easy to modify and extend.
Run: python annotation_gui.py

To set default paths, edit config.py
"""

from __future__ import annotations

import argparse
import json
import subprocess
import threading
from pathlib import Path
from typing import Any

import gradio as gr
from PIL import Image, ImageDraw

# Try to load config, use fallbacks if not available
try:
    from config import (
        DEFAULT_IMAGES_DIR, DEFAULT_MODEL_WEIGHTS, DEFAULT_PORT,
        DEFAULT_DETECTION_THRESHOLD, DEFAULT_TRAIN_EPOCHS,
        DEFAULT_BATCH_SIZE, DEFAULT_LEARNING_RATE, DEFAULT_MODEL_SIZE
    )
except ImportError:
    DEFAULT_IMAGES_DIR = None
    DEFAULT_MODEL_WEIGHTS = None
    DEFAULT_PORT = 7860
    DEFAULT_DETECTION_THRESHOLD = 0.5
    DEFAULT_TRAIN_EPOCHS = 20
    DEFAULT_BATCH_SIZE = 4
    DEFAULT_LEARNING_RATE = 1e-4
    DEFAULT_MODEL_SIZE = "small"


class AnnotationApp:
    def __init__(self, images_dir: Path | None = None, model_weights: Path | None = None):
        self.images_dir = images_dir if images_dir else Path.cwd()
        self.current_model_path = model_weights
        self.image_paths = []
        self.current_idx = 0
        self.annotations = {}  # image_name -> list of boxes
        self.model = None
        self.training_process = None
        self.training_thread = None
        self.training_status = "Not training"
        
        # Load images if directory provided
        if images_dir and images_dir.exists():
            self._load_images(images_dir)
        
        if model_weights and model_weights.exists():
            self._load_model(model_weights)
    
    def _load_images(self, images_dir: Path):
        """Load images from directory."""
        self.images_dir = images_dir
        self.image_paths = sorted(
            list(images_dir.glob("*.jpg")) + list(images_dir.glob("*.png"))
        )
        self.current_idx = 0
        
        # Load existing annotations if present
        self.ann_file = images_dir / "annotations.json"
        if self.ann_file.exists():
            with self.ann_file.open("r") as f:
                self.annotations = json.load(f)
        else:
            self.annotations = {}
        
        return f"✓ Loaded {len(self.image_paths)} images from {images_dir}"
    
    def _load_model(self, weights_path: Path):
        """Load YOLO/YOLOX model for auto-labeling (Ultralytics format)."""
        try:
            from ultralytics import YOLO
            print(f"Loading model from {weights_path}...")
            self.model = YOLO(str(weights_path))
            self.current_model_path = weights_path
            print("✓ Model loaded")
            return f"✓ Model loaded from {weights_path.name}"
        except Exception as e:
            error_msg = f"⚠ Could not load model: {e}"
            print(error_msg)
            self.model = None
            return error_msg
    
    def load_new_model(self, weights_path: str) -> str:
        """Load a new model from the GUI."""
        path = Path(weights_path)
        if not path.exists():
            return f"❌ File not found: {weights_path}"
        
        return self._load_model(path)
    
    def load_new_images_dir(self, images_dir: str) -> tuple[Image.Image | None, str, str]:
        """Load a new images directory from the GUI."""
        path = Path(images_dir)
        if not path.exists():
            return None, "", f"❌ Directory not found: {images_dir}"
        
        if not path.is_dir():
            return None, "", f"❌ Not a directory: {images_dir}"
        
        result = self._load_images(path)
        
        # Load first image
        if self.image_paths:
            img, filename = self.get_current_image()
            boxes = self.annotations.get(filename, [])
            img_with_boxes = self.draw_boxes_on_image(img, boxes) if boxes else img
            boxes_text = self._format_boxes_text(boxes)
            info = f"{result}\nImage 1/{len(self.image_paths)}: {filename}"
            return img_with_boxes, boxes_text, info
        else:
            return None, "", f"{result}\n⚠️ No .jpg or .png images found in directory"
    
    def get_current_model_info(self) -> str:
        """Get info about currently loaded model."""
        if self.model and self.current_model_path:
            return f"📦 Loaded: {self.current_model_path}"
        elif self.model:
            return "📦 Model loaded (pretrained)"
        else:
            return "⚠️ No model loaded"
    
    def get_current_dir_info(self) -> str:
        """Get info about current images directory."""
        return f"📁 {self.images_dir} ({len(self.image_paths)} images)"
    
    def get_current_image(self) -> tuple[Image.Image, str]:
        """Get current image and filename."""
        if not self.image_paths:
            return None, ""
        path = self.image_paths[self.current_idx]
        img = Image.open(path).convert("RGB")
        return img, path.name
    
    def draw_boxes_on_image(self, img: Image.Image, boxes: list[dict]) -> Image.Image:
        """Draw bounding boxes on image."""
        img_draw = img.copy()
        draw = ImageDraw.Draw(img_draw)
        
        for box in boxes:
            x1, y1, x2, y2 = box["bbox"]
            label = box.get("label", "knot")
            conf = box.get("confidence", 1.0)
            
            # Draw box
            draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
            
            # Draw label
            text = f"{label} {conf:.2f}" if conf < 1.0 else label
            draw.text((x1, y1 - 20), text, fill="red")
        
        return img_draw
    
    def auto_label_current(self, threshold: float = 0.5) -> tuple[Image.Image, str, str]:
        """Auto-label current image with model."""
        if not self.model:
            img, filename = self.get_current_image()
            info = f"⚠ No model loaded | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
            return img, "", info
        
        img, filename = self.get_current_image()
        if not img:
            return None, "", "No images"
        
        # Run inference with Ultralytics YOLO
        results = self.model.predict(img, conf=threshold, verbose=False)
        
        # Convert to our format
        boxes = []
        if len(results) > 0:
            result = results[0]  # First image result
            if result.boxes is not None and len(result.boxes) > 0:
                for box in result.boxes:
                    xyxy = box.xyxy[0].cpu().numpy().tolist()  # [x1, y1, x2, y2]
                    conf = float(box.conf[0].cpu().numpy())
                    cls = int(box.cls[0].cpu().numpy())
                    
                    # Get class name if available
                    label = result.names.get(cls, f"class_{cls}") if hasattr(result, 'names') else f"class_{cls}"
                    
                    boxes.append({
                        "bbox": xyxy,
                        "label": label,
                        "confidence": conf,
                        "source": "auto"
                    })
        
        # Save
        self.annotations[filename] = boxes
        self._save_annotations()
        
        # Draw boxes on image
        img_with_boxes = self.draw_boxes_on_image(img, boxes)
        
        # Info with image index
        info = f"✓ Auto-labeled: {len(boxes)} boxes detected | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        boxes_text = self._format_boxes_text(boxes)
        
        return img_with_boxes, boxes_text, info
    
    def _format_boxes_text(self, boxes: list[dict]) -> str:
        """Format boxes for display."""
        if not boxes:
            return "No annotations"
        
        lines = []
        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = box["bbox"]
            conf = box.get("confidence", 1.0)
            source = box.get("source", "manual")
            lines.append(f"{i}: [{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}] conf={conf:.2f} ({source})")
        
        return "\n".join(lines)
    
    def load_image(self, direction: str = "current") -> tuple[Image.Image, str, str]:
        """Load image (current/next/prev)."""
        if direction == "next":
            self.current_idx = min(self.current_idx + 1, len(self.image_paths) - 1)
        elif direction == "prev":
            self.current_idx = max(self.current_idx - 1, 0)
        
        img, filename = self.get_current_image()
        if not img:
            return None, "", "No images"
        
        # Load existing annotations
        boxes = self.annotations.get(filename, [])
        img_with_boxes = self.draw_boxes_on_image(img, boxes) if boxes else img
        boxes_text = self._format_boxes_text(boxes)
        info = f"Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        
        return img_with_boxes, boxes_text, info
    
    def add_box_manual(self, x1: int, y1: int, x2: int, y2: int) -> tuple[Image.Image, str, str]:
        """Manually add a bounding box."""
        img, filename = self.get_current_image()
        if not img:
            return None, "", "No images"
        
        # Add box
        box = {
            "bbox": [float(x1), float(y1), float(x2), float(y2)],
            "label": "knot",
            "confidence": 1.0,
            "source": "manual"
        }
        
        if filename not in self.annotations:
            self.annotations[filename] = []
        self.annotations[filename].append(box)
        self._save_annotations()
        
        # Redraw
        boxes = self.annotations[filename]
        img_with_boxes = self.draw_boxes_on_image(img, boxes)
        boxes_text = self._format_boxes_text(boxes)
        info = f"✓ Added box: {len(boxes)} total | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        
        return img_with_boxes, boxes_text, info
    
    def delete_last_box(self) -> tuple[Image.Image, str, str]:
        """Delete the last box from current image."""
        img, filename = self.get_current_image()
        if not img:
            return None, "", "No images"
        
        if filename in self.annotations and self.annotations[filename]:
            self.annotations[filename].pop()
            self._save_annotations()
        
        # Redraw
        boxes = self.annotations.get(filename, [])
        img_with_boxes = self.draw_boxes_on_image(img, boxes) if boxes else img
        boxes_text = self._format_boxes_text(boxes)
        info = f"✓ Deleted last box: {len(boxes)} remaining | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        
        return img_with_boxes, boxes_text, info
    
    def clear_boxes(self) -> tuple[Image.Image, str, str]:
        """Clear all boxes from current image."""
        img, filename = self.get_current_image()
        if not img:
            return None, "", "No images"
        
        self.annotations[filename] = []
        self._save_annotations()
        
        boxes_text = "No annotations"
        info = f"✓ Cleared all boxes | Image {self.current_idx + 1}/{len(self.image_paths)}: {filename}"
        
        return img, boxes_text, info
    
    def _save_annotations(self):
        """Save annotations to JSON file."""
        with self.ann_file.open("w") as f:
            json.dump(self.annotations, f, indent=2)
    
    def export_to_coco(self, output_path: Path):
        """Export annotations to COCO format."""
        coco_data = {
            "images": [],
            "annotations": [],
            "categories": [{"id": 0, "name": "knot", "supercategory": "defect"}]
        }
        
        ann_id = 0
        for img_id, img_path in enumerate(self.image_paths):
            filename = img_path.name
            img = Image.open(img_path)
            width, height = img.size
            
            coco_data["images"].append({
                "id": img_id,
                "file_name": filename,
                "width": width,
                "height": height
            })
            
            # Add annotations
            boxes = self.annotations.get(filename, [])
            for box in boxes:
                x1, y1, x2, y2 = box["bbox"]
                w = x2 - x1
                h = y2 - y1
                
                coco_data["annotations"].append({
                    "id": ann_id,
                    "image_id": img_id,
                    "category_id": 0,
                    "bbox": [x1, y1, w, h],
                    "area": w * h,
                    "iscrowd": 0,
                    "score": box.get("confidence", 1.0)
                })
                ann_id += 1
        
        with output_path.open("w") as f:
            json.dump(coco_data, f, indent=2)
        
        return f"✓ Exported {len(coco_data['annotations'])} annotations to {output_path}"
    
    def prepare_training_dataset(self, output_dir: Path, train_split: float = 0.8, valid_split: float = 0.1):
        """Prepare dataset in RF-DETR format (train/valid/test splits)."""
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # Create splits
        import random
        annotated_images = [img for img in self.image_paths if img.name in self.annotations and self.annotations[img.name]]
        
        if len(annotated_images) < 10:
            return f"⚠️ Need at least 10 annotated images, have {len(annotated_images)}"
        
        random.shuffle(annotated_images)
        n = len(annotated_images)
        train_n = int(n * train_split)
        valid_n = int(n * valid_split)
        
        splits = {
            "train": annotated_images[:train_n],
            "valid": annotated_images[train_n:train_n + valid_n],
            "test": annotated_images[train_n + valid_n:]
        }
        
        # Create directories and copy images
        import shutil
        for split_name, split_images in splits.items():
            split_dir = output_dir / split_name
            split_dir.mkdir(exist_ok=True)
            
            # Prepare COCO JSON for this split
            coco_data = {
                "images": [],
                "annotations": [],
                "categories": [{"id": 0, "name": "knot", "supercategory": "defect"}]
            }
            
            ann_id = 0
            for img_id, img_path in enumerate(split_images):
                # Copy image
                dest = split_dir / img_path.name
                shutil.copy2(img_path, dest)
                
                # Add to COCO
                img = Image.open(img_path)
                width, height = img.size
                
                coco_data["images"].append({
                    "id": img_id,
                    "file_name": img_path.name,
                    "width": width,
                    "height": height
                })
                
                # Add annotations
                boxes = self.annotations.get(img_path.name, [])
                for box in boxes:
                    x1, y1, x2, y2 = box["bbox"]
                    w = x2 - x1
                    h = y2 - y1
                    
                    coco_data["annotations"].append({
                        "id": ann_id,
                        "image_id": img_id,
                        "category_id": 0,
                        "bbox": [x1, y1, w, h],
                        "area": w * h,
                        "iscrowd": 0
                    })
                    ann_id += 1
            
            # Save COCO JSON
            with (split_dir / "_annotations.coco.json").open("w") as f:
                json.dump(coco_data, f, indent=2)
        
        return f"✓ Dataset prepared: {len(splits['train'])} train, {len(splits['valid'])} valid, {len(splits['test'])} test"
    
    def start_training(self, framework: str, dataset_dir: str, output_dir: str, model_size: str, 
                      epochs: int, batch_size: int, lr: float, progress=gr.Progress()):
        """Start training in background."""
        dataset_path = Path(dataset_dir)
        output_path = Path(output_dir)
        
        if not dataset_path.exists():
            return "❌ Dataset directory not found"
        
        if self.training_process and self.training_process.poll() is None:
            return "⚠️ Training already in progress"
        
        output_path.mkdir(parents=True, exist_ok=True)
        
        # Build training command based on framework
        venv_python = Path(__file__).parent / ".venv/bin/python"
        
        if framework == "RT-DETR":
            train_script = Path(__file__).parent / "train_rtdetr.py"
            # Map sizes: nano->r18, small->r34, medium->r50, base->l
            size_map = {"nano": "rtdetr-r18", "small": "rtdetr-r34", "medium": "rtdetr-r50", "base": "rtdetr-l"}
            model_arg = size_map.get(model_size, "rtdetr-r18")
            
            cmd = [
                str(venv_python),
                str(train_script),
                "--dataset-dir", str(dataset_path),
                "--output-dir", str(output_path),
                "--model", model_arg,
                "--epochs", str(epochs),
                "--batch-size", str(batch_size),
                "--lr", str(lr)
            ]
        elif framework == "YOLOv6":
            train_script = Path(__file__).parent / "train_yolov6.py"
            # Map sizes: nano->n, small->s, medium->m, base->l
            size_map = {"nano": "yolov6n", "small": "yolov6s", "medium": "yolov6m", "base": "yolov6l"}
            model_arg = size_map.get(model_size, "yolov6n")
            
            cmd = [
                str(venv_python),
                str(train_script),
                "--dataset-dir", str(dataset_path),
                "--output-dir", str(output_path),
                "--model", model_arg,
                "--epochs", str(epochs),
                "--batch-size", str(batch_size),
                "--lr", str(lr)
            ]
        elif framework == "YOLOX":
            train_script = Path(__file__).parent / "train_yolox.py"
            # Map sizes: nano->nano, small->s, medium->m, base->l
            size_map = {"nano": "yolox-nano", "small": "yolox-s", "medium": "yolox-m", "base": "yolox-l"}
            model_arg = size_map.get(model_size, "yolox-nano")
            
            cmd = [
                str(venv_python),
                str(train_script),
                "--dataset-dir", str(dataset_path),
                "--output-dir", str(output_path),
                "--model", model_arg,
                "--epochs", str(epochs),
                "--batch-size", str(batch_size),
                "--lr", str(lr)
            ]
        else:
            return f"❌ Unknown framework: {framework}"
        
        # Start training process
        log_file = output_path / "training.log"
        self.training_status = f"🚀 Starting {framework} training..."
        
        def run_training():
            try:
                with log_file.open("w") as f:
                    self.training_process = subprocess.Popen(
                        cmd,
                        stdout=f,
                        stderr=subprocess.STDOUT,
                        text=True
                    )
                    self.training_status = f"⏳ Training in progress (PID: {self.training_process.pid})"
                    self.training_process.wait()
                    
                    if self.training_process.returncode == 0:
                        self.training_status = "✅ Training completed successfully!"
                        # Reload model with new weights
                        best_weights = output_path / "checkpoint_best_total.pth"
                        if best_weights.exists():
                            self._load_model(best_weights)
                    else:
                        self.training_status = f"❌ Training failed (exit code {self.training_process.returncode})"
            except Exception as e:
                self.training_status = f"❌ Error: {e}"
        
        self.training_thread = threading.Thread(target=run_training, daemon=True)
        self.training_thread.start()
        
        return f"✓ Training started! Check {log_file} for progress"
    
    def get_training_status(self):
        """Get current training status."""
        return self.training_status
    
    def stop_training(self):
        """Stop the training process."""
        if self.training_process and self.training_process.poll() is None:
            self.training_process.terminate()
            self.training_status = "⏹️ Training stopped by user"
            return "✓ Training process terminated"
        return "⚠️ No training in progress"


def create_ui(app: AnnotationApp) -> gr.Blocks:
    """Create Gradio UI."""
    
    with gr.Blocks(title="Knot Annotation Tool") as demo:
        gr.Markdown("""
        # 🪵 Wood Knot Annotation Tool
        **Label → Train → Auto-Label → Repeat**
        
        - Manually annotate images or use **Auto-Label** with your trained model
        - Export and prepare dataset for training
        - Train **RT-DETR, YOLOv6, or YOLOX** (all free for commercial use!)
        - Optimized for OAK-D camera deployment
        - Use trained model to auto-label more images
        """)
        
        # Settings section at the top
        with gr.Accordion("⚙️ Settings", open=False):
            with gr.Row():
                with gr.Column():
                    images_dir_input = gr.Textbox(
                        label="Images Directory",
                        value=str(app.images_dir),
                        placeholder="/path/to/images"
                    )
                    load_images_btn = gr.Button("📁 Load Images Directory")
                    dir_info = gr.Textbox(label="Current Directory", value=app.get_current_dir_info(), interactive=False)
                
                with gr.Column():
                    model_weights_input = gr.Textbox(
                        label="Model Weights Path",
                        value=str(app.current_model_path) if app.current_model_path else "",
                        placeholder="runs/training/checkpoint_best_total.pth"
                    )
                    load_model_btn = gr.Button("🤖 Load Model Weights")
                    model_info = gr.Textbox(label="Current Model", value=app.get_current_model_info(), interactive=False)
        
        with gr.Row():
            with gr.Column(scale=3):
                image_display = gr.Image(label="Current Image", type="pil")
                
                with gr.Row():
                    prev_btn = gr.Button("⬅️ Previous")
                    next_btn = gr.Button("Next ➡️")
                    auto_label_btn = gr.Button("🤖 Auto-Label", variant="primary")
                
                with gr.Row():
                    threshold_slider = gr.Slider(0.1, 0.9, DEFAULT_DETECTION_THRESHOLD, label="Detection Threshold")
            
            with gr.Column(scale=1):
                info_text = gr.Textbox(label="Status", lines=2)
                boxes_text = gr.Textbox(label="Annotations", lines=10)
                
                gr.Markdown("### Manual Annotation")
                with gr.Row():
                    x1_input = gr.Number(label="x1", value=100)
                    y1_input = gr.Number(label="y1", value=100)
                with gr.Row():
                    x2_input = gr.Number(label="x2", value=200)
                    y2_input = gr.Number(label="y2", value=200)
                
                add_box_btn = gr.Button("➕ Add Box")
                delete_btn = gr.Button("🗑️ Delete Last")
                clear_btn = gr.Button("❌ Clear All")
                
                gr.Markdown("### Export & Training")
                export_path = gr.Textbox(
                    label="Export Path",
                    value="annotations_coco.json"
                )
                export_btn = gr.Button("💾 Export COCO")
                export_result = gr.Textbox(label="Export Result", lines=1)
        
        # Training tab
        with gr.Tab("🎯 Training"):
            gr.Markdown("""
            ### Train Object Detection Model
            
            **Choose your framework:**
            - **RT-DETR** (Apache 2.0): Modern transformer, great accuracy
            - **YOLOv6** (MIT): Fast, proven on OAK cameras
            - **YOLOX** (MIT): Similar to YOLOv6, slight differences
            
            **All MIT/Apache 2.0 licensed - free for commercial use!** ✅
            
            **Steps:**
            1. Annotate at least 50-100 images in the Annotation tab
            2. Click "Prepare Dataset" to create train/valid/test splits
            3. Select your framework and configure training parameters
            4. Click "Start Training" (runs in background)
            5. After training, export for OAK-D deployment
            """)
            
            with gr.Row():
                with gr.Column():
                    dataset_prep_dir = gr.Textbox(
                        label="Dataset Output Directory",
                        value="dataset_prepared"
                    )
                    train_split = gr.Slider(0.5, 0.9, 0.8, label="Train Split Ratio")
                    valid_split = gr.Slider(0.05, 0.3, 0.1, label="Valid Split Ratio")
                    prep_btn = gr.Button("📦 Prepare Dataset", variant="secondary")
                    prep_result = gr.Textbox(label="Preparation Result", lines=2)
                
                with gr.Column():
                    gr.Markdown("### Training Configuration")
                    model_framework = gr.Dropdown(
                        choices=["RT-DETR", "YOLOv6", "YOLOX"],
                        value="RT-DETR",
                        label="Model Framework",
                        info="All MIT/Apache 2.0 licensed - free for commercial use. Optimized for OAK cameras."
                    )
                    train_dataset_dir = gr.Textbox(
                        label="Dataset Directory",
                        value="dataset_prepared"
                    )
                    train_output_dir = gr.Textbox(
                        label="Output Directory",
                        value="runs/gui_training"
                    )
                    model_size = gr.Dropdown(
                        choices=["nano", "small", "medium", "base"],
                        value=DEFAULT_MODEL_SIZE,
                        label="Model Size"
                    )
                    epochs = gr.Slider(5, 100, DEFAULT_TRAIN_EPOCHS, step=5, label="Epochs")
                    batch_size = gr.Slider(1, 16, DEFAULT_BATCH_SIZE, step=1, label="Batch Size")
                    learning_rate = gr.Number(value=DEFAULT_LEARNING_RATE, label="Learning Rate")
            
            with gr.Row():
                start_train_btn = gr.Button("🚀 Start Training", variant="primary")
                stop_train_btn = gr.Button("⏹️ Stop Training", variant="stop")
                refresh_status_btn = gr.Button("🔄 Refresh Status")
            
            training_status = gr.Textbox(
                label="Training Status",
                value="Not training",
                lines=3
            )
            
            gr.Markdown("""
            **Note**: Training runs in the background. You can continue annotating while training.
            Check the training log file for detailed progress.
            """)
        
        # Event handlers
        def on_load():
            return app.load_image("current")
        
        # Settings handlers
        load_images_btn.click(
            app.load_new_images_dir,
            inputs=[images_dir_input],
            outputs=[image_display, boxes_text, info_text]
        ).then(
            lambda: (app.get_current_dir_info(), app.get_current_model_info()),
            outputs=[dir_info, model_info]
        )
        
        load_model_btn.click(
            app.load_new_model,
            inputs=[model_weights_input],
            outputs=[model_info]
        )
        
        prev_btn.click(
            lambda: app.load_image("prev"),
            outputs=[image_display, boxes_text, info_text]
        )
        
        next_btn.click(
            lambda: app.load_image("next"),
            outputs=[image_display, boxes_text, info_text]
        )
        
        auto_label_btn.click(
            lambda t: app.auto_label_current(t),
            inputs=[threshold_slider],
            outputs=[image_display, boxes_text, info_text]
        )
        
        add_box_btn.click(
            app.add_box_manual,
            inputs=[x1_input, y1_input, x2_input, y2_input],
            outputs=[image_display, boxes_text, info_text]
        )
        
        delete_btn.click(
            app.delete_last_box,
            outputs=[image_display, boxes_text, info_text]
        )
        
        clear_btn.click(
            app.clear_boxes,
            outputs=[image_display, boxes_text, info_text]
        )
        
        export_btn.click(
            lambda path: app.export_to_coco(Path(path)),
            inputs=[export_path],
            outputs=[export_result]
        )
        
        # Training handlers
        prep_btn.click(
            lambda out, train, valid: app.prepare_training_dataset(Path(out), train, valid),
            inputs=[dataset_prep_dir, train_split, valid_split],
            outputs=[prep_result]
        )
        
        start_train_btn.click(
            app.start_training,
            inputs=[model_framework, train_dataset_dir, train_output_dir, model_size, epochs, batch_size, learning_rate],
            outputs=[training_status]
        )
        
        stop_train_btn.click(
            app.stop_training,
            outputs=[training_status]
        )
        
        refresh_status_btn.click(
            app.get_training_status,
            outputs=[training_status]
        )
        
        # Load first image on start
        demo.load(on_load, outputs=[image_display, boxes_text, info_text])
    
    return demo


def main():
    parser = argparse.ArgumentParser(description="Simple annotation GUI with auto-labeling")
    parser.add_argument(
        "--images-dir",
        type=Path,
        default=Path(DEFAULT_IMAGES_DIR) if DEFAULT_IMAGES_DIR else None,
        help="Default directory with images (can be changed in GUI)"
    )
    parser.add_argument(
        "--model-weights",
        type=Path,
        default=Path(DEFAULT_MODEL_WEIGHTS) if DEFAULT_MODEL_WEIGHTS else None,
        help="Default trained model for auto-labeling (can be changed in GUI)"
    )
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="Port for web interface")
    args = parser.parse_args()
    
    # Validate paths if provided
    if args.images_dir and not args.images_dir.exists():
        print(f"⚠️ Warning: Images directory not found: {args.images_dir}")
        print("You can load a different directory from the GUI Settings")
        args.images_dir = None
    
    if args.model_weights and not args.model_weights.exists():
        print(f"⚠️ Warning: Model weights not found: {args.model_weights}")
        print("You can load different weights from the GUI Settings")
        args.model_weights = None
    
    # Create app
    app = AnnotationApp(args.images_dir, args.model_weights)
    
    # Create and launch UI
    demo = create_ui(app)
    
    print(f"\n{'='*60}")
    print(f"🚀 Starting annotation tool...")
    if args.images_dir:
        print(f"📁 Default images: {args.images_dir} ({len(app.image_paths)} images)")
    else:
        print(f"📁 No default images - load directory from Settings")
    if app.model:
        print(f"🤖 Model: Loaded from {args.model_weights}")
    else:
        print(f"⚠️  No model loaded - load from Settings or train one")
    print(f"💡 You can change images directory and model weights from the Settings panel")
    print(f"{'='*60}\n")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=args.port,
        share=False
    )


if __name__ == "__main__":
    main()