#!/usr/bin/env python3
"""
Probe Tip Annotation Tool

Shows cropped probe images and lets annotators click on the probe tip.
Saves (x, y) in crop pixel coordinates.

Usage:
    python3 annotate_tips.py --data-dir ./crop_cache [--port 5555]

The --data-dir should contain:
  - manifest.json  (generated by the pre-crop step)
  - *.png          (cropped images)
"""

from flask import Flask, Response, request, jsonify, render_template_string, send_file
import json
import os
from datetime import datetime
import argparse
import threading

# HF repo-based persistence (optional, for Spaces deployment)
# Use a separate DATASET repo so syncs don't trigger Space rebuilds
HF_REPO_ID = os.environ.get('SPACE_ID')  # auto-set on HF Spaces
HF_DATA_REPO = 'chfeng/probe-tip-annotations-data'  # dataset repo for annotations
_hf_api = None
_save_lock = threading.Lock()

def _get_hf_api():
    global _hf_api
    if _hf_api is None and HF_REPO_ID:
        try:
            from huggingface_hub import HfApi
            _hf_api = HfApi()
        except ImportError:
            pass
    return _hf_api

def _sync_to_repo(annotator):
    """Upload annotation file to the HF Space repo for persistence."""
    api = _get_hf_api()
    if not api:
        return
    path = os.path.join(ANNOTATION_DIR, f'{annotator}.json')
    if not os.path.exists(path):
        return
    try:
        api.upload_file(
            path_or_fileobj=path,
            path_in_repo=f'{annotator}.json',
            repo_id=HF_DATA_REPO,
            repo_type='dataset',
            commit_message=f'Update {annotator} annotations',
        )
    except Exception as e:
        print(f'Warning: failed to sync {annotator}.json to repo: {e}')

def _load_annotations_from_repo():
    """On startup, download any existing annotations from the repo."""
    api = _get_hf_api()
    if not api:
        return
    os.makedirs(ANNOTATION_DIR, exist_ok=True)
    for name in ANNOTATORS:
        try:
            path = api.hf_hub_download(
                repo_id=HF_DATA_REPO,
                filename=f'{name}.json',
                repo_type='dataset',
            )
            import shutil
            shutil.copy(path, os.path.join(ANNOTATION_DIR, f'{name}.json'))
            print(f'  Loaded {name}.json from repo')
        except Exception:
            pass

app = Flask(__name__)

DATA_DIR = './crop_cache'
ANNOTATION_DIR = './tip_annotations'
ANNOTATORS = ['andrew', 'ayush', 'xuanchen', 'chao']

manifest = {}


def load_manifest():
    global manifest
    path = os.path.join(DATA_DIR, 'manifest.json')
    with open(path) as f:
        manifest = json.load(f)
    print(f"Loaded manifest: {manifest['total']} samples")
    for name in ANNOTATORS:
        print(f"  {name}: {len(manifest['assignments'][name])} images")


def get_annotations(annotator):
    path = os.path.join(ANNOTATION_DIR, f'{annotator}.json')
    if os.path.exists(path):
        with open(path) as f:
            return json.load(f)
    return {}


def save_annotation_file(annotator, data):
    path = os.path.join(ANNOTATION_DIR, f'{annotator}.json')
    with open(path, 'w') as f:
        json.dump(data, f, indent=2)
    threading.Thread(target=_sync_to_repo, args=(annotator,), daemon=True).start()


# ─── Routes ───

@app.route('/')
def index():
    stats = {}
    for name in ANNOTATORS:
        anns = get_annotations(name)
        total = len(manifest['assignments'][name])
        done = sum(1 for v in anns.values() if not v.get('skipped', False))
        skipped = sum(1 for v in anns.values() if v.get('skipped', False))
        stats[name] = {
            'total': total, 'done': done, 'skipped': skipped,
            'pct': int(100 * (done + skipped) / total) if total > 0 else 0
        }
    return render_template_string(INDEX_HTML, stats=stats, annotators=ANNOTATORS)


@app.route('/annotate/<annotator>')
def annotate(annotator):
    if annotator not in manifest['assignments']:
        return "Unknown annotator", 404
    return render_template_string(ANNOTATE_HTML, annotator=annotator)


