"""
auto_label.py — Auto-classify crops + Web Review UI
=====================================================
Runs the full pipeline in one command:

  Step 1  Run YOLOv8n on your source images → crop every detection
  Step 2  Run the classifier on every crop → propose a class label
  Step 3  Open a local web UI so you review, confirm, fix, or reject
  Step 4  On "Commit", write the final data/ folder ready for 1b_split_dataset.py

Output layout (after commit)
-----------------------------
data/
    cola_can/        ← confirmed crops
    pepsi_can/
    …
    _rejected_/      ← crops you marked as reject
    _unreviewed_/    ← anything you closed without reviewing

Usage
-----
python auto_label.py \
    --source            shelf_images/ \
    --detector_weights  detector/best.pt \
    --classifier_weights runs/classify/best.pt \
    --output_dir        data \
    --det_conf          0.25 \
    --cls_conf          0.50 \
    --port              5000
"""

import argparse
import base64
import json
import os
import shutil
import sys
import threading
import time
import webbrowser
from pathlib import Path

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from flask import Flask, jsonify, request, send_file
from torchvision import models, transforms
from tqdm import tqdm
from ultralytics import YOLO

# ─────────────────────────── args ────────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser(description="Auto-label crops and review in browser")
    p.add_argument("--source",               required=True)
    p.add_argument("--detector_weights",     default="detector/best.pt")
    p.add_argument("--classifier_weights",   default="runs/classify/best.pt")
    p.add_argument("--output_dir",           default="data")
    p.add_argument("--staging_dir",          default=".autolabel_staging",
                   help="Temp folder for staged crops before commit (hidden by default)")
    p.add_argument("--det_conf",   type=float, default=0.25)
    p.add_argument("--det_iou",    type=float, default=0.45)
    p.add_argument("--cls_conf",   type=float, default=0.75)
    p.add_argument("--img_size",   type=int,   default=640)
    p.add_argument("--padding",    type=int,   default=10)
    p.add_argument("--min_size",   type=int,   default=32)
    p.add_argument("--port",       type=int,   default=5000)
    return p.parse_args()

ARGS = parse_args()

# ─────────────────────────── device / OOM ────────────────────────────────────

def get_device():
    if torch.cuda.is_available():   return torch.device("cuda")
    if torch.backends.mps.is_available(): return torch.device("mps")
    return torch.device("cpu")

DEVICE = get_device()

def _is_oom(exc):
    if isinstance(exc, MemoryError): return True
    if torch.cuda.is_available() and isinstance(exc, torch.cuda.OutOfMemoryError): return True
    if isinstance(exc, RuntimeError):
        m = str(exc).lower()
        return "out of memory" in m or ("memory" in m and "alloc" in m)
    return False

def _free_cache():
    if DEVICE.type == "cuda": torch.cuda.empty_cache()
    elif DEVICE.type == "mps" and hasattr(torch.mps, "empty_cache"): torch.mps.empty_cache()

# ─────────────────────────── classifier ──────────────────────────────────────

