"""
class_labeller.py — Class Labelling Webapp
============================================
A high-performance web app for assigning class labels to crop images.
Handles 14,000+ images via:
  • Parallel thumbnail generation (multiprocessing)
  • On-disk thumbnail + metadata cache (skips re-processing)
  • Virtual scrolling (only visible DOM nodes rendered)
  • Color-grouped layout so visually similar products sit together

Workflow
--------
1. Run this script pointing at your data/unknown/ (or any flat image folder)
2. Browser opens — images are sorted by dominant color into groups
3. Click / Shift-click / Ctrl-click to select images
4. Type a class name (or pick existing) → Assign
5. All assigned images are MOVED to data/<class_name>/
6. Remaining unassigned images stay in the source folder

Usage
-----
python class_labeller.py \
    --source     data/unknown \
    --data_dir   data \
    --port       5001 \
    --workers    8 \
    --thumb_size 112
"""

import argparse
import hashlib
import json
import math
import multiprocessing
import os
import shutil
import time
from functools import partial
from pathlib import Path

import cv2
import numpy as np
from flask import Flask, Response, jsonify, request, send_file


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

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--source",     default="data/unknown",
                   help="Folder of unlabelled images")
    p.add_argument("--data_dir",   default="data",
                   help="Root data folder — assigned images go to data/<class>/")
    p.add_argument("--port",       type=int, default=5001)
    p.add_argument("--workers",    type=int,
                   default=max(1, multiprocessing.cpu_count() - 1),
                   help="Parallel workers for thumbnail generation")
    p.add_argument("--thumb_size", type=int, default=112,
                   help="Thumbnail size in pixels (square)")
    p.add_argument("--cache_dir",  default=".label_cache",
                   help="Where thumbnails and metadata are cached")
    p.add_argument("--n_colors",   type=int, default=16,
                   help="Number of color buckets for grouping")
    return p.parse_args()

ARGS = parse_args()

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

# ─────────────────────────── color helpers ────────────────────────────────────

# Named color buckets (hue ranges in OpenCV HSV, hue 0-179)
COLOR_BUCKETS = [
    (  0,  10, "Red"   ),
    ( 10,  20, "Orange"),
    ( 20,  35, "Yellow"),
    ( 35,  80, "Green" ),
    ( 80,  95, "Cyan"  ),
    ( 95, 130, "Blue"  ),
    (130, 150, "Purple"),
    (150, 170, "Pink"  ),
    (170, 180, "Red"   ),   # red wraps around
]

def dominant_color_hsv(img_bgr: np.ndarray) -> tuple[int, int, int]:
    """
    Return the dominant color of an image as (H, S, V) in OpenCV ranges.
    Strategy: resize to 32×32, convert to HSV, take the median of each channel.
    Fast enough to run on thousands of images without a GPU.
    """
    small = cv2.resize(img_bgr, (32, 32), interpolation=cv2.INTER_AREA)
    hsv   = cv2.cvtColor(small, cv2.COLOR_BGR2HSV)
    h = int(np.median(hsv[:, :, 0]))
    s = int(np.median(hsv[:, :, 1]))
    v = int(np.median(hsv[:, :, 2]))
    return h, s, v

def color_bucket(h: int, s: int, v: int) -> str:
    """Map (H, S, V) to a human-readable color name."""
    if v < 50:                      return "Black"
    if s < 40:
        return "White" if v > 180 else "Gray"
    for lo, hi, name in COLOR_BUCKETS:
        if lo <= h < hi:
            return name
    return "Red"

BUCKET_ORDER = [
    "Red","Orange","Yellow","Green","Cyan",
    "Blue","Purple","Pink","White","Gray","Black"
]

def sort_key(meta: dict) -> tuple:
    """Sort key: bucket index → hue → saturation descending → value."""
    b = meta["bucket"]
    idx = BUCKET_ORDER.index(b) if b in BUCKET_ORDER else 99
    return (idx, meta["h"], -meta["s"], -meta["v"])

# ─────────────────────────── thumbnail generation ────────────────────────────

def process_one(img_path_str: str, thumb_dir: str, size: int) -> dict | None:
    """
    Worker function (runs in subprocess):
    Read one image, make a thumbnail, extract dominant color.
    Returns metadata dict or None on failure.
    """
    try:
        img_path = Path(img_path_str)
        img      = cv2.imread(img_path_str)
        if img is None:
            return None

        # Thumbnail
        h, w   = img.shape[:2]
        scale  = size / max(h, w)
        th, tw = int(h * scale), int(w * scale)
        thumb  = cv2.resize(img, (tw, th), interpolation=cv2.INTER_AREA)

        # Center-pad to square
        canvas = np.zeros((size, size, 3), dtype=np.uint8)
        y0 = (size - th) // 2
        x0 = (size - tw) // 2
        canvas[y0:y0+th, x0:x0+tw] = thumb

        # Save thumbnail — use a hash of the original path as filename
        # so re-runs are skipped cleanly
        key      = hashlib.md5(img_path_str.encode()).hexdigest()
        out_path = Path(thumb_dir) / f"{key}.jpg"
        if not out_path.exists():
            cv2.imwrite(str(out_path), canvas,
                        [cv2.IMWRITE_JPEG_QUALITY, 82])

        dh, ds, dv = dominant_color_hsv(img)
        bucket = color_bucket(dh, ds, dv)

        return {
            "id":       key,
            "path":     img_path_str,
            "name":     img_path.name,
            "bucket":   bucket,
            "h": dh, "s": ds, "v": dv,
            "w": w,   "height": h,
        }
    except Exception:
        return None