@app.route('/api/samples/<annotator>')
def api_samples(annotator):
    if annotator not in manifest['assignments']:
        return jsonify({'error': 'Unknown annotator'}), 404
    samples = manifest['assignments'][annotator]
    anns = get_annotations(annotator)
    return jsonify({
        'samples': samples,
        'annotations': anns,
        'total': len(samples),
        'done': sum(1 for v in anns.values() if not v.get('skipped', False)),
        'skipped': sum(1 for v in anns.values() if v.get('skipped', False))
    })


@app.route('/api/crop/<stem>')
def api_crop(stem):
    path = os.path.join(DATA_DIR, f'{stem}.png')
    if os.path.exists(path):
        return send_file(path, mimetype='image/png')
    return "Not found", 404


@app.route('/api/save', methods=['POST'])
def api_save():
    data = request.json
    annotator = data['annotator']
    stem = data['stem']
    tip_x = data['tip_x']
    tip_y = data['tip_y']

    meta = manifest['metadata'].get(stem)
    if not meta:
        return jsonify({'error': 'unknown image'}), 404

    anns = get_annotations(annotator)
    anns[stem] = {
        'tip_x': round(tip_x, 1),
        'tip_y': round(tip_y, 1),
        'crop_x1': meta['crop_x1'],
        'crop_y1': meta['crop_y1'],
        'crop_w': meta['crop_w'],
        'crop_h': meta['crop_h'],
        'img_w': meta['img_w'],
        'img_h': meta['img_h'],
        'full_x': round(meta['crop_x1'] + tip_x, 1),
        'full_y': round(meta['crop_y1'] + tip_y, 1),
        'timestamp': datetime.now().isoformat()
    }

    save_annotation_file(annotator, anns)
    done = sum(1 for v in anns.values() if not v.get('skipped', False))
    return jsonify({'ok': True, 'done': done})


@app.route('/api/skip', methods=['POST'])
def api_skip():
    data = request.json
    annotator = data['annotator']
    stem = data['stem']

    anns = get_annotations(annotator)
    anns[stem] = {'skipped': True, 'timestamp': datetime.now().isoformat()}
    save_annotation_file(annotator, anns)
    done = sum(1 for v in anns.values() if not v.get('skipped', False))
    return jsonify({'ok': True, 'done': done})


@app.route('/api/delete', methods=['POST'])
def api_delete():
    data = request.json
    annotator = data['annotator']
    stem = data['stem']
    anns = get_annotations(annotator)
    if stem in anns:
        del anns[stem]
        save_annotation_file(annotator, anns)
    done = sum(1 for v in anns.values() if not v.get('skipped', False))
    return jsonify({'ok': True, 'done': done})


# ─── Inline HTML templates ───