class ProductClassifier:
    def __init__(self, weights_path, device, img_size=224):
        ckpt = torch.load(weights_path, map_location=device)
        self.class_names = ckpt["class_names"]
        n = len(self.class_names)
        name = ckpt.get("args", {}).get("model", "efficientnet_b0")

        m = {
            "efficientnet_b0":    models.efficientnet_b0,
            "efficientnet_b2":    models.efficientnet_b2,
            "mobilenet_v3_small": models.mobilenet_v3_small,
            "resnet50":           models.resnet50,
        }[name](weights=None)

        if "efficientnet" in name: m.classifier[1] = torch.nn.Linear(m.classifier[1].in_features, n)
        elif "mobilenet"  in name: m.classifier[3] = torch.nn.Linear(m.classifier[3].in_features, n)
        else:                      m.fc             = torch.nn.Linear(m.fc.in_features, n)

        m.load_state_dict(ckpt["model_state"])
        self.model  = m.to(device).eval()
        self.device = device
        self.tf = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize(int(img_size * 1.15)),
            transforms.CenterCrop(img_size),
            transforms.ToTensor(),
            transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),
        ])
        self._bs = self._probe(img_size)

    def _probe(self, sz):
        if self.device.type == "cpu": return 64
        dummy = torch.zeros(1,3,sz,sz,device=self.device)
        safe, probe = 1, 2
        with torch.no_grad():
            while probe <= 512:
                try:
                    self.model(dummy.expand(probe,-1,-1,-1))
                    safe, probe = probe, probe*2
                except Exception as e:
                    if _is_oom(e): break
                    raise
                finally: _free_cache()
        return max(1, int(safe*0.8))

    @torch.no_grad()
    def predict(self, crops_bgr):
        if not crops_bgr: return []
        tensors = [self.tf(cv2.cvtColor(c, cv2.COLOR_BGR2RGB)) for c in crops_bgr]
        results, i, bs = [], 0, self._bs
        while i < len(tensors):
            batch = torch.stack(tensors[i:i+bs]).to(self.device)
            try:
                probs = F.softmax(self.model(batch), dim=1)
                confs, idxs = probs.max(dim=1)
                # Return top-3 predictions per crop for the review UI
                top3_probs, top3_idxs = torch.topk(probs, min(3, probs.shape[1]), dim=1)
                for j in range(len(idxs)):
                    results.append({
                        "class":      self.class_names[idxs[j].item()],
                        "confidence": round(confs[j].item(), 4),
                        "top3": [
                            {"class": self.class_names[top3_idxs[j][k].item()],
                             "conf":  round(top3_probs[j][k].item(), 4)}
                            for k in range(top3_probs.shape[1])
                        ]
                    })
                i += bs
            except Exception as e:
                if not _is_oom(e): raise
                _free_cache()
                bs = self._bs = max(1, bs // 2)
            finally:
                del batch
        return results

# ─────────────────────────── pipeline ────────────────────────────────────────

IMAGE_EXTS = {".jpg",".jpeg",".png",".bmp",".webp",".tiff"}

def get_images(source):
    p = Path(source)
    if p.is_file() and p.suffix.lower() in IMAGE_EXTS: return [p]
    if p.is_dir():
        imgs = sorted(f for f in p.rglob("*") if f.suffix.lower() in IMAGE_EXTS)
        if not imgs: raise FileNotFoundError(f"No images in {p}")
        return imgs
    raise FileNotFoundError(f"Cannot open source: {source}")

def img_to_b64(arr):
    _, buf = cv2.imencode(".jpg", arr, [cv2.IMWRITE_JPEG_QUALITY, 88])
    return base64.b64encode(buf).decode()

def run_pipeline(detector, classifier, args):
    """
    Detect → crop → classify → stage.
    Returns list of crop metadata dicts that the review UI reads.
    """
    staging = Path(args.staging_dir)
    if staging.exists(): shutil.rmtree(staging)
    staging.mkdir(parents=True)

    image_paths = get_images(args.source)
    print(f"\n[pipeline] {len(image_paths)} source image(s) found")

    crops_meta = []          # master list for the review UI
    total_det = total_skip = 0

    for img_path in tqdm(image_paths, desc="Detecting"):
        frame = cv2.imread(str(img_path))
        if frame is None:
            print(f"  ⚠  Cannot read {img_path.name}")
            continue

        h, w = frame.shape[:2]
        results = detector(frame, conf=args.det_conf, iou=args.det_iou,
                           imgsz=args.img_size, verbose=False)[0]

        boxes, crops_bgr = [], []
        for box in results.boxes:
            x1,y1,x2,y2 = map(int, box.xyxy[0].tolist())
            x1=max(0,x1-args.padding); y1=max(0,y1-args.padding)
            x2=min(w,x2+args.padding); y2=min(h,y2+args.padding)
            crop = frame[y1:y2, x1:x2]
            if crop.size==0 or (x2-x1)<args.min_size or (y2-y1)<args.min_size:
                total_skip += 1
                continue
            boxes.append((x1,y1,x2,y2))
            crops_bgr.append(crop)

        if not crops_bgr:
            continue

        predictions = classifier.predict(crops_bgr)

        for i, (box, pred, crop) in enumerate(zip(boxes, predictions, crops_bgr)):
            crop_id   = f"{img_path.stem}_crop{i:04d}"
            cls_name  = pred["class"]
            conf      = pred["confidence"]

            # Save crop to staging/<predicted_class>/
            cls_dir = staging / cls_name
            cls_dir.mkdir(exist_ok=True)
            crop_path = cls_dir / f"{crop_id}.jpg"
            cv2.imwrite(str(crop_path), crop)

            crops_meta.append({
                "id":           crop_id,
                "file":         str(crop_path),
                "source_image": img_path.name,
                "pred_class":   cls_name,
                "confidence":   conf,
                "top3":         pred["top3"],
                "bbox":         list(box),
                # review fields (mutated by the UI)
                "label":        cls_name,   # user's chosen label
                "status":       "auto",     # "auto" | "confirmed" | "rejected" | "reclassified"
            })
            total_det += 1

    print(f"[pipeline] {total_det} crops staged  |  {total_skip} skipped")
    return crops_meta

# ─────────────────────────── Flask app ───────────────────────────────────────

flask_app = Flask(__name__)
STATE = {
    "crops":   [],        # list of crop metadata dicts
    "classes": [],        # classifier class names
    "done":    False,     # set to True when user commits
}

@flask_app.route("/")
def index():
    return HTML_UI

@flask_app.route("/api/crops")
def api_crops():
    return jsonify(STATE["crops"])

@flask_app.route("/api/classes")
def api_classes():
    return jsonify(STATE["classes"])

@flask_app.route("/api/crop_image/<crop_id>")
def crop_image(crop_id):
    crop = next((c for c in STATE["crops"] if c["id"] == crop_id), None)
    if not crop:
        return "not found", 404
    return send_file(crop["file"], mimetype="image/jpeg")

@flask_app.route("/api/update", methods=["POST"])
def api_update():
    """Receive bulk label updates from the UI."""
    updates = request.json  # list of {id, label, status}
    lookup  = {c["id"]: c for c in STATE["crops"]}
    for u in updates:
        if u["id"] in lookup:
            lookup[u["id"]]["label"]  = u["label"]
            lookup[u["id"]]["status"] = u["status"]
    return jsonify({"ok": True})

@flask_app.route("/api/commit", methods=["POST"])
def api_commit():
    """Move staged crops to final output_dir based on user decisions."""
    out_root = Path(ARGS.output_dir)

    moved = rejected = unreviewed = 0
    for crop in STATE["crops"]:
        src = Path(crop["file"])
        if not src.exists():
            continue

        status = crop["status"]
        label  = crop["label"]

        if status == "rejected":
            dst_dir = out_root / "_rejected_"
        elif status == "unreviewed" or status == "auto":
            dst_dir = out_root / "_unreviewed_"
        else:
            dst_dir = out_root / label

        dst_dir.mkdir(parents=True, exist_ok=True)
        shutil.copy2(str(src), str(dst_dir / src.name))

        if status == "rejected":         rejected    += 1
        elif status in ("auto","unreviewed"): unreviewed += 1
        else:                            moved       += 1

    summary = {
        "committed":  moved,
        "rejected":   rejected,
        "unreviewed": unreviewed,
        "output_dir": str(out_root.resolve()),
    }
    STATE["done"] = True
    return jsonify(summary)

@flask_app.route("/api/stats")
def api_stats():
    crops = STATE["crops"]
    by_class = {}
    for c in crops:
        by_class.setdefault(c["pred_class"], 0)
        by_class[c["pred_class"]] += 1
    statuses = {}
    for c in crops:
        statuses[c["status"]] = statuses.get(c["status"], 0) + 1
    return jsonify({"by_class": by_class, "by_status": statuses, "total": len(crops)})

# ─────────────────────────── embedded UI ─────────────────────────────────────

HTML_UI = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>Auto-Label Review</title>
<style>
  @import url('https://fonts.googleapis.com/css2?family=DM+Mono:wght@400;500&family=Syne:wght@700;800&display=swap');

  :root {
    --bg:      #0d0f12;
    --surface: #161a20;
    --border:  #252a33;
    --accent:  #00e5a0;
    --warn:    #f5c542;
    --danger:  #ff4c6a;
    --muted:   #4a5260;
    --text:    #dde3ee;
    --radius:  6px;
  }

  *, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }

  body {
    background: var(--bg);
    color: var(--text);
    font-family: 'DM Mono', monospace;
    font-size: 13px;
    min-height: 100vh;
    display: grid;
    grid-template-rows: 56px 1fr;
  }

  /* ── header ── */
  header {
    display: flex;
    align-items: center;
    gap: 20px;
    padding: 0 24px;
    border-bottom: 1px solid var(--border);
    background: var(--surface);
    position: sticky; top: 0; z-index: 50;
  }
  header h1 {
    font-family: 'Syne', sans-serif;
    font-size: 17px;
    font-weight: 800;
    letter-spacing: -.3px;
    color: var(--accent);
    white-space: nowrap;
  }
  .stats-bar { display: flex; gap: 16px; font-size: 11px; color: var(--muted); flex: 1; }
  .stats-bar span { display: flex; align-items: center; gap: 5px; }
  .dot { width: 7px; height: 7px; border-radius: 50%; display: inline-block; }
  .dot-auto       { background: var(--muted); }
  .dot-confirmed  { background: var(--accent); }
  .dot-reclassify { background: var(--warn); }
  .dot-rejected   { background: var(--danger); }

  .header-actions { display: flex; gap: 10px; margin-left: auto; }
  btn { display: inline-flex; align-items: center; gap: 6px;
        padding: 7px 14px; border-radius: var(--radius);
        font-family: 'DM Mono', monospace; font-size: 12px;
        cursor: pointer; border: none; font-weight: 500; transition: .15s; }
  .btn-outline { background: transparent; border: 1px solid var(--border); color: var(--text); }
  .btn-outline:hover { border-color: var(--accent); color: var(--accent); }
  .btn-primary { background: var(--accent); color: #000; }
  .btn-primary:hover { filter: brightness(1.1); }
  .btn-danger  { background: var(--danger); color: #fff; }
  .btn-danger:hover { filter: brightness(1.1); }

  /* ── layout ── */
  .layout { display: grid; grid-template-columns: 220px 1fr; height: calc(100vh - 56px); overflow: hidden; }

  /* ── sidebar ── */
  aside {
    border-right: 1px solid var(--border);
    overflow-y: auto;
    padding: 12px 8px;
    background: var(--surface);
  }
  .sidebar-title { font-size: 10px; color: var(--muted); padding: 4px 8px 8px; letter-spacing: 1px; text-transform: uppercase; }
  .cls-btn {
    width: 100%; text-align: left; background: none; border: none;
    color: var(--text); font-family: 'DM Mono', monospace; font-size: 12px;
    padding: 6px 10px; border-radius: var(--radius); cursor: pointer;
    display: flex; justify-content: space-between; align-items: center;
    transition: .12s;
  }
  .cls-btn:hover  { background: var(--border); }
  .cls-btn.active { background: rgba(0,229,160,.12); color: var(--accent); }
  .cls-badge { font-size: 10px; background: var(--border); border-radius: 99px;
               padding: 1px 7px; color: var(--muted); }
  .cls-btn.active .cls-badge { background: rgba(0,229,160,.2); color: var(--accent); }

  /* ── main ── */
  main { overflow-y: auto; padding: 20px 24px; }

  .toolbar {
    display: flex; align-items: center; gap: 10px;
    margin-bottom: 16px; flex-wrap: wrap;
  }
  .toolbar-label { color: var(--muted); font-size: 11px; margin-right: 4px; }
  .filt-btn {
    padding: 4px 10px; border-radius: 99px; font-size: 11px;
    cursor: pointer; border: 1px solid var(--border); background: transparent;
    color: var(--muted); font-family: 'DM Mono', monospace; transition: .12s;
  }
  .filt-btn:hover { border-color: var(--text); color: var(--text); }
  .filt-btn.on    { border-color: var(--accent); color: var(--accent); background: rgba(0,229,160,.08); }

  .conf-filter { display: flex; align-items: center; gap: 8px; margin-left: auto; font-size: 11px; color: var(--muted); }
  .conf-filter input { accent-color: var(--accent); width: 90px; }

  /* bulk actions */
  .bulk-bar {
    display: none; align-items: center; gap: 10px; padding: 10px 14px;
    background: rgba(0,229,160,.06); border: 1px solid rgba(0,229,160,.2);
    border-radius: var(--radius); margin-bottom: 14px; font-size: 12px;
  }
  .bulk-bar.visible { display: flex; }

  /* grid */
  .grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(148px, 1fr)); gap: 10px; }

  .card {
    position: relative; border-radius: var(--radius); overflow: hidden;
    border: 2px solid var(--border); cursor: pointer;
    transition: border-color .15s, transform .15s;
    background: var(--surface);
  }
  .card:hover { transform: translateY(-2px); border-color: var(--muted); }
  .card.selected { border-color: var(--accent) !important; }
  .card.confirmed  { border-color: var(--accent); }
  .card.reclassify { border-color: var(--warn); }
  .card.rejected   { border-color: var(--danger); opacity: .45; }

  .card img { width: 100%; aspect-ratio: 1; object-fit: cover; display: block; }

  .card-body { padding: 6px 8px 8px; }
  .card-class { font-size: 11px; color: var(--text); white-space: nowrap;
                overflow: hidden; text-overflow: ellipsis; font-weight: 500; }
  .card-conf  { font-size: 10px; color: var(--muted); margin-top: 1px; }
  .card-status { font-size: 9px; margin-top: 3px; text-transform: uppercase; letter-spacing: .5px; }
  .card-status.auto        { color: var(--muted); }
  .card-status.confirmed   { color: var(--accent); }
  .card-status.reclassify  { color: var(--warn); }
  .card-status.rejected    { color: var(--danger); }

  .card .sel-check {
    position: absolute; top: 5px; left: 5px;
    width: 18px; height: 18px; border-radius: 4px;
    border: 1.5px solid rgba(255,255,255,.3); background: rgba(0,0,0,.4);
    display: flex; align-items: center; justify-content: center;
    font-size: 11px; transition: .12s;
  }
  .card.selected .sel-check { background: var(--accent); border-color: var(--accent); color: #000; }

  /* ── modal ── */
  .overlay { position: fixed; inset: 0; background: rgba(0,0,0,.75); z-index: 200;
             display: none; align-items: center; justify-content: center; }
  .overlay.open { display: flex; }
  .modal {
    background: var(--surface); border: 1px solid var(--border);
    border-radius: 10px; width: 440px; max-width: 95vw;
    padding: 24px; position: relative;
  }
  .modal h2 { font-family:'Syne',sans-serif; font-size:15px; margin-bottom:14px; }
  .modal img { width: 100%; border-radius: var(--radius); margin-bottom: 14px; object-fit: contain; max-height: 220px; }
  .modal-meta { font-size: 11px; color: var(--muted); margin-bottom: 14px; line-height: 1.8; }
  .modal-meta strong { color: var(--text); }
  .top3 { display: flex; gap: 8px; margin-bottom: 16px; }
  .top3-item { flex: 1; padding: 8px; border-radius: var(--radius); border: 1px solid var(--border);
               font-size: 10px; cursor: pointer; text-align: center; transition: .12s; }
  .top3-item:hover { border-color: var(--warn); color: var(--warn); }
  .top3-item .t3c { font-weight: 500; font-size: 11px; white-space: nowrap;
                    overflow: hidden; text-overflow: ellipsis; }
  .top3-item .t3p { color: var(--muted); margin-top: 2px; }

  .cls-select { width: 100%; background: var(--bg); border: 1px solid var(--border);
                color: var(--text); padding: 7px 10px; border-radius: var(--radius);
                font-family:'DM Mono',monospace; font-size:12px; margin-bottom: 14px; }
  .cls-select:focus { outline: 1px solid var(--accent); }

  .modal-actions { display: flex; gap: 8px; }
  .modal-close { position: absolute; top: 14px; right: 14px; background: none; border: none;
                 color: var(--muted); cursor: pointer; font-size: 18px; line-height: 1; }

  /* ── commit modal ── */
  .commit-result { font-size: 12px; line-height: 2; }
  .commit-result .k { color: var(--muted); }
  .commit-result .v { color: var(--accent); font-weight: 500; }
  .commit-result .vw { color: var(--warn); font-weight: 500; }
  .commit-result .vr { color: var(--danger); font-weight: 500; }

  /* empty state */
  .empty { text-align: center; color: var(--muted); padding: 60px 20px; }
  .empty .big { font-size: 40px; margin-bottom: 12px; }

  /* progress overlay */
  .loading {
    position: fixed; inset: 0; background: var(--bg);
    display: flex; flex-direction: column; align-items: center; justify-content: center;
    z-index: 999; gap: 16px;
  }
  .loading h2 { font-family:'Syne',sans-serif; font-size: 18px; color: var(--accent); }
  .loading p  { color: var(--muted); font-size: 12px; }
  .spinner { width: 36px; height: 36px; border: 3px solid var(--border);
             border-top-color: var(--accent); border-radius: 50%; animation: spin .7s linear infinite; }
  @keyframes spin { to { transform: rotate(360deg); } }

  scrollbar-width: thin;
  ::-webkit-scrollbar { width: 5px; }
  ::-webkit-scrollbar-track { background: transparent; }
  ::-webkit-scrollbar-thumb { background: var(--border); border-radius: 99px; }
</style>
</head>
<body>

<!-- loading -->
<div class="loading" id="loading">
  <div class="spinner"></div>
  <h2>Auto-Label Review</h2>
  <p id="loading-msg">Loading crops…</p>
</div>

<!-- header -->
<header>
  <h1>⬡ AutoLabel</h1>
  <div class="stats-bar" id="stats-bar"></div>
  <div class="header-actions">
    <button class="btn-outline" style="display:inline-flex;align-items:center;gap:6px;padding:7px 14px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="confirmAll()">✓ Confirm all visible</button>
    <button class="btn-primary" style="display:inline-flex;align-items:center;gap:6px;padding:7px 16px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--accent);color:#000;font-weight:500;" onclick="openCommit()">Commit dataset →</button>
  </div>
</header>

<div class="layout">

  <!-- sidebar: class list -->
  <aside id="sidebar">
    <div class="sidebar-title">Classes</div>
    <button class="cls-btn active" data-cls="__all__" onclick="filterClass(this)">
      All <span class="cls-badge" id="badge-all">0</span>
    </button>
    <div id="class-list"></div>
  </aside>

  <!-- main content -->
  <main>
    <!-- toolbar -->
    <div class="toolbar">
      <span class="toolbar-label">Status:</span>
      <button class="filt-btn on" data-st="auto"        onclick="toggleStatus(this)">Auto</button>
      <button class="filt-btn on" data-st="confirmed"   onclick="toggleStatus(this)">Confirmed</button>
      <button class="filt-btn on" data-st="reclassify"  onclick="toggleStatus(this)">Reclassified</button>
      <button class="filt-btn on" data-st="rejected"    onclick="toggleStatus(this)">Rejected</button>

      <div class="conf-filter">
        Min conf: <input type="range" min="0" max="1" step=".05" value="0" oninput="confThresh=+this.value;renderGrid()">
        <span id="conf-val">0.00</span>
      </div>
    </div>

    <!-- bulk action bar -->
    <div class="bulk-bar" id="bulk-bar">
      <span id="bulk-count">0 selected</span>
      <button class="btn-outline" style="display:inline-flex;align-items:center;gap:6px;padding:5px 12px;border-radius:6px;font-family:'DM Mono',monospace;font-size:11px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="bulkAction('confirmed')">✓ Confirm</button>
      <button class="btn-outline" style="display:inline-flex;align-items:center;gap:6px;padding:5px 12px;border-radius:6px;font-family:'DM Mono',monospace;font-size:11px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="openBulkReclassify()">↷ Reclassify</button>
      <button style="display:inline-flex;align-items:center;gap:6px;padding:5px 12px;border-radius:6px;font-family:'DM Mono',monospace;font-size:11px;cursor:pointer;border:none;background:var(--danger);color:#fff;" onclick="bulkAction('rejected')">✕ Reject</button>
      <button class="btn-outline" style="margin-left:auto;display:inline-flex;align-items:center;gap:6px;padding:5px 12px;border-radius:6px;font-family:'DM Mono',monospace;font-size:11px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="clearSelection()">Clear</button>
    </div>

    <div class="grid" id="grid"></div>
  </main>
</div>

<!-- crop detail modal -->
<div class="overlay" id="modal-overlay" onclick="if(event.target===this)closeModal()">
  <div class="modal">
    <button class="modal-close" onclick="closeModal()">✕</button>
    <h2 id="modal-title">Crop detail</h2>
    <img id="modal-img" src="" alt="">
    <div class="modal-meta" id="modal-meta"></div>
    <div style="font-size:10px;color:var(--muted);margin-bottom:6px;text-transform:uppercase;letter-spacing:.8px;">Top-3 suggestions</div>
    <div class="top3" id="modal-top3"></div>
    <div style="font-size:10px;color:var(--muted);margin-bottom:6px;text-transform:uppercase;letter-spacing:.8px;">Override class</div>
    <select class="cls-select" id="modal-cls-select"></select>
    <div class="modal-actions">
      <button class="btn-primary" style="flex:1;display:inline-flex;align-items:center;justify-content:center;gap:6px;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--accent);color:#000;font-weight:500;" onclick="modalConfirm()">✓ Confirm</button>
      <button class="btn-outline" style="flex:1;display:inline-flex;align-items:center;justify-content:center;gap:6px;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="modalReclassify()">↷ Apply override</button>
      <button style="flex:1;display:inline-flex;align-items:center;justify-content:center;gap:6px;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--danger);color:#fff;" onclick="modalReject()">✕ Reject</button>
    </div>
  </div>
</div>

<!-- bulk reclassify modal -->
<div class="overlay" id="bulk-modal-overlay" onclick="if(event.target===this)closeBulkModal()">
  <div class="modal">
    <button class="modal-close" onclick="closeBulkModal()">✕</button>
    <h2>Reclassify <span id="bulk-modal-count"></span> crops</h2>
    <div style="font-size:11px;color:var(--muted);margin-bottom:12px;">All selected crops will be moved to this class.</div>
    <select class="cls-select" id="bulk-cls-select"></select>
    <div class="modal-actions">
      <button class="btn-primary" style="flex:1;display:inline-flex;align-items:center;justify-content:center;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--accent);color:#000;font-weight:500;" onclick="applyBulkReclassify()">Apply</button>
      <button class="btn-outline" style="flex:1;display:inline-flex;align-items:center;justify-content:center;padding:9px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="closeBulkModal()">Cancel</button>
    </div>
  </div>
</div>

<!-- commit modal -->
<div class="overlay" id="commit-overlay" onclick="if(event.target===this)closeCommit()">
  <div class="modal">
    <button class="modal-close" onclick="closeCommit()">✕</button>
    <h2 id="commit-title">Commit dataset</h2>
    <div id="commit-body"></div>
  </div>
</div>

<script>
let ALL_CROPS  = [];
let ALL_CLASSES= [];
let activeClass= '__all__';
let activeStatuses = new Set(['auto','confirmed','reclassify','rejected']);
let confThresh = 0;
let selected   = new Set();   // crop ids
let modalCropId= null;

// ── boot ──────────────────────────────────────────────────────────────────
async function boot() {
  document.getElementById('loading-msg').textContent = 'Fetching crops from server…';
  const [cropsRes, clsRes] = await Promise.all([fetch('/api/crops'), fetch('/api/classes')]);
  ALL_CROPS   = await cropsRes.json();
  ALL_CLASSES = await clsRes.json();

  // Populate selects
  [document.getElementById('modal-cls-select'),
   document.getElementById('bulk-cls-select')].forEach(sel => {
    sel.innerHTML = ALL_CLASSES.map(c => `<option value="${c}">${c}</option>`).join('');
  });

  buildSidebar();
  renderGrid();
  renderStats();
  document.getElementById('loading').style.display = 'none';
}

// ── sidebar ───────────────────────────────────────────────────────────────
function buildSidebar() {
  const counts = {};
  ALL_CROPS.forEach(c => counts[c.pred_class] = (counts[c.pred_class]||0)+1);
  const list = document.getElementById('class-list');
  list.innerHTML = ALL_CLASSES
    .filter(cls => counts[cls])
    .sort()
    .map(cls => `
      <button class="cls-btn" data-cls="${cls}" onclick="filterClass(this)">
        <span style="overflow:hidden;text-overflow:ellipsis;white-space:nowrap">${cls}</span>
        <span class="cls-badge" id="badge-${cls}">${counts[cls]||0}</span>
      </button>`).join('');
  document.getElementById('badge-all').textContent = ALL_CROPS.length;
}

function filterClass(btn) {
  document.querySelectorAll('.cls-btn').forEach(b=>b.classList.remove('active'));
  btn.classList.add('active');
  activeClass = btn.dataset.cls;
  clearSelection();
  renderGrid();
}

function toggleStatus(btn) {
  btn.classList.toggle('on');
  const st = btn.dataset.st;
  activeStatuses.has(st) ? activeStatuses.delete(st) : activeStatuses.add(st);
  renderGrid();
}

// ── grid ──────────────────────────────────────────────────────────────────
function visibleCrops() {
  return ALL_CROPS.filter(c => {
    if (activeClass !== '__all__' && c.pred_class !== activeClass) return false;
    if (!activeStatuses.has(c.status)) return false;
    if (c.confidence < confThresh) return false;
    return true;
  });
}

function renderGrid() {
  const crops = visibleCrops();
  const grid  = document.getElementById('grid');

  // Update conf label
  document.getElementById('conf-val').textContent = confThresh.toFixed(2);

  if (!crops.length) {
    grid.innerHTML = `<div class="empty" style="grid-column:1/-1">
      <div class="big">🔍</div>No crops match the current filters.</div>`;
    renderStats();
    return;
  }

  grid.innerHTML = crops.map(c => `
    <div class="card ${c.status} ${selected.has(c.id)?'selected':''}"
         id="card-${c.id}"
         onclick="handleCardClick(event,'${c.id}')">
      <div class="sel-check">${selected.has(c.id)?'✓':''}</div>
      <img loading="lazy" src="/api/crop_image/${c.id}" alt="">
      <div class="card-body">
        <div class="card-class" title="${c.label}">${c.label}</div>
        <div class="card-conf">${(c.confidence*100).toFixed(1)}%</div>
        <div class="card-status ${c.status}">${c.status}</div>
      </div>
    </div>`).join('');
  renderStats();
}

function handleCardClick(evt, id) {
  if (evt.shiftKey || evt.ctrlKey || evt.metaKey) {
    // multi-select
    selected.has(id) ? selected.delete(id) : selected.add(id);
    renderGrid();
    updateBulkBar();
  } else {
    openModal(id);
  }
}

// ── selection / bulk ──────────────────────────────────────────────────────
function clearSelection() { selected.clear(); renderGrid(); updateBulkBar(); }

function updateBulkBar() {
  const bar = document.getElementById('bulk-bar');
  document.getElementById('bulk-count').textContent = `${selected.size} selected`;
  bar.classList.toggle('visible', selected.size > 0);
}

function bulkAction(status) {
  selected.forEach(id => {
    const c = ALL_CROPS.find(x=>x.id===id);
    if (c) c.status = status;
  });
  syncUpdates([...selected]);
  clearSelection();
  renderGrid();
}

function openBulkReclassify() {
  document.getElementById('bulk-modal-count').textContent = selected.size;
  document.getElementById('bulk-modal-overlay').classList.add('open');
}
function closeBulkModal() { document.getElementById('bulk-modal-overlay').classList.remove('open'); }
function applyBulkReclassify() {
  const cls = document.getElementById('bulk-cls-select').value;
  selected.forEach(id => {
    const c = ALL_CROPS.find(x=>x.id===id);
    if (c) { c.label = cls; c.status = 'reclassify'; }
  });
  syncUpdates([...selected]);
  closeBulkModal();
  clearSelection();
  renderGrid();
}

function confirmAll() {
  const ids = visibleCrops().filter(c=>c.status==='auto').map(c=>c.id);
  ids.forEach(id => { const c = ALL_CROPS.find(x=>x.id===id); if(c) c.status='confirmed'; });
  syncUpdates(ids);
  renderGrid();
}

// ── modal ─────────────────────────────────────────────────────────────────
function openModal(id) {
  modalCropId = id;
  const c = ALL_CROPS.find(x=>x.id===id);
  document.getElementById('modal-title').textContent = c.id;
  document.getElementById('modal-img').src = `/api/crop_image/${id}`;
  document.getElementById('modal-meta').innerHTML = `
    <strong>Predicted:</strong> ${c.pred_class}<br>
    <strong>Label:</strong>     ${c.label}<br>
    <strong>Confidence:</strong>${(c.confidence*100).toFixed(1)}%<br>
    <strong>Source:</strong>    ${c.source_image}`;

  document.getElementById('modal-top3').innerHTML = (c.top3||[]).map(t=>`
    <div class="top3-item" onclick="pickTop3('${t.class}')">
      <div class="t3c">${t.class}</div>
      <div class="t3p">${(t.conf*100).toFixed(1)}%</div>
    </div>`).join('');

  // set select to current label
  document.getElementById('modal-cls-select').value = c.label;
  document.getElementById('modal-overlay').classList.add('open');
}
function closeModal() { document.getElementById('modal-overlay').classList.remove('open'); }

function pickTop3(cls) {
  document.getElementById('modal-cls-select').value = cls;
}
function modalConfirm() {
  const c = ALL_CROPS.find(x=>x.id===modalCropId);
  if(c) { c.status='confirmed'; syncUpdates([modalCropId]); }
  closeModal(); renderGrid();
}
function modalReclassify() {
  const cls = document.getElementById('modal-cls-select').value;
  const c   = ALL_CROPS.find(x=>x.id===modalCropId);
  if(c) { c.label=cls; c.status='reclassify'; syncUpdates([modalCropId]); }
  closeModal(); renderGrid();
}
function modalReject() {
  const c = ALL_CROPS.find(x=>x.id===modalCropId);
  if(c) { c.status='rejected'; syncUpdates([modalCropId]); }
  closeModal(); renderGrid();
}

// ── stats bar ─────────────────────────────────────────────────────────────
function renderStats() {
  const st = { auto:0, confirmed:0, reclassify:0, rejected:0 };
  ALL_CROPS.forEach(c => { if(st[c.status]!==undefined) st[c.status]++; });
  document.getElementById('stats-bar').innerHTML = `
    <span><span class="dot dot-auto"></span> Auto ${st.auto}</span>
    <span><span class="dot dot-confirmed"></span> Confirmed ${st.confirmed}</span>
    <span><span class="dot dot-reclassify"></span> Reclassified ${st.reclassify}</span>
    <span><span class="dot dot-rejected"></span> Rejected ${st.rejected}</span>
    <span style="margin-left:8px"> Total ${ALL_CROPS.length}</span>`;
}

// ── sync to server ─────────────────────────────────────────────────────────
async function syncUpdates(ids) {
  const updates = ids.map(id => {
    const c = ALL_CROPS.find(x=>x.id===id);
    return c ? { id: c.id, label: c.label, status: c.status } : null;
  }).filter(Boolean);
  await fetch('/api/update', { method:'POST',
    headers:{'Content-Type':'application/json'}, body:JSON.stringify(updates) });
}

// ── commit ─────────────────────────────────────────────────────────────────
function openCommit() {
  const auto = ALL_CROPS.filter(c=>c.status==='auto').length;
  document.getElementById('commit-title').textContent = 'Commit dataset';
  document.getElementById('commit-body').innerHTML = `
    <div style="font-size:12px;color:var(--muted);margin-bottom:16px;line-height:1.7">
      ${auto > 0
        ? `<span style="color:var(--warn)">⚠ ${auto} crops still have status "auto" (unreviewed).<br>They will be saved to <code>_unreviewed_/</code>.</span><br><br>`
        : ''}
      This will copy all crops to <code style="color:var(--accent)">${'data/'}</code>.<br>
      Confirmed + reclassified → their class folder.<br>
      Rejected → <code>_rejected_/</code><br>
      Unreviewed → <code>_unreviewed_/</code>
    </div>
    <div style="display:flex;gap:8px">
      <button class="btn-primary" style="flex:1;display:inline-flex;align-items:center;justify-content:center;padding:10px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:none;background:var(--accent);color:#000;font-weight:500;" onclick="doCommit()">Commit now</button>
      <button class="btn-outline" style="display:inline-flex;align-items:center;justify-content:center;padding:10px 16px;border-radius:6px;font-family:'DM Mono',monospace;font-size:12px;cursor:pointer;border:1px solid var(--border);background:transparent;color:var(--text);" onclick="closeCommit()">Cancel</button>
    </div>`;
  document.getElementById('commit-overlay').classList.add('open');
}
function closeCommit() { document.getElementById('commit-overlay').classList.remove('open'); }

async function doCommit() {
  document.getElementById('commit-body').innerHTML = '<div style="text-align:center;padding:20px"><div class="spinner" style="margin:0 auto"></div></div>';
  const res  = await fetch('/api/commit', { method:'POST' });
  const data = await res.json();
  document.getElementById('commit-title').textContent = '✓ Done!';
  document.getElementById('commit-body').innerHTML = `
    <div class="commit-result">
      <div><span class="k">Committed  </span><span class="v">${data.committed}</span></div>
      <div><span class="k">Rejected   </span><span class="vr">${data.rejected}</span></div>
      <div><span class="k">Unreviewed </span><span class="vw">${data.unreviewed}</span></div>
      <div style="margin-top:12px"><span class="k">Saved to   </span><code style="color:var(--accent)">${data.output_dir}</code></div>
    </div>
    <div style="margin-top:18px;font-size:11px;color:var(--muted)">
      Next step:<br>
      <code style="color:var(--text)">python 1b_split_dataset.py --data_dir data --output_dir crops_dataset</code>
    </div>`;
}

// keyboard shortcuts
document.addEventListener('keydown', e => {
  if (!modalCropId) return;
  if (e.key==='ArrowRight'||e.key==='ArrowLeft'||e.key==='d'||e.key==='a'||e.key==='q') {
    const visible = visibleCrops();
    const idx = visible.findIndex(c=>c.id===modalCropId);
    const next = (e.key==='ArrowRight' || e.key==='d') ? idx+1 : idx-1;
    if (next>=0 && next<visible.length) openModal(visible[next].id);
  }
  if (e.key==='c') modalConfirm();
  if (e.key==='r') modalReject();
  if (e.key==='Escape') closeModal();
});

boot();
</script>
</body>
</html>"""

# ─────────────────────────── entry point ─────────────────────────────────────

def main():
    args = ARGS

    print(f"[startup] Device     : {DEVICE}")
    print(f"[startup] Detector   : {args.detector_weights}")
    print(f"[startup] Classifier : {args.classifier_weights}")

    detector   = YOLO(args.detector_weights)
    classifier = ProductClassifier(args.classifier_weights, DEVICE)

    print("\n[step 1/2] Running detection + auto-classification …")
    crops = run_pipeline(detector, classifier, args)

    if not crops:
        print("⚠  No crops were generated. Check --source and --det_conf.")
        sys.exit(1)

    STATE["crops"]   = crops
    STATE["classes"] = classifier.class_names

    print(f"\n[step 2/2] Launching review UI on http://localhost:{args.port}")
    print("  Keyboard shortcuts inside a crop:")
    print("    C           → confirm")
    print("    R           → reject")
    print("    ← →         → previous / next crop")
    print("    Esc         → close modal")
    print("  Ctrl/Shift+click → multi-select for bulk actions")
    print("\n  Close the browser tab or press Ctrl+C to stop.\n")

    # Open browser after a short delay so Flask is ready
    def _open():
        time.sleep(1.2)
        webbrowser.open(f"http://localhost:{args.port}")
    threading.Thread(target=_open, daemon=True).start()

    # Run Flask (single-threaded to keep GPU access safe)
    flask_app.run(host="0.0.0.0", port=args.port, threaded=False, use_reloader=False)


if __name__ == "__main__":
    main()