def build_metadata(source_dir: Path, cache_dir: Path,
                   thumb_size: int, n_workers: int) -> list[dict]:
    """
    Generate thumbnails for every image in source_dir (parallel),
    cache results to disk, return sorted metadata list.
    Re-running only processes new images.
    """
    thumb_dir  = cache_dir / "thumbs"
    thumb_dir.mkdir(parents=True, exist_ok=True)
    meta_file  = cache_dir / "metadata.json"

    # Load existing cache
    cached: dict[str, dict] = {}
    if meta_file.exists():
        try:
            for m in json.loads(meta_file.read_text()):
                cached[m["path"]] = m
        except Exception:
            pass

    # Find all images in source_dir
    all_paths = sorted(
        str(p) for p in source_dir.rglob("*")
        if p.suffix.lower() in IMAGE_EXTS
    )

    # Only process images not in cache
    new_paths = [p for p in all_paths if p not in cached]

    if new_paths:
        print(f"[cache] Processing {len(new_paths)} new images with {n_workers} workers…")
        worker = partial(process_one, thumb_dir=str(thumb_dir), size=thumb_size)

        with multiprocessing.Pool(processes=n_workers) as pool:
            chunk   = max(1, len(new_paths) // (n_workers * 4))
            results = []
            done    = 0
            for batch in [new_paths[i:i+chunk] for i in range(0, len(new_paths), chunk)]:
                batch_results = pool.map(worker, batch)
                for r in batch_results:
                    if r:
                        cached[r["path"]] = r
                done += len(batch)
                pct   = done / len(new_paths) * 100
                print(f"  {done}/{len(new_paths)}  ({pct:.0f}%)", end="\r", flush=True)
        print()
    else:
        print(f"[cache] All {len(all_paths)} images already cached.")

    # Keep only images that still exist in source_dir
    meta = [cached[p] for p in all_paths if p in cached]

    # Sort by color
    meta.sort(key=sort_key)

    # Persist cache
    meta_file.write_text(json.dumps(meta))
    print(f"[cache] {len(meta)} images ready.")
    return meta

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

flask_app  = Flask(__name__)
STATE: dict = {}   # populated in main()

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

@flask_app.route("/api/meta")
def api_meta():
    """
    Return lightweight metadata (no pixel data).
    The client paginates via ?offset=&limit=
    """
    offset = int(request.args.get("offset", 0))
    limit  = int(request.args.get("limit",  500))
    bucket = request.args.get("bucket", "")

    items = STATE["meta"]
    if bucket and bucket != "__all__":
        items = [m for m in items if m["bucket"] == bucket]

    total  = len(items)
    page   = items[offset : offset + limit]
    # strip heavy fields before sending
    slim   = [{"id":m["id"],"name":m["name"],"bucket":m["bucket"],
               "h":m["h"],"s":m["s"],"v":m["v"]} for m in page]
    return jsonify({"total": total, "offset": offset, "items": slim})

@flask_app.route("/api/buckets")
def api_buckets():
    counts: dict[str, int] = {}
    for m in STATE["meta"]:
        counts[m["bucket"]] = counts.get(m["bucket"], 0) + 1
    ordered = [{"name": b, "count": counts[b]}
               for b in BUCKET_ORDER if b in counts]
    return jsonify(ordered)

@flask_app.route("/api/classes")
def api_classes():
    data_dir = Path(ARGS.data_dir)
    classes  = sorted(
        d.name for d in data_dir.iterdir()
        if d.is_dir() and not d.name.startswith("_") and d.name != "unknown"
    ) if data_dir.exists() else []
    return jsonify(classes)

@flask_app.route("/thumb/<img_id>")
def thumb(img_id: str):
    path = Path(ARGS.cache_dir) / "thumbs" / f"{img_id}.jpg"
    if not path.exists():
        return "", 404
    return send_file(str(path), mimetype="image/jpeg",
                     max_age=86400)   # cache 24 h in browser

@flask_app.route("/api/assign", methods=["POST"])
def api_assign():
    """
    Move selected images into data/<class_name>/.
    Body: { "ids": [...], "class_name": "cola_can" }
    """
    body       = request.json
    ids        = set(body.get("ids", []))
    class_name = body["class_name"].strip().replace(" ", "_")

    if not class_name:
        return jsonify({"error": "class_name is required"}), 400

    dst_dir = Path(ARGS.data_dir) / class_name
    dst_dir.mkdir(parents=True, exist_ok=True)

    moved   = 0
    errors  = []
    id_map  = {m["id"]: m for m in STATE["meta"]}

    for img_id in ids:
        meta = id_map.get(img_id)
        if not meta:
            continue
        src = Path(meta["path"])
        if not src.exists():
            continue
        dst = dst_dir / src.name
        # If name collision, append a short hash suffix
        if dst.exists():
            stem  = src.stem
            ext   = src.suffix
            dst   = dst_dir / f"{stem}_{img_id[:6]}{ext}"
        try:
            shutil.move(str(src), str(dst))
            moved += 1
        except Exception as e:
            errors.append(str(e))

    # Remove moved images from in-memory metadata
    STATE["meta"] = [m for m in STATE["meta"] if m["id"] not in ids]

    return jsonify({"moved": moved, "errors": errors,
                    "remaining": len(STATE["meta"])})

@flask_app.route("/api/stats")
def api_stats():
    return jsonify({
        "total":    len(STATE["meta"]),
        "buckets":  len(set(m["bucket"] for m in STATE["meta"])),
    })

# ─────────────────────────── HTML / JS UI ────────────────────────────────────

HTML = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>Class Labeller</title>
<style>
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500;600&family=Epilogue:wght@800;900&display=swap');

:root{
  --bg:#0a0b0e; --s1:#111318; --s2:#1a1d24; --bd:#24282f;
  --text:#cdd5e0; --muted:#4e5566; --accent:#00ffb2;
  --sel:#1a3d2e; --sel-bd:#00ffb2;
  --red:#ff4558; --amber:#ffb347;
  --radius:6px; --thumb:112px;
}
*{box-sizing:border-box;margin:0;padding:0}
html,body{height:100%;overflow:hidden}
body{background:var(--bg);color:var(--text);font-family:'IBM Plex Mono',monospace;
     font-size:13px;display:grid;grid-template-rows:52px 1fr;grid-template-columns:240px 1fr}

/* ── header ── */
header{grid-column:1/-1;display:flex;align-items:center;gap:16px;
       padding:0 20px;border-bottom:1px solid var(--bd);background:var(--s1);z-index:10}
header h1{font-family:'Epilogue',sans-serif;font-size:16px;font-weight:900;
          color:var(--accent);letter-spacing:-.5px;white-space:nowrap}
.hstats{font-size:11px;color:var(--muted);flex:1}
.hstats b{color:var(--text)}

/* assign bar */
.assign-bar{display:flex;gap:8px;align-items:center;margin-left:auto}
#sel-count{font-size:11px;color:var(--muted);white-space:nowrap;min-width:80px;text-align:right}
.cls-wrap{position:relative}
#cls-input{background:var(--s2);border:1px solid var(--bd);color:var(--text);
           padding:6px 12px;border-radius:var(--radius);font-family:inherit;
           font-size:12px;width:200px;outline:none}
