"""
Manual annotation tool for camera calibration aids.

Use this when automatic corner detection is not practical. It lets you freeze a
camera/video frame or open an image, then manually mark corners, lines,
rectangles, and ROI boxes. Saved JSON coordinates are in original image pixels.
"""

import argparse
import json
import os
import sys
from datetime import datetime

os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"

import cv2
import numpy as np


def _fix_qt_plugin_path():
    """OpenCV can redirect Qt to cv2/qt/plugins; prefer PyQt5's plugins."""
    try:
        from PyQt5.QtCore import QLibraryInfo

        os.environ["QT_QPA_PLATFORM_PLUGIN_PATH"] = QLibraryInfo.location(
            QLibraryInfo.PluginsPath
        )
    except Exception:
        path = os.environ.get("QT_QPA_PLATFORM_PLUGIN_PATH", "")
        if "cv2" in path.replace("\\", "/"):
            os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH", None)


_fix_qt_plugin_path()

from PyQt5.QtCore import QPoint, QRect, Qt, QTimer
from PyQt5.QtGui import QColor, QFont, QImage, QPainter, QPen, QPixmap
from PyQt5.QtWidgets import (
    QAction,
    QApplication,
    QButtonGroup,
    QFileDialog,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QMainWindow,
    QMessageBox,
    QPushButton,
    QSizePolicy,
    QToolButton,
    QVBoxLayout,
    QWidget,
)


COLORS = {
    "corner": QColor(0, 229, 255),
    "line": QColor(255, 213, 79),
    "rect": QColor(129, 199, 132),
    "roi": QColor(255, 112, 67),
}

MODE_TEXT = {
    "corner": "角点",
    "line": "直线",
    "rect": "矩形",
    "roi": "ROI",
}


def cv_to_qpixmap(frame_bgr):
    rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
    h, w, ch = rgb.shape
    qimg = QImage(rgb.data, w, h, ch * w, QImage.Format_RGB888).copy()
    return QPixmap.fromImage(qimg)


def normalize_rect(p0, p1):
    x0, y0 = p0
    x1, y1 = p1
    left, right = sorted((int(round(x0)), int(round(x1))))
    top, bottom = sorted((int(round(y0)), int(round(y1))))
    return [left, top, max(0, right - left), max(0, bottom - top)]


