"""
Control Box 라벨링 서버 v2 — 전체 이미지 + bbox 오버레이 방식
사용법: python tools/labeling_server.py \
          --json "data/역사이미지/slope/DJI_20260306113838_0004_everything.json"
          [--reset]  DB 초기화
브라우저: http://localhost:7001
"""
import argparse, json, sqlite3, sys
from collections import defaultdict
from pathlib import Path

import cv2, numpy as np
from fastapi import FastAPI
from fastapi.responses import HTMLResponse, JSONResponse, Response
import uvicorn

ROOT       = Path(__file__).parent.parent
DB         = ROOT / "labels" / "labeling.db"
MIN_VOTES  = 3
TRUE_RATIO = 0.6
MAX_DIM    = 3000  # served image max dimension (px)

CONTROL_LABELS = {
    "small dark object on ballast", "manhole cover", "trackside enclosure",
    "dark rectangle on ground", "small square box", "small metal cabinet",
    "small cube on gravel", "compact trackside junction box",
    "small square gray metal box beside rail",
    "small near-square electrical enclosure on the ground",
}

app = FastAPI()
_img_data: bytes = None
_orig_w = _orig_h = _disp_w = _disp_h = 0


# ── DB ────────────────────────────────────────────────────────────────────────
def get_db():
    con = sqlite3.connect(str(DB))
    con.row_factory = sqlite3.Row
    return con


def init_db(candidates: list):
    DB.parent.mkdir(parents=True, exist_ok=True)
    con = get_db()
    con.executescript("""
        CREATE TABLE IF NOT EXISTS candidates (
            id         INTEGER PRIMARY KEY,
            json_idx   INTEGER,
            label      TEXT,
            score      REAL,
            bbox       TEXT,
            image_path TEXT
        );
        CREATE TABLE IF NOT EXISTS votes (
            id           INTEGER PRIMARY KEY AUTOINCREMENT,
            candidate_id INTEGER,
            user         TEXT,
            vote         INTEGER,
            ts           DATETIME DEFAULT CURRENT_TIMESTAMP,
            UNIQUE(candidate_id, user)
        );
    """)
    if con.execute("SELECT COUNT(*) FROM candidates").fetchone()[0] > 0:
        n = con.execute("SELECT COUNT(*) FROM candidates").fetchone()[0]
        print(f"DB 기존 데이터 유지. candidates={n}")
        con.close()
        return
    for c in candidates:
        con.execute(
            "INSERT INTO candidates(json_idx,label,score,bbox,image_path) VALUES(?,?,?,?,?)",
            (c["json_idx"], c["label"], c["score"], json.dumps(c["bbox"]), c["image_path"]),
        )
    con.commit()
    print(f"DB 등록: {len(candidates)}개")
    con.close()


# ── 후보 로드 ─────────────────────────────────────────────────────────────────
def load_candidates(json_path: Path) -> list:
    data = json.loads(json_path.read_text(encoding="utf-8"))
    stem = json_path.stem.replace("_everything", "")
    img_path = next(
        (json_path.parent / f"{stem}{ext}" for ext in (".JPG", ".jpg", ".png")
         if (json_path.parent / f"{stem}{ext}").exists()),
        None,
    )
    if img_path is None:
        raise FileNotFoundError(f"원본 이미지 없음: {json_path.parent / stem}.*")

    candidates = []
    for idx, seg in enumerate(data.get("segments", [])):
        label = seg.get("label", "").strip()
        if label not in CONTROL_LABELS:
            continue
        x0, y0, x1, y1 = [float(v) for v in seg["bbox"]]
        if (x1 - x0) < 4 or (y1 - y0) < 4:
            continue
        candidates.append({
            "json_idx":   idx,
            "label":      label,
            "score":      float(seg.get("score", 0)),
            "bbox":       [x0, y0, x1, y1],
            "image_path": str(img_path),
        })

    segs_total = len(data.get("segments", []))
    print(f"필터링: {segs_total}개 → {len(candidates)}개 control_box 후보")
    return candidates


