"""everything.json을 레이블별로 분리해 bbox를 원본 이미지에 표시.

사용:
    python tools/render_everything_by_label.py \
        --image data/역사이미지/slope/DJI_20260306113838_0004_tophalf.jpg \
        --json  output/raamen/DJI_20260306113838_0004_tophalf_everything.json \
        --output output/everything_by_label \
        [--scale 0.3] \
        [--label small_dark_object_on_ballast ...] \
        [--rail-offset 200]   # railroad_track/railway_rail bbox + N픽셀 안쪽 필터
"""
import argparse
import json
import re
import cv2
import numpy as np
from pathlib import Path

RAIL_LABELS = {"railroad track", "railway rail"}


def safe_name(label: str) -> str:
    return re.sub(r"[^\w]", "_", label)


def build_rail_zone(segments: list, offset: int, img_h: int, img_w: int):
    """railroad_track/railway_rail 세그먼트 합산 bbox + offset → (x1,y1,x2,y2)."""
    xs1, ys1, xs2, ys2 = [], [], [], []
    for seg in segments:
        if seg.get("label", "").strip().lower() in RAIL_LABELS:
            bx1, by1, bx2, by2 = seg["bbox"]
            xs1.append(bx1); ys1.append(by1)
            xs2.append(bx2); ys2.append(by2)
    if not xs1:
        return None
    return (
        max(0,     int(min(xs1)) - offset),
        max(0,     int(min(ys1)) - offset),
        min(img_w, int(max(xs2)) + offset),
        min(img_h, int(max(ys2)) + offset),
    )


def in_zone(seg_bbox, zone):
    """세그먼트 bbox 중심이 zone 안에 있으면 True."""
    bx1, by1, bx2, by2 = seg_bbox
    cx, cy = (bx1 + bx2) / 2, (by1 + by2) / 2
    zx1, zy1, zx2, zy2 = zone
    return zx1 <= cx <= zx2 and zy1 <= cy <= zy2


def render_label(img_full, segments, label: str, output_path: Path, scale: float,
                 zone=None):
    h, w = img_full.shape[:2]
    canvas = cv2.resize(img_full, (int(w * scale), int(h * scale)),
                        interpolation=cv2.INTER_AREA)

    color     = (0, 80, 255)
    font      = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = max(0.4, scale * 2.0)
    thickness  = max(1, int(scale * 5))

    if zone:
        zx1, zy1, zx2, zy2 = [int(v * scale) for v in zone]
        cv2.rectangle(canvas, (zx1, zy1), (zx2, zy2), (0, 255, 0), 2)

    for seg in segments:
        if seg.get("points"):
            pts = np.array([[int(x * scale), int(y * scale)] for x, y in seg["points"]], dtype=np.int32)
            overlay = canvas.copy()
            cv2.fillPoly(overlay, [pts], color)
            cv2.addWeighted(overlay, 0.25, canvas, 0.75, 0, canvas)
            cv2.polylines(canvas, [pts], True, color, thickness)
        else:
            x1, y1, x2, y2 = [int(v * scale) for v in seg["bbox"]]
            cv2.rectangle(canvas, (x1, y1), (x2, y2), color, thickness)

    header = f"{label}  [{len(segments)}]"
    cv2.putText(canvas, header, (10, 40), font, font_scale * 1.2,
                (0, 220, 0), max(1, int(scale * 3)) + 1)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    ret, buf = cv2.imencode(".jpg", canvas, [cv2.IMWRITE_JPEG_QUALITY, 88])
    if not ret:
        raise RuntimeError("JPEG 인코딩 실패")
    output_path.write_bytes(buf.tobytes())
    print(f"  {label:40s} {len(segments):5d}개  → {output_path.name}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image",       required=True, type=Path)
    ap.add_argument("--json",        required=True, type=Path)
    ap.add_argument("--output",      required=True, type=Path)
    ap.add_argument("--scale",       type=float, default=0.25)
    ap.add_argument("--label",       nargs="*", default=None)
    ap.add_argument("--rail-offset", type=int, default=0,
                    help="railroad_track/railway_rail 합산 bbox + N픽셀 zone 필터 (0=비활성)")
    ap.add_argument("--zone", type=int, nargs=4, metavar=("X1","Y1","X2","Y2"), default=None,
                    help="수동 zone 지정 (--rail-offset보다 우선)")
    args = ap.parse_args()

    buf = np.fromfile(str(args.image), dtype=np.uint8)
    img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(f"이미지 읽기 실패: {args.image}")
    img_h, img_w = img.shape[:2]

    with open(args.json, encoding="utf-8") as f:
        data = json.load(f)

    groups: dict[str, list] = {}
    for seg in data["segments"]:
        groups.setdefault(seg["label"], []).append(seg)

    zone = None
    if args.zone:
        zone = tuple(args.zone)
        print(f"수동 zone: {zone}")
    elif args.rail_offset > 0:
        zone = build_rail_zone(data["segments"], args.rail_offset, img_h, img_w)
        if zone:
            print(f"레일 zone (offset={args.rail_offset}px): {zone}")
        else:
            print("경고: railroad_track/railway_rail 세그먼트 없음 → zone 필터 미적용")

    target_labels = args.label if args.label else sorted(groups)

    print(f"원본: {img_w}x{img_h}  scale={args.scale}")
    print(f"레이블 {len(target_labels)}개 렌더링 → {args.output}/")

    for lbl in target_labels:
        if lbl not in groups:
            print(f"  [skip] {lbl} - no segments")
            continue
        segs = sorted(groups[lbl], key=lambda s: s["score"], reverse=True)
        if zone:
            segs = [s for s in segs if in_zone(s["bbox"], zone)]
        fname = f"{safe_name(lbl)}.jpg"
        render_label(img, segs, lbl, args.output / fname, args.scale, zone=zone)

    print("완료.")


if __name__ == "__main__":
    main()