#cls-input:focus{border-color:var(--accent)}
#cls-drop{position:absolute;top:calc(100% + 4px);left:0;right:0;
          background:var(--s2);border:1px solid var(--bd);border-radius:var(--radius);
          max-height:200px;overflow-y:auto;display:none;z-index:100}
#cls-drop.open{display:block}
.cls-opt{padding:7px 12px;cursor:pointer;font-size:12px}
.cls-opt:hover{background:var(--sel);color:var(--accent)}
.btn{display:inline-flex;align-items:center;gap:6px;padding:6px 16px;
     border-radius:var(--radius);font-family:inherit;font-size:12px;
     cursor:pointer;border:none;font-weight:600;transition:.12s}
.btn-primary{background:var(--accent);color:#000}
.btn-primary:hover{filter:brightness(1.1)}
.btn-primary:disabled{background:var(--bd);color:var(--muted);cursor:not-allowed;filter:none}
.btn-ghost{background:transparent;border:1px solid var(--bd);color:var(--muted)}
.btn-ghost:hover{border-color:var(--text);color:var(--text)}

/* ── sidebar ── */
aside{background:var(--s1);border-right:1px solid var(--bd);
      overflow-y:auto;padding:12px 8px}
.side-title{font-size:10px;letter-spacing:1px;text-transform:uppercase;
            color:var(--muted);padding:4px 8px 8px}
.bkt-btn{width:100%;background:none;border:none;color:var(--text);
         font-family:inherit;font-size:12px;padding:7px 10px;
         border-radius:var(--radius);cursor:pointer;
         display:flex;justify-content:space-between;align-items:center;transition:.1s}
.bkt-btn:hover{background:var(--s2)}
.bkt-btn.active{background:rgba(0,255,178,.08);color:var(--accent)}
.bkt-swatch{width:10px;height:10px;border-radius:2px;flex-shrink:0;margin-right:6px}
.bkt-left{display:flex;align-items:center;gap:0}
.bkt-count{font-size:10px;color:var(--muted);background:var(--bd);
           border-radius:99px;padding:1px 7px}
.bkt-btn.active .bkt-count{background:rgba(0,255,178,.15);color:var(--accent)}

/* ── main ── */
main{overflow:hidden;display:flex;flex-direction:column;position:relative}

/* toolbar */
.toolbar{display:flex;align-items:center;gap:10px;padding:10px 16px;
         border-bottom:1px solid var(--bd);flex-shrink:0;background:var(--s1)}
.sel-bar{display:none;align-items:center;gap:8px;
         background:rgba(0,255,178,.06);border:1px solid rgba(0,255,178,.2);
         border-radius:var(--radius);padding:6px 12px;font-size:12px}
.sel-bar.on{display:flex}
.filter-input{background:var(--s2);border:1px solid var(--bd);color:var(--text);
              padding:5px 10px;border-radius:var(--radius);
              font-family:inherit;font-size:12px;outline:none;width:180px}
.filter-input:focus{border-color:var(--accent)}
.toolbar-right{display:flex;align-items:center;gap:10px;margin-left:auto;font-size:11px;color:var(--muted)}

/* scroller */
.scroll-host{flex:1;overflow-y:auto;overflow-x:hidden;position:relative}
#spacer{width:1px} /* height set by JS */
#grid-rows{position:absolute;top:0;left:0;right:0;padding:12px 14px}

/* row */
.grid-row{display:flex;flex-wrap:wrap;gap:6px;margin-bottom:6px;contain:layout style paint}

/* card */
.card{width:var(--thumb);flex-shrink:0;border-radius:var(--radius);
      overflow:hidden;cursor:pointer;border:2px solid transparent;
      background:var(--s2);transition:border-color .1s,transform .1s;
      position:relative;user-select:none;contain:layout style paint}
.card:hover{border-color:var(--muted);transform:translateY(-2px)}
.card.sel{border-color:var(--sel-bd);background:var(--sel)}
.card img{width:100%;height:var(--thumb);object-fit:cover;display:block;
          background:var(--bd)}
.card .cname{font-size:10px;padding:3px 6px;white-space:nowrap;
             overflow:hidden;text-overflow:ellipsis;color:var(--muted)}
.card.sel .cname{color:var(--accent)}
.card .chk{position:absolute;top:4px;left:4px;width:16px;height:16px;
           border-radius:3px;border:1.5px solid rgba(255,255,255,.25);
           background:rgba(0,0,0,.35);display:flex;align-items:center;
           justify-content:center;font-size:9px;transition:.1s}
.card.sel .chk{background:var(--accent);border-color:var(--accent);color:#000}

/* bucket label row */
.bucket-label{width:100%;font-size:10px;letter-spacing:.8px;
              text-transform:uppercase;color:var(--muted);
              padding:10px 2px 4px;display:flex;align-items:center;gap:8px}
.bucket-label::after{content:'';flex:1;height:1px;background:var(--bd)}

/* loading overlay */
.loading{position:absolute;inset:0;background:var(--bg);
         display:flex;flex-direction:column;align-items:center;
         justify-content:center;gap:14px;z-index:50}
.loading h2{font-family:'Epilogue',sans-serif;font-size:18px;color:var(--accent)}
.loading p{font-size:12px;color:var(--muted)}
.bar-wrap{width:260px;height:4px;background:var(--bd);border-radius:99px;overflow:hidden}
.bar-fill{height:100%;background:var(--accent);border-radius:99px;
          transition:width .3s;width:0}
.spinner{width:28px;height:28px;border:2px solid var(--bd);
         border-top-color:var(--accent);border-radius:50%;
         animation:spin .6s linear infinite}
@keyframes spin{to{transform:rotate(360deg)}}

/* toast */
.toast{position:fixed;bottom:24px;right:24px;background:var(--s2);
       border:1px solid var(--bd);border-radius:var(--radius);
       padding:10px 18px;font-size:12px;opacity:0;
       transform:translateY(8px);transition:.25s;pointer-events:none;z-index:200}
.toast.show{opacity:1;transform:none}
.toast.ok{border-color:var(--accent);color:var(--accent)}
.toast.err{border-color:var(--red);color:var(--red)}

/* empty */
.empty{text-align:center;padding:60px 20px;color:var(--muted);
       font-size:13px;line-height:2}

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

<!-- header -->
<header>
  <h1>⬡ ClassLabeller</h1>
  <div class="hstats" id="hstats">Loading…</div>
  <div class="assign-bar">
    <span id="sel-count"></span>
    <div class="cls-wrap">
      <input id="cls-input" type="text" placeholder="class name…" autocomplete="off"
             oninput="filterDrop()" onkeydown="clsKeydown(event)">
      <div id="cls-drop"></div>
    </div>
    <button class="btn btn-primary" id="assign-btn" disabled onclick="doAssign()">Assign →</button>
    <button class="btn btn-ghost" onclick="clearSel()">Clear</button>
  </div>
</header>

<!-- sidebar -->
<aside id="sidebar">
  <div class="side-title">Color groups</div>
  <div id="bucket-list"></div>
</aside>

<!-- main -->
<main>
  <div class="toolbar">
    <input class="filter-input" type="text" id="name-filter"
           placeholder="filter by name…" oninput="applyFilter()">
    <div class="sel-bar" id="sel-bar">
      <span id="sel-bar-txt"></span>
      <button class="btn btn-ghost" style="padding:3px 10px;font-size:11px"
              onclick="selectAll()">Select all visible</button>
    </div>
    <div class="toolbar-right">
      <span id="visible-count"></span>
    </div>
  </div>

  <!-- virtual scroll host -->
  <div class="scroll-host" id="scroll-host">
    <div id="spacer"></div>
    <div id="grid-rows"></div>

    <!-- loading state -->
    <div class="loading" id="loading">
      <div class="spinner"></div>
      <h2>Class Labeller</h2>
      <p id="load-msg">Building thumbnail cache…</p>
      <div class="bar-wrap"><div class="bar-fill" id="bar-fill"></div></div>
    </div>
  </div>
</main>

<div class="toast" id="toast"></div>

<script>
// ─── state ────────────────────────────────────────────────────────────────
const THUMB  = parseInt(getComputedStyle(document.documentElement)
               .getPropertyValue('--thumb')) || 112;
const GAP    = 6;
const PAD    = 28;

let ALL      = [];
let VISIBLE  = [];
let SEL      = new Set();
let CLASSES  = [];
let BUCKETS  = [];
let activeBucket = '__all__';
let nameFilter   = '';
let cols         = 1;

const CARD_W   = THUMB + 4;
const CARD_H   = THUMB + 22;
const ROW_H    = CARD_H + GAP;
const OVERSCAN = 8;

let lastFirstRow = -1;
let lastLastRow  = -1;
const loadedThumbs = new Set();

// ─── boot ─────────────────────────────────────────────────────────────────
async function boot() {
  try {
    setLoadMsg('Fetching metadata...', 10);
    const [buckRes, clsRes] = await Promise.all([
      fetch('/api/buckets'), fetch('/api/classes')
    ]);
    BUCKETS = await buckRes.json();
    CLASSES = await clsRes.json();
    buildSidebar();
    populateClsDrop();

    setLoadMsg('Loading image list...', 30);
    await loadAll();

    document.getElementById('loading').style.display = 'none';
    applyFilter();
    updateStats();
  } catch(e) {
    document.getElementById('load-msg').textContent = 'Error: ' + e.message;
    console.error(e);
  }
}

function setLoadMsg(msg, pct) {
  document.getElementById('load-msg').textContent = msg;
  document.getElementById('bar-fill').style.width = pct + '%';
}

async function loadAll() {
  let offset = 0;
  const LIMIT = 2000;
  let total = Infinity;
  while (offset < total) {
    const r = await fetch('/api/meta?offset=' + offset + '&limit=' + LIMIT);
    const d = await r.json();
    total = d.total;
    ALL.push(...d.items);
    offset += LIMIT;
    const pct = Math.round(30 + (Math.min(offset, total) / total) * 60);
    setLoadMsg('Loaded ' + Math.min(offset, total).toLocaleString() + ' / ' + total.toLocaleString() + ' images...', pct);
  }
}

// ─── sidebar ──────────────────────────────────────────────────────────────
const SWATCH = {
  Red:'#e05050',Orange:'#e09050',Yellow:'#d4c040',Green:'#50a050',
  Cyan:'#30b0b0',Blue:'#4060d0',Purple:'#8060c0',Pink:'#c060a0',
  White:'#e0e0e0',Gray:'#808080',Black:'#303030'
};

function buildSidebar() {
  const list = document.getElementById('bucket-list');
  const total = BUCKETS.reduce((a,b)=>a+b.count,0);
  let html = `<button class="bkt-btn active" data-bkt="__all__" onclick="setBucket(this)">
    <span class="bkt-left"><span class="bkt-swatch" style="background:linear-gradient(135deg,#e05050,#50a050,#4060d0)"></span>All</span>
    <span class="bkt-count">${total.toLocaleString()}</span></button>`;
  for (const b of BUCKETS) {
    html += `<button class="bkt-btn" data-bkt="${b.name}" onclick="setBucket(this)">
      <span class="bkt-left">
        <span class="bkt-swatch" style="background:${SWATCH[b.name]||'#666'}"></span>
        ${b.name}
      </span>
      <span class="bkt-count">${b.count.toLocaleString()}</span>
    </button>`;
  }
  list.innerHTML = html;
}

function setBucket(btn) {
  document.querySelectorAll('.bkt-btn').forEach(b=>b.classList.remove('active'));
  btn.classList.add('active');
  activeBucket = btn.dataset.bkt;
  SEL.clear();
  applyFilter();
}

// ─── filter ───────────────────────────────────────────────────────────────
function applyFilter() {
  const nf = nameFilter.toLowerCase();
  VISIBLE = ALL.filter(m => {
    if (activeBucket !== '__all__' && m.bucket !== activeBucket) return false;
    if (nf && !m.name.toLowerCase().includes(nf)) return false;
    return true;
  });
  document.getElementById('visible-count').textContent =
    VISIBLE.length.toLocaleString() + ' images';
  computeLayout();
  renderViewport(true);
  updateSelBar();
}

document.getElementById('name-filter').addEventListener('input', e => {
  nameFilter = e.target.value;
  applyFilter();
});

// ─── layout ───────────────────────────────────────────────────────────────
function computeLayout() {
  const w  = document.getElementById('scroll-host').clientWidth - PAD;
  cols     = Math.max(1, Math.floor((w + GAP) / (CARD_W + GAP)));
  const rows = Math.ceil(VISIBLE.length / cols);
  const totalH = rows * ROW_H;
  document.getElementById('spacer').style.height = totalH + 'px';
}

window.addEventListener('resize', () => { computeLayout(); renderViewport(true); });

// ─── virtual scroll ───────────────────────────────────────────────────────
const scrollHost = document.getElementById('scroll-host');
let rafPending   = false;

//scrollHost.addEventListener('scroll', () => {
//  if (!rafPending) {
//    rafPending = true;
//    requestAnimationFrame(() => { renderViewport(true); rafPending = false; });
//  }
//});

scrollHost.addEventListener('scroll', () => {
  if (!rafPending) {
    rafPending = true;
    requestAnimationFrame(() => { renderViewport(true); rafPending = false; });
  }
}, { passive: true });

// Image loading queue — limits concurrent fetches
const IMG_CONCURRENCY = 12;
let   loadingNow      = 0;
const loadQueue       = [];

//function enqueueImg(img, src) {
//  const doLoad = () => {
//    loadingNow++;
//    img.src = src;
//    img.onload = img.onerror = () => {
//      loadingNow--;
//      if (loadQueue.length) loadQueue.shift()();
//    };
//  };
//  if (loadingNow < IMG_CONCURRENCY) doLoad();
//  else loadQueue.push(doLoad);
//}

function enqueueImg(img, src) {
  const id = src.split('/').pop();
  if (loadedThumbs.has(id)) {
    img.src = src;
    return;
  }
  const doLoad = () => {
    loadingNow++;
    img.src = src;
    img.onload = img.onerror = () => {
      loadedThumbs.add(id);
      loadingNow--;
      if (loadQueue.length) loadQueue.shift()();
    };
  };
  if (loadingNow < IMG_CONCURRENCY) doLoad();
  else loadQueue.push(doLoad);
}

//function renderViewport() {
//  const host     = scrollHost;
//  const scrollY  = host.scrollTop;
//  const vpH      = host.clientHeight;
//  const firstRow = Math.max(0, Math.floor(scrollY / ROW_H) - OVERSCAN);
//  const lastRow  = Math.min(
//    Math.ceil(VISIBLE.length / cols),
//    Math.ceil((scrollY + vpH) / ROW_H) + OVERSCAN
//  );
//
//  const container = document.getElementById('grid-rows');
//  container.style.top = (firstRow * ROW_H) + 'px';
//
//  // Build HTML for visible rows
//  let html = '';
//  let prevBucket = null;
//  for (let row = firstRow; row < lastRow; row++) {
//    const startIdx = row * cols;
//    if (startIdx >= VISIBLE.length) break;
//
//    // Bucket separator
//    const firstInRow = VISIBLE[startIdx];
//    if (firstInRow && firstInRow.bucket !== prevBucket && activeBucket === '__all__') {
//      html += `<div class="bucket-label">${firstInRow.bucket}</div>`;
//      prevBucket = firstInRow.bucket;
//    }
//
//    html += '<div class="grid-row">';
//    for (let c = 0; c < cols; c++) {
//      const idx  = startIdx + c;
//      if (idx >= VISIBLE.length) break;
//      const m    = VISIBLE[idx];
//      const sel  = SEL.has(m.id) ? ' sel' : '';
//      html += `<div class="card${sel}" id="c-${m.id}"
//                    onclick="cardClick(event,'${m.id}')"
//                    data-id="${m.id}">
//        <div class="chk">${SEL.has(m.id)?'✓':''}</div>
//        <img data-src="/thumb/${m.id}" alt="${m.name}">
//        <div class="cname" title="${m.name}">${m.name}</div>
//      </div>`;
//    }
//    html += '</div>';
//  }
//
//  if (!html && VISIBLE.length === 0) {
//    html = `<div class="empty">No images in this group.<br>
//            Try a different filter or colour group.</div>`;
//  }
//
//  container.innerHTML = html;
//
//  // Lazy-load images via queue
//  container.querySelectorAll('img[data-src]').forEach(img => {
//    enqueueImg(img, img.dataset.src);
//    img.removeAttribute('data-src');
//  });
//}

function renderViewport(force = false) {
  const host     = scrollHost;
  const scrollY  = host.scrollTop;
  const vpH      = host.clientHeight;
  const firstRow = Math.max(0, Math.floor(scrollY / ROW_H) - OVERSCAN);
  const lastRow  = Math.min(
    Math.ceil(VISIBLE.length / cols),
    Math.ceil((scrollY + vpH) / ROW_H) + OVERSCAN
  );

  if (!force && firstRow === lastFirstRow && lastRow === lastLastRow) return;
  lastFirstRow = firstRow;
  lastLastRow  = lastRow;

  const container = document.getElementById('grid-rows');
  container.style.top = (firstRow * ROW_H) + 'px';

  let html = '';
  let prevBucket = null;
  for (let row = firstRow; row < lastRow; row++) {
    const startIdx = row * cols;
    if (startIdx >= VISIBLE.length) break;

    const firstInRow = VISIBLE[startIdx];
    if (firstInRow && firstInRow.bucket !== prevBucket && activeBucket === '__all__') {
      html += `<div class="bucket-label">${firstInRow.bucket}</div>`;
      prevBucket = firstInRow.bucket;
    }

    html += '<div class="grid-row">';
    for (let c = 0; c < cols; c++) {
      const idx = startIdx + c;
      if (idx >= VISIBLE.length) break;
      const m   = VISIBLE[idx];
      const sel = SEL.has(m.id) ? ' sel' : '';
      const url = `/thumb/${m.id}`;
      const img = loadedThumbs.has(m.id)
        ? `<img src="${url}" alt="${m.name}">`
        : `<img data-src="${url}" alt="${m.name}">`;
      html += `<div class="card${sel}" id="c-${m.id}"
                    onclick="cardClick(event,'${m.id}')"
                    data-id="${m.id}">
        <div class="chk">${SEL.has(m.id)?'✓':''}</div>
        ${img}
        <div class="cname" title="${m.name}">${m.name}</div>
      </div>`;
    }
    html += '</div>';
  }

  if (!html && VISIBLE.length === 0) {
    html = `<div class="empty">No images in this group.<br>
            Try a different filter or colour group.</div>`;
  }

  container.innerHTML = html;

  container.querySelectorAll('img[data-src]').forEach(img => {
    enqueueImg(img, img.dataset.src);
    img.removeAttribute('data-src');
  });
}

// ─── selection ────────────────────────────────────────────────────────────
let lastClickIdx = null;

function cardClick(evt, id) {
  const idx   = VISIBLE.findIndex(m => m.id === id);
  const card  = document.getElementById('c-' + id);

  if (evt.shiftKey && lastClickIdx !== null) {
    // Range select
    const lo = Math.min(lastClickIdx, idx);
    const hi = Math.max(lastClickIdx, idx);
    for (let i = lo; i <= hi; i++) SEL.add(VISIBLE[i].id);
  } else if (evt.ctrlKey || evt.metaKey) {
    SEL.has(id) ? SEL.delete(id) : SEL.add(id);
  } else {
    SEL.has(id) ? SEL.delete(id) : SEL.add(id);
  }
  lastClickIdx = idx;
  updateSelBar();
  renderViewport(true);
}

function clearSel() { SEL.clear(); updateSelBar(); renderViewport(true); }

function selectAll() {
  VISIBLE.forEach(m => SEL.add(m.id));
  updateSelBar();
  renderViewport(true);
}

function updateSelBar() {
  const n   = SEL.size;
  const bar = document.getElementById('sel-bar');
  bar.classList.toggle('on', n > 0);
  document.getElementById('sel-bar-txt').textContent = `${n.toLocaleString()} selected`;
  document.getElementById('sel-count').textContent   = n > 0 ? `${n} selected` : '';
  document.getElementById('assign-btn').disabled     = n === 0;
}

// ─── class dropdown ───────────────────────────────────────────────────────
function populateClsDrop() {
  const drop = document.getElementById('cls-drop');
  drop.innerHTML = CLASSES.map(c =>
    `<div class="cls-opt" onclick="pickCls('${c}')">${c}</div>`
  ).join('');
}

function filterDrop() {
  const val  = document.getElementById('cls-input').value.toLowerCase();
  const drop = document.getElementById('cls-drop');
  const opts = CLASSES.filter(c => c.toLowerCase().includes(val));
  drop.innerHTML = opts.map(c =>
    `<div class="cls-opt" onclick="pickCls('${c}')">${c}</div>`
  ).join('');
  drop.classList.toggle('open', opts.length > 0 && val.length > 0);
  document.getElementById('assign-btn').disabled = SEL.size === 0;
}

function pickCls(name) {
  document.getElementById('cls-input').value = name;
  document.getElementById('cls-drop').classList.remove('open');
}

function clsKeydown(e) {
  if (e.key === 'Enter') doAssign();
  if (e.key === 'Escape') document.getElementById('cls-drop').classList.remove('open');
}

document.addEventListener('click', e => {
  if (!(e.target instanceof Element) || !e.target.closest('.cls-wrap'))
    document.getElementById('cls-drop').classList.remove('open');
});

// ─── assign ───────────────────────────────────────────────────────────────
async function doAssign() {
  const cls = document.getElementById('cls-input').value.trim();
  if (!cls) { toast('Enter a class name', 'err'); return; }
  if (SEL.size === 0) { toast('Nothing selected', 'err'); return; }

  const ids = [...SEL];
  const res = await fetch('/api/assign', {
    method: 'POST',
    headers: {'Content-Type':'application/json'},
    body: JSON.stringify({ids, class_name: cls})
  });
  const d = await res.json();

  // Remove from local state
  ALL = ALL.filter(m => !SEL.has(m.id));
  SEL.clear();
  document.getElementById('cls-input').value = '';

  // Add class to dropdown if new
  if (!CLASSES.includes(cls)) {
    CLASSES.push(cls);
    CLASSES.sort();
    populateClsDrop();
  }

  applyFilter();
  updateStats();
  toast(`✓ Moved ${d.moved} images to ${cls}/`, 'ok');
}

// ─── stats ────────────────────────────────────────────────────────────────
function updateStats() {
  document.getElementById('hstats').innerHTML =
    `<b>${ALL.length.toLocaleString()}</b> remaining · <b>${CLASSES.length}</b> classes`;
}

// ─── toast ────────────────────────────────────────────────────────────────
let toastTimer = null;
function toast(msg, type) {
  if (!type) type = 'ok';
  const el = document.getElementById('toast');
  el.textContent = msg;
  el.className = `toast ${type} show`;
  if (toastTimer) clearTimeout(toastTimer);
  toastTimer = setTimeout(() => el.classList.remove('show'), 3000);
}

// ─── keyboard shortcuts ───────────────────────────────────────────────────
document.addEventListener('keydown', function(e) {
  const tag = document.activeElement && document.activeElement.id;
  if (tag === 'cls-input' || tag === 'name-filter') return;
  if (e.key === 'a' && (e.ctrlKey || e.metaKey)) { e.preventDefault(); selectAll(); }
  if (e.key === 'Escape') clearSel();
  if (e.key === 'Enter' && SEL.size > 0) doAssign();
});

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

# ─────────────────────────── main ────────────────────────────────────────────

def main():
    args = ARGS

    source_dir = Path(args.source)
    cache_dir  = Path(args.cache_dir)
    data_dir   = Path(args.data_dir)

    if not source_dir.exists():
        print(f"[error] Source directory not found: {source_dir}")
        print("        Create it and add your unlabelled crops, then re-run.")
        return

    data_dir.mkdir(parents=True, exist_ok=True)

    print(f"""
╔══════════════════════════════════════════╗
║         Class Labeller                  ║
╠══════════════════════════════════════════╣
  Source   : {source_dir}
  Data dir : {data_dir}
  Cache    : {cache_dir}
  Workers  : {args.workers}
  Thumb    : {args.thumb_size}px
╚══════════════════════════════════════════╝
""")

    # Build / update thumbnail cache (blocking, before server starts)
    meta = build_metadata(source_dir, cache_dir, args.thumb_size, args.workers)

    STATE["meta"] = meta

    import threading
    import webbrowser
    url = f"http://localhost:{args.port}"

    def _open():
        time.sleep(1.0)
        webbrowser.open(url)
    threading.Thread(target=_open, daemon=True).start()

    print(f"\n[server] Starting at {url}")
    print("  Keyboard shortcuts:")
    print("    Ctrl+A        → select all visible")
    print("    Shift+click   → range select")
    print("    Ctrl+click    → toggle single")
    print("    Enter         → assign selected")
    print("    Esc           → clear selection")
    print("\n  Press Ctrl+C to stop.\n")

    flask_app.run(host="0.0.0.0", port=args.port,
                  threaded=True, use_reloader=False)


if __name__ == "__main__":
    main()