from PySide6.QtCore import Qt, QPointF, QRectF, Signal, QSize, QLineF
from PySide6.QtGui import (QPixmap, QPainter, QPen, QBrush, QCursor, QColor, QImage,
                          QPainterPath, QTransform)
from PySide6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, 
                             QGraphicsView, QGraphicsScene, QGraphicsItem,
                             QGraphicsRectItem, QGraphicsEllipseItem, 
                             QGraphicsPixmapItem, QSizePolicy)
import math
class ResizableGraphicsItem(QGraphicsRectItem):
    handle_radius = 12          # 手柄半径（视觉+点击区域）
    rotate_offset = 0          # 旋转手柄离矩形顶边的距离

    def __init__(self, x, y, w, h, parent=None):
        super().__init__(x, y, w, h, parent)
        self.setFlags(QGraphicsItem.ItemIsSelectable |
                      QGraphicsItem.ItemIsMovable |
                      QGraphicsItem.ItemSendsGeometryChanges)
        self.setAcceptHoverEvents(True)
        self.setPen(QPen(Qt.blue, 2))
        self.setBrush(QBrush(Qt.transparent))

        self.setTransformOriginPoint(self.rect().center())

        self._rotating = False
        self._resize_corner = None
        self._orig_rect = QRectF()
        self._press_scene = QPointF()
        self._orig_rotation = 0.0

        self._rotate_anchor = QPointF()   # 旋转中心（item 坐标）
        self._last_vec      = QPointF()   # 上一帧鼠标相对向量

    def setRect(self, rect):
        super().setRect(rect)
        self.setTransformOriginPoint(rect.center())

    def paint(self, painter, option, widget=None):
        super().paint(painter, option, widget)
        if not self.isSelected():
            return

        painter.save()
        painter.setRenderHint(QPainter.Antialiasing)   # 消除残影
        painter.setPen(QPen(Qt.black, 1))

        for name, pt in self._handles_local().items():
            color = Qt.red if name == 'rotate' else Qt.green
            painter.setBrush(QBrush(color))
            painter.drawEllipse(pt, self.handle_radius, self.handle_radius)
        painter.restore()

    def drawTransformHandles(self, painter):
        """绘制变换控制点"""
        rect = self.rect()
        handle_size = self.handle_radius * 2
        
        # 四个角的控制点
        corners = [
            rect.topLeft(),
            rect.topRight(),
            rect.bottomLeft(),
            rect.bottomRight()
        ]
        
        # 旋转控制点（顶部中间）
        rotate_pos = QPointF(rect.center().x(), rect.top() - 20)
        
        # 绘制控制点
        painter.setBrush(QBrush(Qt.green))
        painter.setPen(QPen(Qt.black, 1))
        for corner in corners:
            painter.drawEllipse(corner, handle_size, handle_size)
        
        # 绘制旋转控制点
        painter.setBrush(QBrush(Qt.red))
        painter.drawEllipse(rotate_pos, handle_size, handle_size)

    def mouseMoveEvent(self, event):
        if self._rotating:
            # 当前鼠标向量
            cur_vec = event.scenePos() - self._rotate_center
            # 计算角度增量
            old_angle = math.atan2(self._press_vec.y(), self._press_vec.x())
            new_angle = math.atan2(cur_vec.y(), cur_vec.x())
            delta = math.degrees(new_angle - old_angle)

            # 直接累加旋转角，Qt 会按 TransformOriginPoint 旋转
            self.setRotation(self.rotation() + delta)

            self._press_vec = cur_vec
            event.accept()
            return

        if self._resize_corner:
            delta = event.pos() - self._press_local
            r = self._orig_rect
            new_r = QRectF(r)
            if self._resize_corner == 'topLeft':
                new_r.setTopLeft(r.topLeft() + delta)
            elif self._resize_corner == 'topRight':
                new_r.setTopRight(r.topRight() + delta)
            elif self._resize_corner == 'bottomLeft':
                new_r.setBottomLeft(r.bottomLeft() + delta)
            elif self._resize_corner == 'bottomRight':
                new_r.setBottomRight(r.bottomRight() + delta)

            if new_r.width() >= 10 and new_r.height() >= 10:
                self.setRect(new_r)
            event.accept()
            return

        super().mouseMoveEvent(event)

    def mousePressEvent(self, event):
        if event.button() != Qt.LeftButton:
            super().mousePressEvent(event)
            return

        hit = self._hit_handle(event.pos())
        if hit == 'rotate':
            self._rotating = True
            # 记录旋转中心（场景坐标）
            self._rotate_center = self.mapToScene(self.rect().center())
            # 记录鼠标第一帧向量（场景坐标）
            self._press_vec = event.scenePos() - self._rotate_center
            event.accept()
            return

        if hit in ('topLeft', 'topRight', 'bottomLeft', 'bottomRight'):
            self._resize_corner = hit
            self._orig_rect = self.rect()
            self._press_local = event.pos()
            event.accept()
            return
        
        if hit == 'rotate':
            self._rotating = True
            self._rotate_anchor = self.rect().center()          # 旋转中心
            self._last_vec = event.pos() - self._rotate_anchor  # 第一帧向量
            event.accept()
            return

        super().mousePressEvent(event)

    def mouseReleaseEvent(self, event):
        self._rotating = False
        self._resize_corner = None
        self.setCursor(Qt.ArrowCursor)
        super().mouseReleaseEvent(event)

    def hoverMoveEvent(self, event):
        if not self.isSelected():
            super().hoverMoveEvent(event)
            return
        hit = self._hit_handle(event.pos())
        self.setCursor(Qt.PointingHandCursor if hit == 'rotate' else
                       Qt.SizeFDiagCursor if hit in ('topLeft', 'bottomRight') else
                       Qt.SizeBDiagCursor if hit in ('topRight', 'bottomLeft') else
                       Qt.ArrowCursor)
        super().hoverMoveEvent(event)

    def _handles_local(self):
        r = self.rect()
        return {
            'topLeft'    : r.topLeft(),
            'topRight'   : r.topRight(),
            'bottomLeft' : r.bottomLeft(),
            'bottomRight': r.bottomRight(),
            'rotate'     : QPointF(r.center().x(), r.top() - self.rotate_offset)
        }

    def boundingRect(self):
        m = self.handle_radius + 5
        return self.rect().adjusted(-m, -m - self.rotate_offset, m, m)

    def _handles(self):
        """返回所有手柄的 scene 坐标"""
        rect = self.rect()
        center = self.mapToScene(rect.center())
        top_mid = self.mapToScene(QPointF(rect.center().x(), rect.top()))
        rotate = top_mid + QPointF(0, -20)  # 旋转手柄
        return {
            'topLeft'    : self.mapToScene(rect.topLeft()),
            'topRight'   : self.mapToScene(rect.topRight()),
            'bottomLeft' : self.mapToScene(rect.bottomLeft()),
            'bottomRight': self.mapToScene(rect.bottomRight()),
            'rotate'     : rotate
        }

    def _hit_handle(self, pos_local):
        for name, pt in self._handles_local().items():
            if (pos_local - pt).manhattanLength() <= self.handle_radius:
                return name
        return None