INDEX_HTML = '''<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Probe Tip Annotation</title>
<style>
  * { margin: 0; padding: 0; box-sizing: border-box; }
  body {
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    background: #0f172a; color: #e2e8f0; min-height: 100vh;
  }
  .header {
    background: #1e293b; border-bottom: 1px solid #334155;
    padding: 24px 40px; text-align: center;
  }
  .header h1 { font-size: 28px; font-weight: 600; color: #f1f5f9; }
  .header p { color: #94a3b8; margin-top: 8px; font-size: 15px; }
  .stats-bar {
    display: flex; justify-content: center; gap: 32px;
    margin-top: 16px; font-size: 14px; color: #94a3b8;
  }
  .stats-bar span { color: #38bdf8; font-weight: 600; }
  .grid {
    display: grid; grid-template-columns: repeat(auto-fit, minmax(280px, 1fr));
    gap: 24px; max-width: 1200px; margin: 40px auto; padding: 0 40px;
  }
  .card {
    background: #1e293b; border: 1px solid #334155; border-radius: 12px;
    padding: 32px; transition: transform 0.15s, border-color 0.15s;
  }
  .card:hover { transform: translateY(-2px); border-color: #38bdf8; }
  .card h2 {
    font-size: 22px; font-weight: 600; color: #f1f5f9;
    text-transform: capitalize; margin-bottom: 20px;
  }
  .progress-wrap {
    background: #0f172a; border-radius: 8px; height: 10px;
    overflow: hidden; margin-bottom: 12px;
  }
  .progress-fill {
    height: 100%; border-radius: 8px; transition: width 0.3s;
  }
  .progress-fill.low { background: #ef4444; }
  .progress-fill.mid { background: #f59e0b; }
  .progress-fill.high { background: #22c55e; }
  .info { font-size: 14px; color: #94a3b8; margin-bottom: 4px; }
  .info span { color: #e2e8f0; font-weight: 500; }
  .btn {
    display: block; width: 100%; margin-top: 20px; padding: 12px;
    background: #2563eb; color: white; border: none; border-radius: 8px;
    font-size: 15px; font-weight: 500; cursor: pointer;
    text-decoration: none; text-align: center; transition: background 0.15s;
  }
  .btn:hover { background: #1d4ed8; }
</style>
</head>
<body>
  <div class="header">
    <h1>Probe Tip Annotation Tool</h1>
    <p>Click on the probe tip in each cropped image to annotate its (x, y) position</p>
    <div class="stats-bar">
      {% set total_done = stats.values()|map(attribute='done')|sum %}
      {% set total_all = stats.values()|map(attribute='total')|sum %}
      {% set total_skipped = stats.values()|map(attribute='skipped')|sum %}
      <div>Total: <span>{{ total_all }}</span> images</div>
      <div>Annotated: <span>{{ total_done }}</span></div>
      <div>Skipped: <span>{{ total_skipped }}</span></div>
      <div>Remaining: <span>{{ total_all - total_done - total_skipped }}</span></div>
    </div>
  </div>
  <div class="grid">
    {% for name in annotators %}
    {% set s = stats[name] %}
    <div class="card">
      <h2>{{ name }}</h2>
      <div class="progress-wrap">
        <div class="progress-fill {{ 'high' if s.pct > 66 else ('mid' if s.pct > 33 else 'low') }}"
             style="width: {{ s.pct }}%"></div>
      </div>
      <div class="info">Progress: <span>{{ s.pct }}%</span></div>
      <div class="info">Annotated: <span>{{ s.done }}</span> / {{ s.total }}</div>
      <div class="info">Skipped: <span>{{ s.skipped }}</span></div>
      <div class="info">Remaining: <span>{{ s.total - s.done - s.skipped }}</span></div>
      <a class="btn" href="/annotate/{{ name }}">Start Annotating</a>
    </div>
    {% endfor %}
  </div>
</body>
</html>'''