class ImageCanvas(QLabel):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setMouseTracking(True)
        self.setFocusPolicy(Qt.StrongFocus)
        self.setAlignment(Qt.AlignCenter)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.setMinimumSize(640, 480)

        self.frame = None
        self.pixmap_src = None
        self.annotations = []
        self.mode = "corner"
        self.drag_start = None
        self.drag_current = None
        self.hover_img_pos = None
        self.status_callback = None

    def set_status_callback(self, callback):
        self.status_callback = callback

    def set_frame(self, frame_bgr):
        self.frame = frame_bgr.copy()
        self.pixmap_src = cv_to_qpixmap(self.frame)
        self.update()

    def clear_annotations(self):
        self.annotations.clear()
        self.drag_start = None
        self.drag_current = None
        self.update()

    def undo(self):
        if self.annotations:
            self.annotations.pop()
            self.update()

    def set_mode(self, mode):
        self.mode = mode
        self.drag_start = None
        self.drag_current = None
        self.update()

    def image_size(self):
        if self.frame is None:
            return None
        h, w = self.frame.shape[:2]
        return [w, h]

    def _target_rect(self):
        if self.pixmap_src is None:
            return QRect()
        src_w = self.pixmap_src.width()
        src_h = self.pixmap_src.height()
        if src_w <= 0 or src_h <= 0:
            return QRect()

        widget_w = self.width()
        widget_h = self.height()
        scale = min(widget_w / src_w, widget_h / src_h)
        draw_w = int(src_w * scale)
        draw_h = int(src_h * scale)
        x = (widget_w - draw_w) // 2
        y = (widget_h - draw_h) // 2
        return QRect(x, y, draw_w, draw_h)

    def _widget_to_image(self, pos):
        if self.pixmap_src is None:
            return None
        rect = self._target_rect()
        if not rect.contains(pos):
            return None
        x = (pos.x() - rect.x()) * self.pixmap_src.width() / rect.width()
        y = (pos.y() - rect.y()) * self.pixmap_src.height() / rect.height()
        x = min(max(x, 0), self.pixmap_src.width() - 1)
        y = min(max(y, 0), self.pixmap_src.height() - 1)
        return [float(x), float(y)]

    def _image_to_widget(self, point):
        rect = self._target_rect()
        x = rect.x() + point[0] * rect.width() / self.pixmap_src.width()
        y = rect.y() + point[1] * rect.height() / self.pixmap_src.height()
        return QPoint(int(round(x)), int(round(y)))

    def mousePressEvent(self, event):
        if event.button() != Qt.LeftButton or self.frame is None:
            return
        img_pos = self._widget_to_image(event.pos())
        if img_pos is None:
            return

        if self.mode == "corner":
            self.annotations.append(
                {
                    "type": "corner",
                    "label": f"corner_{self._count_type('corner') + 1}",
                    "points": [[round(img_pos[0], 2), round(img_pos[1], 2)]],
                }
            )
            self.update()
        else:
            self.drag_start = img_pos
            self.drag_current = img_pos
            self.update()

    def mouseMoveEvent(self, event):
        img_pos = self._widget_to_image(event.pos())
        self.hover_img_pos = img_pos
        if self.status_callback:
            if img_pos is None:
                self.status_callback(f"模式: {MODE_TEXT[self.mode]} | 坐标: -")
            else:
                self.status_callback(
                    f"模式: {MODE_TEXT[self.mode]} | 坐标: "
                    f"x={img_pos[0]:.1f}, y={img_pos[1]:.1f}"
                )
        if self.drag_start is not None and img_pos is not None:
            self.drag_current = img_pos
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() != Qt.LeftButton or self.drag_start is None:
            return
        img_pos = self._widget_to_image(event.pos()) or self.drag_current
        if img_pos is None:
            self.drag_start = None
            self.drag_current = None
            self.update()
            return

        p0 = [round(self.drag_start[0], 2), round(self.drag_start[1], 2)]
        p1 = [round(img_pos[0], 2), round(img_pos[1], 2)]
        if self.mode == "line":
            self.annotations.append(
                {
                    "type": "line",
                    "label": f"line_{self._count_type('line') + 1}",
                    "points": [p0, p1],
                }
            )
        elif self.mode in ("rect", "roi"):
            x, y, w, h = normalize_rect(p0, p1)
            if w > 0 and h > 0:
                self.annotations.append(
                    {
                        "type": self.mode,
                        "label": f"{self.mode}_{self._count_type(self.mode) + 1}",
                        "rect": [x, y, w, h],
                        "points": [[x, y], [x + w, y + h]],
                    }
                )
        self.drag_start = None
        self.drag_current = None
        self.update()

    def keyPressEvent(self, event):
        if event.key() == Qt.Key_Z and event.modifiers() & Qt.ControlModifier:
            self.undo()
        else:
            super().keyPressEvent(event)

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.fillRect(self.rect(), QColor("#0b0b0b"))

        if self.pixmap_src is None:
            painter.setPen(QColor("#777777"))
            painter.setFont(QFont("SimHei", 16))
            painter.drawText(self.rect(), Qt.AlignCenter, "打开图片/视频或启动摄像头")
            painter.end()
            return

        target = self._target_rect()
        painter.drawPixmap(target, self.pixmap_src)
        painter.setRenderHint(QPainter.Antialiasing, True)

        for ann in self.annotations:
            self._draw_annotation(painter, ann)

        if self.drag_start is not None and self.drag_current is not None:
            preview = self._preview_annotation()
            if preview is not None:
                self._draw_annotation(painter, preview, preview=True)

        painter.end()

    def _preview_annotation(self):
        if self.mode == "line":
            return {"type": "line", "points": [self.drag_start, self.drag_current]}
        if self.mode in ("rect", "roi"):
            x, y, w, h = normalize_rect(self.drag_start, self.drag_current)
            return {
                "type": self.mode,
                "rect": [x, y, w, h],
                "points": [[x, y], [x + w, y + h]],
            }
        return None

    def _draw_annotation(self, painter, ann, preview=False):
        color = COLORS.get(ann["type"], QColor(255, 255, 255))
        pen = QPen(color, 2 if not preview else 1, Qt.DashLine if preview else Qt.SolidLine)
        painter.setPen(pen)
        painter.setBrush(Qt.NoBrush)

        if ann["type"] == "corner":
            p = self._image_to_widget(ann["points"][0])
            painter.drawEllipse(p, 5, 5)
            painter.drawLine(p.x() - 8, p.y(), p.x() + 8, p.y())
            painter.drawLine(p.x(), p.y() - 8, p.x(), p.y() + 8)
            self._draw_label(painter, p, ann.get("label", "corner"), color)
        elif ann["type"] == "line":
            p0 = self._image_to_widget(ann["points"][0])
            p1 = self._image_to_widget(ann["points"][1])
            painter.drawLine(p0, p1)
            self._draw_label(painter, p0, ann.get("label", "line"), color)
        elif ann["type"] in ("rect", "roi"):
            x, y, w, h = ann["rect"]
            p0 = self._image_to_widget([x, y])
            p1 = self._image_to_widget([x + w, y + h])
            rect = QRect(p0, p1).normalized()
            painter.drawRect(rect)
            self._draw_label(painter, rect.topLeft(), ann.get("label", ann["type"]), color)

    def _draw_label(self, painter, pos, text, color):
        painter.setFont(QFont("Consolas", 10))
        metrics = painter.fontMetrics()
        box = metrics.boundingRect(text).adjusted(-4, -2, 4, 2)
        box.moveTopLeft(pos + QPoint(8, -box.height() - 4))
        painter.fillRect(box, QColor(0, 0, 0, 170))
        painter.setPen(color)
        painter.drawText(box, Qt.AlignCenter, text)

    def _count_type(self, ann_type):
        return sum(1 for ann in self.annotations if ann["type"] == ann_type)


