"""Generate an interactive HTML curation tool for check mark templates.

Reads all raw template PNGs, embeds them as base64 in a self-contained HTML
file.  Open the HTML in a browser to crop/skip templates interactively, then
click Export to download a ZIP of curated crops.

Usage:
    python datasets/checkbox/tools/curate_templates.py
    # Then open temp/curate_templates.html in a browser
"""

import base64
import json
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent
RAW_DIR = ROOT / "templates" / "raw"
OUTPUT_HTML = Path(__file__).resolve().parent.parent.parent.parent / "temp" / "curate_templates.html"


def generate_html():
    OUTPUT_HTML.parent.mkdir(parents=True, exist_ok=True)

    # Load manifest
    manifest_path = RAW_DIR / "manifest.json"
    with open(manifest_path) as f:
        manifest = json.load(f)

    # Encode all images as base64
    images_data = []
    for entry in manifest:
        img_path = RAW_DIR / entry["filename"]
        if not img_path.exists():
            continue
        b64 = base64.b64encode(img_path.read_bytes()).decode("ascii")
        images_data.append({
            "filename": entry["filename"],
            "source_doc": entry["source_doc"],
            "b64": b64,
        })

    logger.info("Embedded %d images into HTML", len(images_data))

    html = _build_html(images_data)
    OUTPUT_HTML.write_text(html)
    logger.info("Wrote: %s", OUTPUT_HTML)
    logger.info("Open in browser, crop templates, then click Export.")