class ResizableEllipseItem(ResizableGraphicsItem):
    def paint(self, painter, option, widget=None):
        # 先画椭圆
        painter.setPen(QPen(Qt.red, 2))
        painter.setBrush(Qt.NoBrush)
        painter.drawEllipse(self.rect())

        # 选中时画控制手柄
        if self.isSelected():
            self.drawTransformHandles(painter)

class CustomGraphicsView(QGraphicsView):
    def __init__(self, scene, editor, parent=None):
        super().__init__(scene, parent)
        self.editor = editor
        self.setTransformationAnchor(QGraphicsView.AnchorUnderMouse)
        self.setResizeAnchor(QGraphicsView.AnchorUnderMouse)
        self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
        self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self._pan = False
        self._pan_start = QPointF()
        self._last_pan_value = QPointF()

    def wheelEvent(self, event):
        """鼠标滚轮缩放"""
        factor = 1.2
        if event.angleDelta().y() < 0:
            factor = 1.0 / factor
        
        self.scale(factor, factor)
        self.editor.scale_factor *= factor
        self.limit_view_to_image()
    
    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            scene_pos = self.mapToScene(event.position().toPoint())
            items = self.scene().items(scene_pos)

            # 如果点到了框线，让框线处理，且禁止 pan
            if any(isinstance(it, (ResizableGraphicsItem, ResizableEllipseItem))
                   for it in items):
                super().mousePressEvent(event)
                return

            # 选择模式下点空白，启动 pan
            if self.editor.mouse_state == "select":
                self.scene().clearSelection()   # 取消所有选中
                self._pan = True
                self._pan_start = event.position()
                self.setCursor(Qt.ClosedHandCursor)
                return

            if self.editor.mouse_state in ["create_rect", "create_ellipse"]:
                self.editor.start_pos = scene_pos
                if self.editor.mouse_state == "create_rect":
                    self.editor.current_item = ResizableGraphicsItem(0, 0, 0, 0)
                else:
                    self.editor.current_item = ResizableEllipseItem(0, 0, 0, 0)
                self.editor.current_item.setRect(QRectF(scene_pos, scene_pos))
                self.scene().addItem(self.editor.current_item)
                self.editor.current_item.setSelected(True)
                return

        super().mousePressEvent(event)


    def mouseMoveEvent(self, event):
        if self._pan:
            delta = event.position() - self._pan_start
            self._pan_start = event.position()
            hs = self.horizontalScrollBar()
            vs = self.verticalScrollBar()
            hs.setValue(hs.value() - int(delta.x()))
            vs.setValue(vs.value() - int(delta.y()))
            return
        if self.editor.current_item and self.editor.start_pos:
            scene_pos = self.mapToScene(event.position().toPoint())
            rect = QRectF(self.editor.start_pos, scene_pos).normalized()
            self.editor.current_item.setRect(rect)
            return

        super().mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.MiddleButton or self._pan:
            self._pan = False
            self.setCursor(Qt.ArrowCursor)
            return
        if self.editor.current_item:
            if self.editor.current_item.rect().width() < 10 or self.editor.current_item.rect().height() < 10:
                self.scene().removeItem(self.editor.current_item)
            else:
                self.editor.set_select_mode()
            self.editor.current_item = None
            self.editor.start_pos = None
            return

        super().mouseReleaseEvent(event)
    
    def limit_view_to_image(self):
        """限制视图范围不超过图片边界"""
        if not self.editor.pixmap_item:
            return
            
        # 确保滚动条不会超出范围
        h_scroll = self.horizontalScrollBar()
        v_scroll = self.verticalScrollBar()
        
        h_scroll.setValue(max(0, min(h_scroll.value(), h_scroll.maximum())))
        v_scroll.setValue(max(0, min(v_scroll.value(), v_scroll.maximum())))