class MainWindow(QMainWindow):
    def __init__(self, args):
        super().__init__()
        self.setWindowTitle("Orisys 手动标定标注工具")
        self.args = args
        self.capture = None
        self.source = None
        self.source_type = None
        self.paused = True

        self.timer = QTimer(self)
        self.timer.timeout.connect(self._read_frame)

        self._setup_ui()
        self._setup_shortcuts()

        if args.source:
            self._open_source(args.source)

    def _setup_ui(self):
        root = QWidget()
        self.setCentralWidget(root)
        layout = QVBoxLayout(root)
        layout.setContentsMargins(8, 8, 8, 8)
        layout.setSpacing(8)

        top = QHBoxLayout()
        top.setSpacing(6)

        self.source_edit = QLineEdit(str(self.args.source if self.args.source is not None else "0"))
        self.source_edit.setPlaceholderText("摄像头编号、视频路径或图片路径")
        top.addWidget(self.source_edit, 1)

        self.btn_open = QPushButton("打开")
        self.btn_open.clicked.connect(self._on_open_clicked)
        top.addWidget(self.btn_open)

        self.btn_file = QPushButton("选文件")
        self.btn_file.clicked.connect(self._choose_file)
        top.addWidget(self.btn_file)

        self.btn_pause = QPushButton("暂停/继续")
        self.btn_pause.clicked.connect(self._toggle_pause)
        top.addWidget(self.btn_pause)

        self.btn_save = QPushButton("保存 JSON")
        self.btn_save.clicked.connect(self._save_json)
        top.addWidget(self.btn_save)

        self.btn_save_image = QPushButton("保存预览图")
        self.btn_save_image.clicked.connect(self._save_preview)
        top.addWidget(self.btn_save_image)

        layout.addLayout(top)

        tools = QHBoxLayout()
        tools.setSpacing(6)

        self.mode_group = QButtonGroup(self)
        self.mode_group.setExclusive(True)
        for mode in ("corner", "line", "rect", "roi"):
            btn = QToolButton()
            btn.setText(MODE_TEXT[mode])
            btn.setCheckable(True)
            btn.clicked.connect(lambda checked, m=mode: self.canvas.set_mode(m))
            self.mode_group.addButton(btn)
            tools.addWidget(btn)
            if mode == "corner":
                btn.setChecked(True)

        tools.addStretch(1)

        self.btn_undo = QPushButton("撤销")
        self.btn_undo.clicked.connect(lambda: self.canvas.undo())
        tools.addWidget(self.btn_undo)

        self.btn_clear = QPushButton("清空")
        self.btn_clear.clicked.connect(self._clear_annotations)
        tools.addWidget(self.btn_clear)

        layout.addLayout(tools)

        self.canvas = ImageCanvas()
        self.canvas.set_status_callback(self._set_status)
        layout.addWidget(self.canvas, 1)

        self.status = QLabel("模式: 角点 | 坐标: -")
        self.status.setMinimumHeight(24)
        layout.addWidget(self.status)

        self.setStyleSheet(
            """
            QMainWindow, QWidget { background: #111111; color: #dddddd; }
            QLineEdit {
                background: #0b0b0b; color: #dddddd; border: 1px solid #333333;
                border-radius: 4px; padding: 5px 8px;
            }
            QPushButton, QToolButton {
                background: #1d1d1d; color: #dddddd; border: 1px solid #333333;
                border-radius: 4px; padding: 6px 12px;
            }
            QPushButton:hover, QToolButton:hover { background: #2a2a2a; }
            QToolButton:checked {
                background: #00323a; color: #00e5ff; border-color: #00a6b8;
            }
            QLabel { color: #aaaaaa; }
            """
        )

    def _setup_shortcuts(self):
        shortcuts = [
            ("Ctrl+S", self._save_json),
            ("Ctrl+Z", self.canvas.undo),
            ("Space", self._toggle_pause),
            ("C", lambda: self._select_mode("corner")),
            ("L", lambda: self._select_mode("line")),
            ("B", lambda: self._select_mode("rect")),
            ("R", lambda: self._select_mode("roi")),
        ]
        for key, callback in shortcuts:
            action = QAction(self)
            action.setShortcut(key)
            action.triggered.connect(callback)
            self.addAction(action)

    def _select_mode(self, mode):
        self.canvas.set_mode(mode)
        for btn in self.mode_group.buttons():
            if btn.text() == MODE_TEXT[mode]:
                btn.setChecked(True)
                break

    def _set_status(self, text):
        total = len(self.canvas.annotations)
        self.status.setText(f"{text} | 标注数: {total}")

    def _on_open_clicked(self):
        self._open_source(self.source_edit.text().strip())

    def _choose_file(self):
        path, _ = QFileDialog.getOpenFileName(
            self,
            "选择图片或视频",
            "data",
            "媒体文件 (*.png *.jpg *.jpeg *.bmp *.tif *.tiff *.mp4 *.avi *.mov);;所有文件 (*)",
        )
        if path:
            self.source_edit.setText(path)
            self._open_source(path)

    def _open_source(self, source):
        self._release_capture()
        self.canvas.clear_annotations()
        self.source = source

        src_for_cv = source
        try:
            src_for_cv = int(source)
        except (TypeError, ValueError):
            pass

        if isinstance(src_for_cv, str) and self._is_image_file(src_for_cv):
            frame = cv2.imread(src_for_cv, cv2.IMREAD_COLOR)
            if frame is None:
                self._show_error(f"无法读取图片: {src_for_cv}")
                return
            self.source_type = "image"
            self.paused = True
            self.canvas.set_frame(frame)
            self._set_status("图片已打开")
            return

        cap = cv2.VideoCapture(src_for_cv)
        if not cap.isOpened():
            self._show_error(f"无法打开输入源: {source}")
            return
        self.capture = cap
        self.source_type = "camera_or_video"
        self.paused = False
        self.timer.start(30)
        self._read_frame()

    def _read_frame(self):
        if self.capture is None or self.paused:
            return
        ok, frame = self.capture.read()
        if not ok:
            self.paused = True
            self.timer.stop()
            self._set_status("视频结束或读取失败")
            return
        self.canvas.set_frame(frame)

    def _toggle_pause(self):
        if self.capture is None:
            return
        self.paused = not self.paused
        if self.paused:
            self.timer.stop()
        else:
            self.timer.start(30)

    def _clear_annotations(self):
        if not self.canvas.annotations:
            return
        reply = QMessageBox.question(
            self,
            "清空标注",
            "确定清空当前所有标注？",
            QMessageBox.Yes | QMessageBox.No,
            QMessageBox.No,
        )
        if reply == QMessageBox.Yes:
            self.canvas.clear_annotations()

    def _save_json(self):
        if self.canvas.frame is None:
            self._show_error("当前没有图像可保存。")
            return

        default_dir = os.path.join("examples", "annotations")
        os.makedirs(default_dir, exist_ok=True)
        default_path = os.path.join(
            default_dir, f"manual_annotation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        )
        path, _ = QFileDialog.getSaveFileName(
            self, "保存标注 JSON", default_path, "JSON (*.json);;所有文件 (*)"
        )
        if not path:
            return
        if not path.lower().endswith(".json"):
            path += ".json"

        payload = {
            "created_at": datetime.now().isoformat(timespec="seconds"),
            "source": self.source,
            "source_type": self.source_type,
            "image_size": self.canvas.image_size(),
            "coordinate_system": "origin_top_left_pixels",
            "annotations": self.canvas.annotations,
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)
        self._set_status(f"已保存: {path}")

    def _save_preview(self):
        if self.canvas.frame is None:
            self._show_error("当前没有图像可保存。")
            return

        preview = self._render_preview_image()
        default_dir = os.path.join("examples", "annotations")
        os.makedirs(default_dir, exist_ok=True)
        default_path = os.path.join(
            default_dir, f"manual_annotation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
        )
        path, _ = QFileDialog.getSaveFileName(
            self, "保存标注预览图", default_path, "PNG (*.png);;所有文件 (*)"
        )
        if not path:
            return
        if not path.lower().endswith(".png"):
            path += ".png"
        cv2.imwrite(path, preview)
        self._set_status(f"已保存预览图: {path}")

    def _render_preview_image(self):
        img = self.canvas.frame.copy()
        for ann in self.canvas.annotations:
            color = COLORS.get(ann["type"], QColor(255, 255, 255))
            bgr = (color.blue(), color.green(), color.red())
            if ann["type"] == "corner":
                x, y = [int(round(v)) for v in ann["points"][0]]
                cv2.drawMarker(img, (x, y), bgr, cv2.MARKER_CROSS, 18, 2)
                cv2.circle(img, (x, y), 5, bgr, 2)
                cv2.putText(img, ann.get("label", "corner"), (x + 8, y - 8),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, bgr, 1, cv2.LINE_AA)
            elif ann["type"] == "line":
                p0 = tuple(int(round(v)) for v in ann["points"][0])
                p1 = tuple(int(round(v)) for v in ann["points"][1])
                cv2.line(img, p0, p1, bgr, 2)
                cv2.putText(img, ann.get("label", "line"), (p0[0] + 8, p0[1] - 8),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, bgr, 1, cv2.LINE_AA)
            elif ann["type"] in ("rect", "roi"):
                x, y, w, h = ann["rect"]
                cv2.rectangle(img, (x, y), (x + w, y + h), bgr, 2)
                cv2.putText(img, ann.get("label", ann["type"]), (x + 8, y + 18),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, bgr, 1, cv2.LINE_AA)
        return img

    def _release_capture(self):
        self.timer.stop()
        if self.capture is not None:
            self.capture.release()
            self.capture = None

    def closeEvent(self, event):
        self._release_capture()
        event.accept()

    @staticmethod
    def _is_image_file(path):
        return os.path.splitext(path.lower())[1] in {
            ".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"
        }

    def _show_error(self, message):
        QMessageBox.warning(self, "错误", message)
        self._set_status(message)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--source",
        "-s",
        default="0",
        help="Camera index, video path, or image path. Default: 0",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
    QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)
    app = QApplication(sys.argv)
    window = MainWindow(args)
    window.resize(1280, 860)
    window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()