def _build_html(images_data: list) -> str:
    images_json = json.dumps(images_data)

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<title>Check Mark Template Curation</title>
<script src="https://cdnjs.cloudflare.com/ajax/libs/jszip/3.10.1/jszip.min.js"></script>
<style>
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; background: #1a1a2e; color: #e0e0e0; display: flex; flex-direction: column; align-items: center; min-height: 100vh; }}
.header {{ padding: 20px; text-align: center; width: 100%; background: #16213e; border-bottom: 1px solid #0f3460; }}
.header h1 {{ font-size: 20px; margin-bottom: 8px; }}
.stats {{ font-size: 14px; color: #a0a0c0; }}
.stats span {{ margin: 0 12px; }}
.stats .kept {{ color: #4ade80; }}
.stats .skipped {{ color: #f87171; }}
.main {{ flex: 1; display: flex; flex-direction: column; align-items: center; justify-content: center; padding: 30px; }}
.canvas-wrap {{ position: relative; border: 2px solid #0f3460; border-radius: 8px; overflow: hidden; background: #fff; cursor: crosshair; }}
canvas {{ display: block; }}
.info {{ margin-top: 16px; font-size: 13px; color: #808090; text-align: center; }}
.controls {{ padding: 20px; text-align: center; width: 100%; background: #16213e; border-top: 1px solid #0f3460; }}
.controls .hint {{ font-size: 13px; color: #808090; margin-bottom: 12px; }}
kbd {{ background: #2a2a4a; border: 1px solid #444; border-radius: 4px; padding: 2px 8px; font-family: monospace; font-size: 12px; }}
button {{ background: #0f3460; color: #e0e0e0; border: 1px solid #1a5276; border-radius: 6px; padding: 8px 24px; font-size: 14px; cursor: pointer; margin: 0 8px; }}
button:hover {{ background: #1a5276; }}
button.export {{ background: #065f46; border-color: #047857; }}
button.export:hover {{ background: #047857; }}
button:disabled {{ opacity: 0.4; cursor: default; }}
.done {{ text-align: center; }}
.done h2 {{ font-size: 24px; margin-bottom: 12px; color: #4ade80; }}
.selection-box {{ position: absolute; border: 2px dashed #f59e0b; background: rgba(245, 158, 11, 0.15); pointer-events: none; }}
.preview-row {{ display: flex; flex-wrap: wrap; gap: 6px; margin-top: 16px; justify-content: center; max-width: 900px; }}
.preview-row img {{ height: 40px; border: 1px solid #333; border-radius: 3px; background: #fff; }}
</style>
</head>
<body>

<div class="header">
  <h1>Check Mark Template Curation</h1>
  <div class="stats">
    <span id="progress">0 / 0</span>
    <span class="kept" id="kept-count">0 kept</span>
    <span class="skipped" id="skipped-count">0 skipped</span>
  </div>
</div>

<div class="main" id="main-area">
  <div class="canvas-wrap" id="canvas-wrap">
    <canvas id="canvas"></canvas>
    <div class="selection-box" id="sel-box" style="display:none;"></div>
  </div>
  <div class="info" id="info-line"></div>
</div>

<div class="controls">
  <div class="hint">
    <kbd>Click + drag</kbd> to crop the check mark &nbsp;|&nbsp;
    <kbd>Space</kbd> skip &nbsp;|&nbsp;
    <kbd>Z</kbd> undo last
  </div>
  <button id="skip-btn" onclick="skip()">Skip (Space)</button>
  <button id="undo-btn" onclick="undo()">Undo (Z)</button>
  <button id="export-btn" class="export" onclick="doExport()" disabled>Export ZIP</button>
</div>

<script>
const IMAGES = {images_json};
const SCALE = 20;
const MIN_DRAG = 4;

let idx = 0;
let kept = [];
let skipped = 0;
let history = [];

let dragging = false;
let startX = 0, startY = 0;

const canvas = document.getElementById('canvas');
const ctx = canvas.getContext('2d');
const wrap = document.getElementById('canvas-wrap');
const selBox = document.getElementById('sel-box');
const mainArea = document.getElementById('main-area');

function loadCurrent() {{
  if (idx >= IMAGES.length) {{
    showDone();
    return;
  }}
  const img = new Image();
  img.onload = function() {{
    const w = img.naturalWidth;
    const h = img.naturalHeight;
    canvas.width = w * SCALE;
    canvas.height = h * SCALE;
    ctx.imageSmoothingEnabled = false;
    ctx.drawImage(img, 0, 0, w * SCALE, h * SCALE);
    wrap.style.width = (w * SCALE) + 'px';
    wrap.style.height = (h * SCALE) + 'px';

    // Store on canvas for cropping
    canvas._img = img;
    canvas._origW = w;
    canvas._origH = h;

    document.getElementById('info-line').textContent =
      IMAGES[idx].filename + '  (' + w + 'x' + h + ')  from doc ' + IMAGES[idx].source_doc;
    updateStats();
  }};
  img.src = 'data:image/png;base64,' + IMAGES[idx].b64;
}}

function updateStats() {{
  document.getElementById('progress').textContent = (idx + 1) + ' / ' + IMAGES.length;
  document.getElementById('kept-count').textContent = kept.length + ' kept';
  document.getElementById('skipped-count').textContent = skipped + ' skipped';
  document.getElementById('export-btn').disabled = kept.length === 0;
}}

function skip() {{
  if (idx >= IMAGES.length) return;
  history.push({{ action: 'skip', idx: idx }});
  skipped++;
  idx++;
  loadCurrent();
}}

function undo() {{
  if (history.length === 0) return;
  const last = history.pop();
  if (last.action === 'skip') {{
    skipped--;
    idx = last.idx;
  }} else if (last.action === 'crop') {{
    kept.pop();
    idx = last.idx;
  }}
  // If we were on the done screen, restore canvas
  const doneEl = document.getElementById('done-screen');
  if (doneEl) doneEl.remove();
  wrap.style.display = '';
  document.getElementById('info-line').style.display = '';
  loadCurrent();
}}

// Mouse events for crop selection
canvas.addEventListener('mousedown', function(e) {{
  if (idx >= IMAGES.length) return;
  dragging = true;
  const rect = canvas.getBoundingClientRect();
  startX = e.clientX - rect.left;
  startY = e.clientY - rect.top;
  selBox.style.display = 'block';
  selBox.style.left = startX + 'px';
  selBox.style.top = startY + 'px';
  selBox.style.width = '0px';
  selBox.style.height = '0px';
}});

canvas.addEventListener('mousemove', function(e) {{
  if (!dragging) return;
  const rect = canvas.getBoundingClientRect();
  const curX = e.clientX - rect.left;
  const curY = e.clientY - rect.top;
  const x = Math.min(startX, curX);
  const y = Math.min(startY, curY);
  const w = Math.abs(curX - startX);
  const h = Math.abs(curY - startY);
  selBox.style.left = x + 'px';
  selBox.style.top = y + 'px';
  selBox.style.width = w + 'px';
  selBox.style.height = h + 'px';
}});

canvas.addEventListener('mouseup', function(e) {{
  if (!dragging) return;
  dragging = false;
  selBox.style.display = 'none';

  const rect = canvas.getBoundingClientRect();
  const endX = e.clientX - rect.left;
  const endY = e.clientY - rect.top;

  // Selection in scaled canvas coords
  let sx = Math.min(startX, endX);
  let sy = Math.min(startY, endY);
  let sw = Math.abs(endX - startX);
  let sh = Math.abs(endY - startY);

  // Convert to original image coords
  const origW = canvas._origW;
  const origH = canvas._origH;
  let ox = Math.max(0, Math.floor(sx / SCALE));
  let oy = Math.max(0, Math.floor(sy / SCALE));
  let ow = Math.min(origW - ox, Math.ceil(sw / SCALE));
  let oh = Math.min(origH - oy, Math.ceil(sh / SCALE));

  if (ow < 2 || oh < 2) return; // Too small, ignore

  // Crop from original image
  const cropCanvas = document.createElement('canvas');
  cropCanvas.width = ow;
  cropCanvas.height = oh;
  const cropCtx = cropCanvas.getContext('2d');
  cropCtx.drawImage(canvas._img, ox, oy, ow, oh, 0, 0, ow, oh);

  const dataUrl = cropCanvas.toDataURL('image/png');
  kept.push({{
    filename: 'curated_' + String(kept.length).padStart(4, '0') + '.png',
    sourceFile: IMAGES[idx].filename,
    sourceDoc: IMAGES[idx].source_doc,
    dataUrl: dataUrl,
    cropRegion: [ox, oy, ow, oh],
  }});

  history.push({{ action: 'crop', idx: idx }});
  idx++;
  loadCurrent();
}});

// Keyboard
document.addEventListener('keydown', function(e) {{
  if (e.code === 'Space') {{
    e.preventDefault();
    skip();
  }} else if (e.code === 'KeyZ' && !e.metaKey && !e.ctrlKey) {{
    undo();
  }}
}});

function showDone() {{
  wrap.style.display = 'none';
  document.getElementById('info-line').style.display = 'none';
  updateStats();

  const div = document.createElement('div');
  div.id = 'done-screen';
  div.className = 'done';
  div.innerHTML = '<h2>All done!</h2><p>' + kept.length + ' templates kept, ' + skipped + ' skipped.</p>';

  if (kept.length > 0) {{
    const row = document.createElement('div');
    row.className = 'preview-row';
    for (const k of kept) {{
      const img = document.createElement('img');
      img.src = k.dataUrl;
      img.title = k.filename;
      row.appendChild(img);
    }}
    div.appendChild(row);
  }}

  mainArea.appendChild(div);
}}

async function doExport() {{
  if (kept.length === 0) return;
  const btn = document.getElementById('export-btn');
  btn.textContent = 'Exporting...';
  btn.disabled = true;

  const zip = new JSZip();
  const manifest = [];

  for (const k of kept) {{
    // dataUrl → binary
    const b64 = k.dataUrl.split(',')[1];
    zip.file(k.filename, b64, {{ base64: true }});
    manifest.push({{
      filename: k.filename,
      source_file: k.sourceFile,
      source_doc: k.sourceDoc,
      crop_region: k.cropRegion,
    }});
  }}

  zip.file('manifest.json', JSON.stringify(manifest, null, 2));

  const blob = await zip.generateAsync({{ type: 'blob' }});
  const url = URL.createObjectURL(blob);
  const a = document.createElement('a');
  a.href = url;
  a.download = 'curated_templates.zip';
  a.click();
  URL.revokeObjectURL(url);

  btn.textContent = 'Exported!';
  setTimeout(() => {{
    btn.textContent = 'Export ZIP';
    btn.disabled = kept.length === 0;
  }}, 2000);
}}

// Init
loadCurrent();
</script>
</body>
</html>""";


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    generate_html()