# ── 이미지 준비 (서버 시작 시 1회) ───────────────────────────────────────────
def prepare_image(img_path: Path):
    global _img_data, _orig_w, _orig_h, _disp_w, _disp_h
    buf = np.fromfile(str(img_path), dtype=np.uint8)
    img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"이미지 로드 실패: {img_path}")
    _orig_h, _orig_w = img.shape[:2]
    scale = min(1.0, MAX_DIM / max(_orig_w, _orig_h))
    _disp_w, _disp_h = int(_orig_w * scale), int(_orig_h * scale)
    if scale < 1.0:
        img = cv2.resize(img, (_disp_w, _disp_h), interpolation=cv2.INTER_LANCZOS4)
    _, enc = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 92])
    _img_data = bytes(enc)
    print(f"이미지 준비: {_orig_w}×{_orig_h} → {_disp_w}×{_disp_h}")


# ── API ───────────────────────────────────────────────────────────────────────
@app.get("/image")
async def serve_image():
    return Response(content=_img_data, media_type="image/jpeg")


@app.get("/api/image_info")
async def api_image_info():
    return {"orig_w": _orig_w, "orig_h": _orig_h, "disp_w": _disp_w, "disp_h": _disp_h}


@app.get("/api/candidates")
async def api_candidates(user: str = ""):
    con = get_db()
    rows = con.execute("""
        SELECT c.id, c.bbox, c.label, c.score, v.vote
        FROM candidates c
        LEFT JOIN votes v ON c.id=v.candidate_id AND v.user=?
    """, (user,)).fetchall()
    con.close()
    return {"candidates": [
        {
            "id":     r["id"],
            "bbox":   json.loads(r["bbox"]),
            "label":  r["label"],
            "score":  round(r["score"], 2),
            "voted":  r["vote"] is not None,
            "is_true": r["vote"] == 1 if r["vote"] is not None else None,
        } for r in rows
    ]}


@app.post("/api/vote")
async def api_vote(data: dict):
    user     = data.get("user", "").strip()
    all_ids  = data.get("all_ids", [])
    true_ids = set(data.get("true_ids", []))
    if not user or not all_ids:
        return {"ok": False, "error": "파라미터 오류"}
    con = get_db()
    for cid in all_ids:
        con.execute(
            "INSERT OR REPLACE INTO votes(candidate_id,user,vote) VALUES(?,?,?)",
            (cid, user, 1 if cid in true_ids else 0),
        )
    con.commit()
    con.close()
    return {"ok": True}


@app.get("/api/stats")
async def api_stats():
    con = get_db()
    total  = con.execute("SELECT COUNT(*) FROM candidates").fetchone()[0]
    voted3 = con.execute(f"""
        SELECT COUNT(*) FROM (
            SELECT candidate_id FROM votes GROUP BY candidate_id HAVING COUNT(*) >= {MIN_VOTES}
        )
    """).fetchone()[0]
    users  = con.execute("SELECT COUNT(DISTINCT user) FROM votes").fetchone()[0]
    tvotes = con.execute("SELECT COUNT(*) FROM votes").fetchone()[0]
    con.close()
    return {"total": total, "voted3plus": voted3, "users": users, "total_votes": tvotes}


@app.post("/api/export")
async def api_export():
    con = get_db()
    confirmed = con.execute(f"""
        SELECT c.id, c.bbox, c.image_path,
               SUM(v.vote) AS true_v, COUNT(v.id) AS total_v
        FROM candidates c JOIN votes v ON c.id=v.candidate_id
        GROUP BY c.id
        HAVING total_v >= {MIN_VOTES} AND (CAST(true_v AS REAL)/total_v) >= {TRUE_RATIO}
    """).fetchall()
    con.close()

    by_image = defaultdict(list)
    for row in confirmed:
        by_image[row["image_path"]].append(row)

    out_dir = ROOT / "labels" / "yolo_export"
    out_dir.mkdir(parents=True, exist_ok=True)

    count = 0
    for img_path, rows in by_image.items():
        buf = np.fromfile(img_path, dtype=np.uint8)
        img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
        H, W = img.shape[:2]
        txt = out_dir / (Path(img_path).stem + ".txt")
        with open(txt, "w") as f:
            for row in rows:
                x0, y0, x1, y1 = json.loads(row["bbox"])
                f.write(f"0 {((x0+x1)/2)/W:.6f} {((y0+y1)/2)/H:.6f} "
                        f"{(x1-x0)/W:.6f} {(y1-y0)/H:.6f}\n")
                count += 1
    return {"ok": True, "labels": count, "dir": str(out_dir)}