ANNOTATE_HTML = '''<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Annotate – {{ annotator }}</title>
<style>
  * { margin: 0; padding: 0; box-sizing: border-box; }
  body {
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    background: #0f172a; color: #e2e8f0; min-height: 100vh;
    display: flex; flex-direction: column;
  }
  .topbar {
    background: #1e293b; border-bottom: 1px solid #334155;
    padding: 12px 24px; display: flex; align-items: center;
    justify-content: space-between; flex-shrink: 0;
  }
  .topbar .left { display: flex; align-items: center; gap: 16px; }
  .topbar a.back {
    color: #94a3b8; text-decoration: none; font-size: 14px;
    padding: 6px 12px; border: 1px solid #334155; border-radius: 6px;
    transition: all 0.15s;
  }
  .topbar a.back:hover { color: #e2e8f0; border-color: #64748b; }
  .topbar h1 { font-size: 18px; font-weight: 600; text-transform: capitalize; }
  .progress-text { font-size: 14px; color: #94a3b8; }
  .progress-text span { color: #38bdf8; font-weight: 600; }
  .progress-mini {
    width: 200px; height: 6px; background: #0f172a;
    border-radius: 3px; overflow: hidden;
  }
  .progress-mini-fill {
    height: 100%; background: #22c55e; border-radius: 3px;
    transition: width 0.3s;
  }
  .main {
    flex: 1; display: flex; flex-direction: column;
    align-items: center; padding: 24px; overflow: auto;
  }
  .image-info {
    font-size: 13px; color: #64748b; margin-bottom: 12px; text-align: center;
  }
  .image-info .filename { color: #94a3b8; font-family: monospace; }
  .image-info .idx { color: #38bdf8; }
  .img-container {
    position: relative; cursor: crosshair; display: inline-block;
    border: 2px solid #334155; border-radius: 4px; overflow: hidden;
    background: #000;
  }
  .img-container img {
    display: block; max-width: 80vw; max-height: 65vh; object-fit: contain;
  }
  .marker {
    position: absolute; pointer-events: none;
    transform: translate(-50%, -50%); z-index: 10;
  }
  .marker .dot {
    width: 12px; height: 12px; background: #ef4444;
    border: 2px solid white; border-radius: 50%;
    position: absolute; top: 50%; left: 50%;
    transform: translate(-50%, -50%);
  }
  .marker .crosshair-h, .marker .crosshair-v {
    position: absolute; background: rgba(239, 68, 68, 0.5);
  }
  .marker .crosshair-h {
    width: 32px; height: 1px; top: 50%; left: 50%;
    transform: translate(-50%, -50%);
  }
  .marker .crosshair-v {
    width: 1px; height: 32px; top: 50%; left: 50%;
    transform: translate(-50%, -50%);
  }
  .coord-display {
    margin-top: 12px; font-size: 14px; color: #94a3b8; height: 20px;
  }
  .coord-display span { color: #f59e0b; font-family: monospace; font-weight: 600; }
  .controls { display: flex; gap: 12px; margin-top: 20px; align-items: center; }
  .btn {
    padding: 10px 24px; border: none; border-radius: 8px;
    font-size: 14px; font-weight: 500; cursor: pointer;
    transition: all 0.15s; display: flex; align-items: center; gap: 6px;
  }
  .btn:disabled { opacity: 0.4; cursor: not-allowed; }
  .btn-primary { background: #2563eb; color: white; }
  .btn-primary:hover:not(:disabled) { background: #1d4ed8; }
  .btn-secondary { background: #334155; color: #e2e8f0; }
  .btn-secondary:hover:not(:disabled) { background: #475569; }
  .btn-danger { background: #7f1d1d; color: #fca5a5; }
  .btn-danger:hover:not(:disabled) { background: #991b1b; }
  .btn-skip { background: #78350f; color: #fcd34d; }
  .btn-skip:hover:not(:disabled) { background: #92400e; }
  .shortcuts {
    margin-top: 16px; font-size: 12px; color: #475569; text-align: center;
  }
  .shortcuts kbd {
    background: #1e293b; border: 1px solid #334155; border-radius: 3px;
    padding: 1px 5px; font-family: monospace; color: #94a3b8;
  }
  .status-badge {
    display: inline-block; padding: 3px 10px; border-radius: 12px;
    font-size: 12px; font-weight: 600; margin-left: 8px;
  }
  .status-badge.done { background: #14532d; color: #86efac; }
  .status-badge.skipped { background: #78350f; color: #fcd34d; }
  .status-badge.new { background: #1e3a5f; color: #7dd3fc; }
  .loading { font-size: 18px; color: #64748b; margin-top: 100px; }
</style>
</head>
<body>
  <div class="topbar">
    <div class="left">
      <a class="back" href="/">Home</a>
      <h1>{{ annotator }}</h1>
    </div>
    <div style="display:flex;align-items:center;gap:16px;">
      <div class="progress-text">
        <span id="done-count">0</span> / <span id="total-count">0</span> done
      </div>
      <div class="progress-mini">
        <div class="progress-mini-fill" id="progress-fill" style="width:0%"></div>
      </div>
    </div>
  </div>
  <div class="main" id="main">
    <div class="loading" id="loading">Loading images...</div>
    <div id="content" style="display:none;text-align:center;">
      <div class="image-info">
        <span class="idx" id="img-idx">1/688</span>
        &middot;
        <span class="filename" id="img-name">image.png</span>
        <span class="status-badge new" id="status-badge">new</span>
      </div>
      <div class="img-container" id="img-container">
        <img id="crop-img" src="" draggable="false">
        <div class="marker" id="marker" style="display:none;">
          <div class="crosshair-h"></div>
          <div class="crosshair-v"></div>
          <div class="dot"></div>
        </div>
      </div>
      <div class="coord-display" id="coord-display"></div>
      <div class="controls">
        <button class="btn btn-secondary" id="btn-prev" disabled>Prev</button>
        <button class="btn btn-primary" id="btn-save" disabled>Save &amp; Next</button>
        <button class="btn btn-skip" id="btn-skip">Skip</button>
        <button class="btn btn-danger" id="btn-clear">Clear</button>
        <button class="btn btn-secondary" id="btn-next">Next</button>
      </div>
      <div class="shortcuts">
        <kbd>Enter</kbd> Save &amp; Next
        <kbd>S</kbd> Skip
        <kbd>&larr;</kbd> Prev
        <kbd>&rarr;</kbd> Next
        <kbd>C</kbd> Clear
        <kbd>G</kbd> Go to #
      </div>
    </div>
    <div id="all-done" style="display:none;text-align:center;margin-top:100px;">
      <h2 style="font-size:28px;color:#22c55e;margin-bottom:12px;">All done!</h2>
      <p style="color:#94a3b8;">You have annotated all assigned images.</p>
      <a class="btn btn-primary" href="/" style="display:inline-block;margin-top:20px;">Back to Home</a>
    </div>
  </div>
<script>
const ANNOTATOR = "{{ annotator }}";
let samples = [], annotations = {}, currentIdx = 0, currentTip = null;

async function init() {
  const resp = await fetch(`/api/samples/${ANNOTATOR}`);
  const data = await resp.json();
  samples = data.samples;
  annotations = data.annotations;
  document.getElementById('total-count').textContent = samples.length;
  updateProgress();
  currentIdx = samples.findIndex(s => !(s in annotations));
  if (currentIdx === -1) currentIdx = 0;
  document.getElementById('loading').style.display = 'none';
  document.getElementById('content').style.display = '';
  loadImage(currentIdx);
}

function updateProgress() {
  const done = Object.values(annotations).filter(a => !a.skipped).length;
  document.getElementById('done-count').textContent = done;
  document.getElementById('total-count').textContent = samples.length;
  const pct = samples.length > 0 ? (done / samples.length * 100) : 0;
  document.getElementById('progress-fill').style.width = pct + '%';
}

function loadImage(idx) {
  if (idx < 0 || idx >= samples.length) return;
  currentIdx = idx;
  currentTip = null;
  const stem = samples[idx];
  document.getElementById('img-idx').textContent = `${idx + 1}/${samples.length}`;
  document.getElementById('img-name').textContent = stem;
  document.getElementById('coord-display').innerHTML = '';
  document.getElementById('marker').style.display = 'none';
  document.getElementById('btn-save').disabled = true;
  document.getElementById('btn-prev').disabled = (idx === 0);
  const badge = document.getElementById('status-badge');
  const ann = annotations[stem];
  if (ann && !ann.skipped) {
    badge.textContent = 'done'; badge.className = 'status-badge done';
    currentTip = { x: ann.tip_x, y: ann.tip_y };
    document.getElementById('btn-save').disabled = false;
  } else if (ann && ann.skipped) {
    badge.textContent = 'skipped'; badge.className = 'status-badge skipped';
  } else {
    badge.textContent = 'new'; badge.className = 'status-badge new';
  }
  const img = document.getElementById('crop-img');
  img.onload = () => { if (currentTip) showMarker(currentTip.x, currentTip.y); };
  img.src = `/api/crop/${stem}`;
  if (idx + 1 < samples.length) { const n = new Image(); n.src = `/api/crop/${samples[idx+1]}`; }
}

function showMarker(cropX, cropY) {
  const img = document.getElementById('crop-img');
  const marker = document.getElementById('marker');
  const sx = img.clientWidth / img.naturalWidth;
  const sy = img.clientHeight / img.naturalHeight;
  marker.style.left = (cropX * sx) + 'px';
  marker.style.top = (cropY * sy) + 'px';
  marker.style.display = '';
  document.getElementById('coord-display').innerHTML =
    `Tip: (<span>${Math.round(cropX)}</span>, <span>${Math.round(cropY)}</span>) px in crop`;
}

document.getElementById('img-container').addEventListener('click', function(e) {
  const img = document.getElementById('crop-img');
  const rect = img.getBoundingClientRect();
  const sx = img.naturalWidth / img.clientWidth;
  const sy = img.naturalHeight / img.clientHeight;
  const cropX = (e.clientX - rect.left) * sx;
  const cropY = (e.clientY - rect.top) * sy;
  currentTip = { x: cropX, y: cropY };
  showMarker(cropX, cropY);
  document.getElementById('btn-save').disabled = false;
});

window.addEventListener('resize', () => { if (currentTip) showMarker(currentTip.x, currentTip.y); });

async function saveAndNext() {
  if (!currentTip) return;
  const stem = samples[currentIdx];
  await fetch('/api/save', {
    method: 'POST', headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ annotator: ANNOTATOR, stem, tip_x: currentTip.x, tip_y: currentTip.y })
  });
  annotations[stem] = { tip_x: currentTip.x, tip_y: currentTip.y };
  updateProgress(); goNext();
}

async function skipAndNext() {
  const stem = samples[currentIdx];
  await fetch('/api/skip', {
    method: 'POST', headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ annotator: ANNOTATOR, stem })
  });
  annotations[stem] = { skipped: true };
  updateProgress(); goNext();
}

function goNext() {
  if (currentIdx + 1 < samples.length) { loadImage(currentIdx + 1); }
  else {
    const remaining = samples.filter(s => !(s in annotations));
    if (remaining.length === 0) {
      document.getElementById('content').style.display = 'none';
      document.getElementById('all-done').style.display = '';
    }
  }
}

document.getElementById('btn-save').addEventListener('click', saveAndNext);
document.getElementById('btn-skip').addEventListener('click', skipAndNext);
document.getElementById('btn-prev').addEventListener('click', () => { if (currentIdx > 0) loadImage(currentIdx - 1); });
document.getElementById('btn-next').addEventListener('click', goNext);
document.getElementById('btn-clear').addEventListener('click', () => {
  currentTip = null;
  document.getElementById('marker').style.display = 'none';
  document.getElementById('coord-display').innerHTML = '';
  document.getElementById('btn-save').disabled = true;
});

document.addEventListener('keydown', (e) => {
  if (e.target.tagName === 'INPUT') return;
  switch(e.key) {
    case 'Enter': e.preventDefault(); if (currentTip) saveAndNext(); break;
    case 's': case 'S': e.preventDefault(); skipAndNext(); break;
    case 'ArrowLeft': e.preventDefault(); if (currentIdx > 0) loadImage(currentIdx - 1); break;
    case 'ArrowRight': e.preventDefault(); goNext(); break;
    case 'c': case 'C': e.preventDefault(); document.getElementById('btn-clear').click(); break;
    case 'g': case 'G':
      e.preventDefault();
      const num = prompt('Go to image # (1-' + samples.length + '):');
      if (num) { const n = parseInt(num) - 1; if (n >= 0 && n < samples.length) loadImage(n); }
      break;
  }
});

init();
</script>
</body>
</html>'''


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Probe Tip Annotation Tool')
    parser.add_argument('--port', type=int, default=5555)
    parser.add_argument('--data-dir', type=str, default='./crop_cache',
                        help='Directory with manifest.json and crop PNGs')
    parser.add_argument('--annotation-dir', type=str, default='./tip_annotations',
                        help='Directory to save annotation JSON files')
    args = parser.parse_args()

    DATA_DIR = os.path.abspath(args.data_dir)
    ANNOTATION_DIR = os.path.abspath(args.annotation_dir)
    os.makedirs(ANNOTATION_DIR, exist_ok=True)

    load_manifest()
    _load_annotations_from_repo()
    print(f"\nStarting server on http://0.0.0.0:{args.port}")
    print(f"Data:        {DATA_DIR}")
    print(f"Annotations: {ANNOTATION_DIR}")
    print(f"HF Repo:     {HF_REPO_ID or 'none (local mode)'}")
    app.run(host='0.0.0.0', port=args.port, debug=False)