"""
热成像威胁标注工具 — 手动标记训练数据

操作方式:
  鼠标左键拖拽  → 绘制矩形框
  P → 标记为 Person (人员)
  B → 标记为 Barrel Overheat (枪管过热)
  W → 标记为 Cold Weapon (冷武器)
  F → 标记为 Flash Residue (闪光残留)
  H → 标记为 Hotspot (热点)
  D → 删除最后一个标注
  C → 清除所有标注
  S → 保存并下一张
  N → 不保存，跳到下一张
  Q → 退出
  左右箭头 → 切换图片

用法:
  python label_tool.py                          # 标注 test/test_images/ 下所有图片
  python label_tool.py path/to/images/          # 标注指定目录
  python label_tool.py photo1.jpg photo2.jpg    # 标注指定图片
  python label_tool.py --review                 # 只查看已有标注
"""

import cv2
import numpy as np
import json
import os
import sys
import glob

THREAT_TYPES = {
    'P': ('Person',          (0, 255, 0)),
    'B': ('Barrel_Overheat', (0, 165, 255)),
    'W': ('Cold_Weapon',     (255, 255, 0)),
    'F': ('Flash_Residue',   (0, 255, 255)),
    'H': ('Hotspot',         (0, 0, 255)),
}

LABEL_FILE = 'labels.json'