# ── HTML ──────────────────────────────────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
async def index():
    return HTML


HTML = r"""<!DOCTYPE html>
<html lang="ko">
<head>
<meta charset="UTF-8">
<title>Control Box 라벨링</title>
<style>
*{box-sizing:border-box;margin:0;padding:0}
body{font-family:'Segoe UI',sans-serif;background:#1a1a2e;color:#e0e0e0;display:flex;flex-direction:column;height:100vh;overflow:hidden}
#toolbar{background:#16213e;border-bottom:1px solid #0f3460;padding:8px 14px;display:flex;align-items:center;gap:10px;flex-shrink:0;flex-wrap:wrap}
h1{color:#e94560;font-size:1rem;white-space:nowrap}
input[type=text]{background:#0f3460;border:1px solid #2a4a7a;color:#e0e0e0;padding:5px 10px;border-radius:4px;font-size:.85rem;width:130px}
input[type=text]:focus{outline:none;border-color:#e94560}
.btn{background:#e94560;color:#fff;border:none;padding:6px 12px;border-radius:5px;cursor:pointer;font-size:.82rem;font-weight:700;white-space:nowrap}
.btn:hover{background:#c73652}
.btn-sec{background:#0f3460;border:1px solid #2a4a7a;color:#e0e0e0}
.btn-sec:hover{background:#1a5299}
#stats{font-size:.75rem;color:#777;white-space:nowrap}
#legend{display:flex;gap:10px;align-items:center;font-size:.72rem;color:#888}
.leg{display:inline-flex;align-items:center;gap:3px}
.lb{width:12px;height:12px;border:2px solid;border-radius:2px;flex-shrink:0}
#wrap{flex:1;position:relative;overflow:hidden;cursor:crosshair}
canvas{display:block}
#tip{position:absolute;background:rgba(0,0,0,.85);color:#fff;font-size:.7rem;padding:3px 8px;border-radius:3px;pointer-events:none;display:none;max-width:300px}
#hint{position:absolute;bottom:8px;left:50%;transform:translateX(-50%);font-size:.72rem;color:#444;pointer-events:none;white-space:nowrap}
</style>
</head>
<body>
<div id="toolbar">
  <h1>🎯 Control Box</h1>
  <input type="text" id="uname" placeholder="이름/사번" autocomplete="off">
  <button class="btn" onclick="init()">시작</button>
  <button class="btn btn-sec" onclick="fitView()">맞춤(F)</button>
  <span id="stats">—</span>
  <div id="legend">
    <span class="leg"><span class="lb" style="border-color:#ffaa00"></span>미투표</span>
    <span class="leg"><span class="lb" style="border-color:#00cc66"></span>컨트롤박스 ✓</span>
    <span class="leg"><span class="lb" style="border-color:#cc4444"></span>아님 ✗</span>
    <span class="leg"><span class="lb" style="border-color:#444"></span>타인투표완료</span>
  </div>
  <div style="flex:1"></div>
  <button class="btn btn-sec" onclick="doExport()" style="font-size:.75rem">YOLO 내보내기</button>
</div>
<div id="wrap">
  <canvas id="cv"></canvas>
  <div id="tip"></div>
  <div id="hint">휠=줌 | 드래그=이동 | 클릭=YES/NO 토글 | F=맞춤</div>
</div>

<script>
const cv = document.getElementById('cv');
const ctx = cv.getContext('2d');
const wrap = document.getElementById('wrap');

let user = '', cands = [], imgScale = 1;
let origW = 0, origH = 0, dispW = 0, dispH = 0;
let sc = 1, ox = 0, oy = 0;
let isDrag = false, dragX = 0, dragY = 0, moved = false;
const img = new Image();

function resizeCv() {
    cv.width = wrap.clientWidth;
    cv.height = wrap.clientHeight;
    render();
}
window.addEventListener('resize', resizeCv);
resizeCv();

async function init() {
    user = document.getElementById('uname').value.trim();
    if (!user) { alert('이름 입력 필요'); return; }
    const info = await fetch('/api/image_info').then(r => r.json());
    origW = info.orig_w; origH = info.orig_h;
    dispW = info.disp_w; dispH = info.disp_h;
    imgScale = dispW / origW;
    img.onload = () => { fitView(); loadCands(); };
    img.src = '/image?' + Date.now();
}

async function loadCands() {
    const d = await fetch('/api/candidates?user=' + encodeURIComponent(user)).then(r => r.json());
    cands = d.candidates.map(c => ({...c, _sel: c.voted ? c.is_true : undefined}));
    updateStats();
    render();
}

function updateStats() {
    const myVotes = cands.filter(c => c._sel !== undefined).length;
    const yes = cands.filter(c => c._sel === true).length;
    const othersVoted = cands.filter(c => c.voted && c._sel === undefined).length;
    document.getElementById('stats').textContent =
        `후보 ${cands.length}개 | 내투표 ${myVotes}개 (YES:${yes}) | 타인완료 ${othersVoted}개`;
}

function fitView() {
    if (!img.naturalWidth) return;
    const ww = wrap.clientWidth, wh = wrap.clientHeight;
    sc = Math.min(ww / dispW, wh / dispH) * 0.97;
    ox = (ww - dispW * sc) / 2;
    oy = (wh - dispH * sc) / 2;
    render();
}

function render() {
    const W = cv.width, H = cv.height;
    ctx.fillStyle = '#0d0d1e';
    ctx.fillRect(0, 0, W, H);
    if (!img.naturalWidth) return;

    ctx.save();
    ctx.translate(ox, oy);
    ctx.scale(sc, sc);
    ctx.drawImage(img, 0, 0, dispW, dispH);

    const lw = Math.max(1, 2 / sc);
    for (const c of cands) {
        const [x0, y0, x1, y1] = c.bbox.map(v => v * imgScale);
        const w = x1 - x0, h = y1 - y0;
        let color, fill = null;
        if (c._sel === true)  { color = '#00cc66'; fill = 'rgba(0,204,102,0.15)'; }
        else if (c._sel === false) { color = '#cc4444'; fill = 'rgba(204,68,68,0.15)'; }
        else if (c.voted) { color = '#444444'; }
        else { color = c.score > 0.5 ? '#ffaa00' : '#ff7700aa'; }

        ctx.lineWidth = lw;
        ctx.strokeStyle = color;
        if (fill) {
            ctx.fillStyle = fill;
            ctx.fillRect(x0, y0, w, h);
        }
        ctx.strokeRect(x0, y0, w, h);
    }
    ctx.restore();
}

// ── wheel zoom ────────────────────────────────────────────────────────────────
wrap.addEventListener('wheel', e => {
    e.preventDefault();
    const rect = wrap.getBoundingClientRect();
    const mx = e.clientX - rect.left, my = e.clientY - rect.top;
    const f = e.deltaY < 0 ? 1.15 : 1/1.15;
    ox = mx - (mx - ox) * f;
    oy = my - (my - oy) * f;
    sc *= f;
    render();
}, {passive: false});

// ── drag pan ──────────────────────────────────────────────────────────────────
wrap.addEventListener('mousedown', e => {
    if (e.button !== 0) return;
    isDrag = true; moved = false;
    dragX = e.clientX; dragY = e.clientY;
    wrap.style.cursor = 'grabbing';
});
window.addEventListener('mousemove', e => {
    if (isDrag) {
        const dx = e.clientX - dragX, dy = e.clientY - dragY;
        if (Math.abs(dx) + Math.abs(dy) > 3) moved = true;
        if (moved) { ox += dx; oy += dy; dragX = e.clientX; dragY = e.clientY; render(); }
    }
    // tooltip
    const c = hitTest(e);
    const tip = document.getElementById('tip');
    if (c) {
        const rect = wrap.getBoundingClientRect();
        tip.style.display = 'block';
        tip.style.left = (e.clientX - rect.left + 14) + 'px';
        tip.style.top  = (e.clientY - rect.top  +  4) + 'px';
        const status = c._sel === true ? '✅ YES' : c._sel === false ? '❌ NO' : c.voted ? '⬜ 타인완료' : '❓ 미투표';
        tip.textContent = `[${c.id}] ${c.label} (${c.score}) — ${status}`;
    } else {
        tip.style.display = 'none';
    }
});
window.addEventListener('mouseup', e => {
    if (!isDrag) return;
    wrap.style.cursor = 'crosshair';
    isDrag = false;
    if (!moved) handleClick(e);
});

// ── hit test ──────────────────────────────────────────────────────────────────
function imgCoords(e) {
    const rect = wrap.getBoundingClientRect();
    return [(e.clientX - rect.left - ox) / sc / imgScale,
            (e.clientY - rect.top  - oy) / sc / imgScale];
}

function hitTest(e) {
    const [ix, iy] = imgCoords(e);
    let best = null, bestArea = Infinity;
    for (const c of cands) {
        const [x0, y0, x1, y1] = c.bbox;
        if (ix >= x0 && ix <= x1 && iy >= y0 && iy <= y1) {
            const a = (x1-x0)*(y1-y0);
            if (a < bestArea) { bestArea = a; best = c; }
        }
    }
    return best;
}

async function handleClick(e) {
    if (!user) return;
    const c = hitTest(e);
    if (!c) return;
    // Toggle: undefined→YES→NO→YES...
    c._sel = c._sel === true ? false : true;
    render();
    updateStats();
    // Save immediately
    await fetch('/api/vote', {
        method: 'POST',
        headers: {'Content-Type': 'application/json'},
        body: JSON.stringify({user, all_ids: [c.id], true_ids: c._sel ? [c.id] : []})
    });
    c.voted = true;
    c.is_true = c._sel;
}

// ── keyboard ──────────────────────────────────────────────────────────────────
window.addEventListener('keydown', e => {
    if (e.key === 'f' || e.key === 'F') fitView();
});

async function doExport() {
    const r = await fetch('/api/export', {method: 'POST'}).then(r => r.json());
    alert(`완료: ${r.labels}개 라벨\n경로: ${r.dir}`);
}
</script>
</body>
</html>"""


# ── 메인 ─────────────────────────────────────────────────────────────────────
def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--json",  required=True, help="*_everything.json 경로")
    ap.add_argument("--port",  type=int, default=7001)
    ap.add_argument("--reset", action="store_true", help="DB 초기화 후 재로드")
    args = ap.parse_args()

    json_path = Path(args.json)
    if not json_path.exists():
        print(f"JSON 없음: {json_path}"); sys.exit(1)

    if args.reset and DB.exists():
        DB.unlink()
        print("DB 초기화 완료.")

    print("후보 로딩 중...")
    candidates = load_candidates(json_path)
    init_db(candidates)

    # Prepare image
    stem = json_path.stem.replace("_everything", "")
    img_path = next(
        (json_path.parent / f"{stem}{ext}" for ext in (".JPG", ".jpg", ".png")
         if (json_path.parent / f"{stem}{ext}").exists()),
        None,
    )
    if img_path is None:
        print(f"원본 이미지 없음: {json_path.parent / stem}.*"); sys.exit(1)
    prepare_image(img_path)

    print(f"\n라벨링 서버 시작: http://localhost:{args.port}")
    uvicorn.run(app, host="0.0.0.0", port=args.port, log_level="warning")


if __name__ == "__main__":
    main()