class DefectMarkEditor(QDialog):
    def __init__(self, parent=None, image_path="", current_description=""):
        super().__init__(parent)
        self.setWindowTitle("缺陷标记编辑器")
        self.setMinimumSize(800, 600)
        
        self.image_path = image_path
        self.current_description = current_description
        self.mouse_state = "select"  # select | create_rect | create_ellipse
        self.start_pos = None
        self.current_item = None
        self.scale_factor = 1.0
        self.is_panning = False
        self.pan_start_pos = QPointF()
        self.max_view_size = QSize(1920, 1080)  # 最大视图尺寸
        self.pixmap_item = None

        self.init_ui()
        self.load_image()
    
    def init_ui(self):
        main_layout = QVBoxLayout(self)
        
        # 工具栏
        toolbar = QHBoxLayout()
        
        self.select_btn = QPushButton("选择")
        self.select_btn.setCheckable(True)
        self.select_btn.setChecked(True)
        self.select_btn.clicked.connect(self.set_select_mode)
        
        self.rect_btn = QPushButton("矩形框")
        self.rect_btn.setCheckable(True)
        self.rect_btn.clicked.connect(self.set_create_rect_mode)
        
        self.ellipse_btn = QPushButton("圆形框")
        self.ellipse_btn.setCheckable(True)
        self.ellipse_btn.clicked.connect(self.set_create_ellipse_mode)
        
        self.clear_btn = QPushButton("清除所有标记")
        self.clear_btn.clicked.connect(self.clear_all_items)
        
        toolbar.addWidget(self.select_btn)
        toolbar.addWidget(self.rect_btn)
        toolbar.addWidget(self.ellipse_btn)
        toolbar.addWidget(self.clear_btn)
        
        # 场景和视图
        self.scene = QGraphicsScene(self)
        self.view = CustomGraphicsView(self.scene, self)
        self.view.setRenderHint(QPainter.Antialiasing)
        self.view.setMouseTracking(True)
        self.view.setDragMode(QGraphicsView.NoDrag)
        
        main_layout.addLayout(toolbar)
        main_layout.addWidget(self.view)
    
    def load_image(self):
        """加载图片到场景"""
        if not self.image_path:
            return
            
        self.scene.clear()
        
        # 加载图片并计算合适的缩放比例
        image = QImage(self.image_path)
        if image.isNull():
            return
            
        # 计算缩放比例以适应视图
        img_size = image.size()
        
        # 如果图片大于最大视图尺寸，则缩放
        if img_size.width() > self.max_view_size.width() or img_size.height() > self.max_view_size.height():
            img_size.scale(self.max_view_size, Qt.KeepAspectRatio)
            scaled_pixmap = QPixmap.fromImage(image.scaled(img_size, Qt.KeepAspectRatio, Qt.SmoothTransformation))
        else:
            scaled_pixmap = QPixmap.fromImage(image)
        
        self.pixmap_item = QGraphicsPixmapItem(scaled_pixmap)
        self.pixmap_item.setFlag(QGraphicsItem.ItemIsMovable, False)
        self.pixmap_item.setFlag(QGraphicsItem.ItemIsSelectable, False)
        self.scene.addItem(self.pixmap_item)
        
        # 设置场景大小为图片大小
        self.scene.setSceneRect(self.pixmap_item.boundingRect())
        
        # 自动调整视图以显示整个图片
        self.view.fitInView(self.pixmap_item, Qt.KeepAspectRatio)
    
    def set_select_mode(self):
        """设置为选择模式"""
        self.mouse_state = "select"
        self.select_btn.setChecked(True)
        self.rect_btn.setChecked(False)
        self.ellipse_btn.setChecked(False)
        self.view.setCursor(Qt.ArrowCursor)
    
    def set_create_rect_mode(self):
        """设置为创建矩形模式"""
        self.mouse_state = "create_rect"
        self.select_btn.setChecked(False)
        self.rect_btn.setChecked(True)
        self.ellipse_btn.setChecked(False)
        self.view.setCursor(Qt.CrossCursor)
    
    def set_create_ellipse_mode(self):
        """设置为创建圆形模式"""
        self.mouse_state = "create_ellipse"
        self.select_btn.setChecked(False)
        self.rect_btn.setChecked(False)
        self.ellipse_btn.setChecked(True)
        self.view.setCursor(Qt.CrossCursor)
    
    def clear_all_items(self):
        """清除所有标记项"""
        for item in self.scene.items():
            if isinstance(item, (ResizableGraphicsItem, ResizableEllipseItem)):
                self.scene.removeItem(item)

    def export_annotated_image(self, output_dir):
        """
        将当前场景（原图 + 所有标记）按原分辨率导出为 PNG/JPG，
        文件名与原图完全一致，保存到 output_dir。
        """
        import os
        from PySide6.QtGui import QPainter, QImage

        if not self.pixmap_item:
            return False

        # 1. 取原图 QImage（保持原始分辨率）
        source_image = QImage(self.image_path)
        if source_image.isNull():
            return False

        # 2. 创建与原图同尺寸的 QImage 作为画布
        out_image = QImage(source_image.size(), QImage.Format_ARGB32)
        out_image.fill(Qt.transparent)

        painter = QPainter(out_image)
        painter.setRenderHint(QPainter.Antialiasing)

        # 3. 先把原图完整画上去
        painter.drawImage(0, 0, source_image)

        # 4. 计算场景 → 原图坐标 的缩放因子
        scene_rect = self.pixmap_item.sceneBoundingRect()
        sx = source_image.width()  / scene_rect.width()
        sy = source_image.height() / scene_rect.height()

        # 5. 逐个绘制标记（仅最终形状，不绘制手柄）
        for item in self.scene.items():
            if isinstance(item, (ResizableGraphicsItem, ResizableEllipseItem)):
                # 关闭选中状态，避免手柄被绘制
                was_selected = item.isSelected()
                item.setSelected(False)

                # 计算 item 在场景中的实际几何
                shape_path = item.mapToScene(item.shape())          # 考虑旋转
                painter.setTransform(QTransform())                  # 重置 painter
                painter.scale(sx, sy)                               # 映射到原图坐标
                painter.translate(-scene_rect.x(), -scene_rect.y()) # 对齐偏移

                # 设置画笔
                pen = item.pen()
                pen.setCosmetic(False)  # 让线宽随缩放
                painter.setPen(pen)
                painter.setBrush(Qt.NoBrush)

                # 绘制形状
                if isinstance(item, ResizableEllipseItem):
                    painter.drawPath(shape_path)
                else:
                    painter.drawPath(shape_path)

                # 恢复选中状态
                item.setSelected(was_selected)

        painter.end()

        # 6. 保存文件
        os.makedirs(output_dir, exist_ok=True)
        base_name = os.path.basename(self.image_path)
        save_path = os.path.join(output_dir, base_name)
        return out_image.save(save_path)
    
    def closeEvent(self, event):
        """关闭窗口时保存标记"""
        self.export_annotated_image("/home/dtyx/桌面/yhh/ReportGenerator/output")
        super().closeEvent(event)

if __name__ == '__main__':
    from PySide6.QtWidgets import QApplication
    import sys
    app = QApplication(sys.argv)
    editor = DefectMarkEditor(image_path=r"/home/dtyx/桌面/yhh/ReportGenerator/测试数据/山东国华无棣风电场叶片外部数据/A1（31301）/1#（1-GW68.6C-I190809C）/DJI_20250606162745_0005_Z.JPG")
    editor.show()
    sys.exit(app.exec())