def load_labels(label_path):
    if os.path.exists(label_path):
        with open(label_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    return {}


def save_labels(label_path, data):
    with open(label_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def draw_labels(image, labels):
    vis = image.copy()
    for lb in labels:
        x, y, w, h = lb['bbox']
        color = (0, 0, 0)
        for key, (name, c) in THREAT_TYPES.items():
            if name == lb['type']:
                color = c
                break
        cv2.rectangle(vis, (x, y), (x + w, y + h), color, 2)
        label_text = lb['type']
        cv2.putText(vis, label_text, (x, y - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
    return vis


class LabelTool:
    def __init__(self, image_paths, label_path, review_only=False):
        self.image_paths = image_paths
        self.label_path = label_path
        self.review_only = review_only
        self.all_labels = load_labels(label_path)
        self.idx = 0

        self.current_labels = []
        self.drawing = False
        self.rect_start = None
        self.rect_end = None
        self.display = None
        self.original = None
        self._pending_rect = None

        self.win = 'Label Tool'
        cv2.namedWindow(self.win, cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.win, self._mouse_cb)

    def _load_current(self):
        path = self.image_paths[self.idx]
        buf = np.fromfile(path, dtype=np.uint8)
        img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
        if img is None:
            img = cv2.imdecode(buf, cv2.IMREAD_GRAYSCALE)
            if img is not None:
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        if img is None:
            print(f"无法读取: {path}")
            return False

        self.original = img.copy()
        fname = os.path.basename(path)
        self.current_labels = self.all_labels.get(fname, [])
        self._redraw()
        return True

    def _redraw(self):
        self.display = draw_labels(self.original, self.current_labels)
        path = self.image_paths[self.idx]
        fname = os.path.basename(path)
        info = f"[{self.idx+1}/{len(self.image_paths)}] {fname}  |  {len(self.current_labels)} labels"
        cv2.putText(self.display, info, (10, 25),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
        help_y = self.display.shape[0] - 10
        help_text = "P:Person B:Barrel W:Weapon F:Flash H:Hotspot D:Del C:Clear S:Save N:Skip Q:Quit"
        cv2.putText(self.display, help_text, (10, help_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, (200, 200, 200), 1)
        cv2.imshow(self.win, self.display)

    def _mouse_cb(self, event, x, y, flags, param):
        if self.review_only:
            return
        if event == cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
            self.rect_start = (x, y)
            self.rect_end = (x, y)
        elif event == cv2.EVENT_MOUSEMOVE and self.drawing:
            self.rect_end = (x, y)
            vis = self.display.copy()
            cv2.rectangle(vis, self.rect_start, self.rect_end, (255, 255, 255), 1)
            cv2.imshow(self.win, vis)
        elif event == cv2.EVENT_LBUTTONUP:
            self.drawing = False
            self.rect_end = (x, y)
            x1, y1 = self.rect_start
            x2, y2 = self.rect_end
            rx, ry = min(x1, x2), min(y1, y2)
            rw, rh = abs(x2 - x1), abs(y2 - y1)
            if rw < 5 or rh < 5:
                return
            self._pending_rect = (rx, ry, rw, rh)
            print(f"  框选: ({rx},{ry}) {rw}x{rh}  -> 按 P/B/W/F/H 选类型")

    def _add_label(self, type_name):
        if not self._pending_rect:
            print("  请先框选区域")
            return
        x, y, w, h = self._pending_rect
        self.current_labels.append({
            'type': type_name,
            'bbox': [x, y, w, h],
        })
        print(f"  + {type_name} at ({x},{y}) {w}x{h}")
        self._pending_rect = None
        self._redraw()

    def _save_current(self):
        path = self.image_paths[self.idx]
        fname = os.path.basename(path)
        self.all_labels[fname] = self.current_labels
        save_labels(self.label_path, self.all_labels)
        print(f"  已保存 {fname}: {len(self.current_labels)} labels")

    def _next(self):
        if self.idx + 1 < len(self.image_paths):
            self.idx += 1
            self._load_current()
        else:
            print("已是最后一张")

    def _prev(self):
        if self.idx > 0:
            self.idx -= 1
            self._load_current()
        else:
            print("已是第一张")

    def run(self):
        if not self.image_paths:
            print("没有可标注的图片")
            return
        if not self._load_current():
            return

        print("\n=== 热成像标注工具 ===")
        print("鼠标拖拽框选 -> 按键选类型 (P/B/W/F/H)")
        print("D=删除最后 C=清空 S=保存 N=跳过 Q=退出")
        print("左右箭头切换图片\n")

        while True:
            key = cv2.waitKey(50) & 0xFF

            if key == ord('q') or key == 27:
                break
            elif key == ord('p'):
                self._add_label('Person')
            elif key == ord('b'):
                self._add_label('Barrel_Overheat')
            elif key == ord('w'):
                self._add_label('Cold_Weapon')
            elif key == ord('f'):
                self._add_label('Flash_Residue')
            elif key == ord('h'):
                self._add_label('Hotspot')
            elif key == ord('d') and not self.review_only:
                if self.current_labels:
                    removed = self.current_labels.pop()
                    print(f"  - 删除: {removed['type']}")
                    self._redraw()
            elif key == ord('c') and not self.review_only:
                self.current_labels = []
                print("  已清除所有标注")
                self._redraw()
            elif key == ord('s'):
                self._save_current()
                self._next()
            elif key == ord('n'):
                self._next()
            elif key == 83 or key == ord('.'):
                self._next()
            elif key == 81 or key == ord(','):
                self._prev()

        path = self.image_paths[self.idx]
        fname = os.path.basename(path)
        if self.current_labels:
            self.all_labels[fname] = self.current_labels
            save_labels(self.label_path, self.all_labels)
            print(f"\n自动保存 {fname}")

        cv2.destroyAllWindows()

        total = sum(len(v) for v in self.all_labels.values())
        print(f"\n标注完成: {len(self.all_labels)} 张图片, {total} 个标注")
        print(f"保存至: {self.label_path}")


def collect_images(source):
    images = []
    if isinstance(source, str) and os.path.isdir(source):
        for ext in ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tif']:
            images.extend(glob.glob(os.path.join(source, ext)))
    elif isinstance(source, list):
        for s in source:
            if os.path.isfile(s):
                images.append(s)
            elif os.path.isdir(s):
                images.extend(collect_images(s))
    images.sort()
    return images


def main():
    review = '--review' in sys.argv
    args = [a for a in sys.argv[1:] if not a.startswith('--')]

    if args:
        if all(os.path.isfile(a) for a in args):
            image_dir = None
            image_paths = args
        else:
            image_dir = args[0]
            image_paths = collect_images(image_dir)
    else:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        image_dir = os.path.join(script_dir, 'test', 'test_images')
        image_paths = collect_images(image_dir)

    if not image_paths:
        print("没有找到可标注的图片")
        print("用法: python label_tool.py [图片目录 | 图片1 图片2 ...]")
        return

    if image_dir:
        label_path = os.path.join(image_dir, '..', LABEL_FILE)
    else:
        label_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), LABEL_FILE)
    label_path = os.path.abspath(label_path)

    print(f"图片: {len(image_paths)} 张")
    print(f"标注文件: {label_path}")

    tool = LabelTool(image_paths, label_path, review_only=review)
    tool.run()


if __name__ == '__main__':
    main()
