import sys
import json
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QLabel, QPushButton, QColorDialog, QInputDialog,
    QFileDialog, QMessageBox, QVBoxLayout, QHBoxLayout, QWidget
)
from PyQt5.QtGui import QImage, QPixmap, QPainter, QPen, QColor
from PyQt5.QtCore import Qt, QPoint

class ImageLabel(QLabel):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setMouseTracking(True)
        self.image = QImage()
        self.rois = []
        self.current_roi = []
        self.drawing = False
        self.current_color = QColor(Qt.red)
        self.current_label = "ROI"
        self.undo_stack = []

    def load_image(self, file_path):
        self.image = QImage(file_path)
        if self.image.isNull():
            raise ValueError("Failed to load image.")
        self.setPixmap(QPixmap.fromImage(self.image))
        self.rois.clear()
        self.current_roi.clear()
        self.update()

    def set_current_color(self, color):
        self.current_color = color

    def set_current_label(self, label):
        self.current_label = label

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.drawing = True
            self.current_roi = [event.pos()]
            self.update()

    def mouseMoveEvent(self, event):
        if self.drawing:
            self.current_roi.append(event.pos())
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self.drawing:
            self.drawing = False
            if len(self.current_roi) > 2:
                self.rois.append({
                    'label': self.current_label,
                    'color': self.current_color,
                    'points': self.current_roi
                })
                self.undo_stack.append(self.rois[-1])
            self.current_roi = []
            self.update()

    def paintEvent(self, event):
        super().paintEvent(event)
        if not self.image.isNull():
            painter = QPainter(self)
            painter.drawImage(self.rect(), self.image, self.image.rect())
            pen = QPen(Qt.red, 2, Qt.SolidLine)
            painter.setPen(pen)
            for roi in self.rois:
                pen.setColor(roi['color'])
                painter.setPen(pen)
                points = [QPoint(p.x(), p.y()) for p in roi['points']]
                painter.drawPolygon(*points)
            if self.current_roi:
                pen.setColor(Qt.blue)
                painter.setPen(pen)
                points = [QPoint(p.x(), p.y()) for p in self.current_roi]
                painter.drawPolyline(*points)
            painter.end()

    def undo(self):
        if self.undo_stack:
            last_action = self.undo_stack.pop()
            if last_action in self.rois:
                self.rois.remove(last_action)
            self.update()

    def save_rois(self, file_path):
        roi_data = []
        for roi in self.rois:
            roi_data.append({
                'label': roi['label'],
                'color': roi['color'].name(),
                'points': [(point.x(), point.y()) for point in roi['points']]
            })
        with open(file_path, 'w') as file:
            json.dump(roi_data, file, indent=4)

class ROIDrawer(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setWindowTitle('ROI Drawer')
        self.setGeometry(100, 100, 1000, 600)

        # Central widget
        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)

        # Layouts
        main_layout = QHBoxLayout(central_widget)
        image_layout = QVBoxLayout()
        button_layout = QVBoxLayout()

        # Image display
        self.image_label = ImageLabel(self)
        image_layout.addWidget(self.image_label)

        # Buttons
        load_button = QPushButton('Load Image', self)
        load_button.clicked.connect(self.load_image)
        button_layout.addWidget(load_button)

        color_button = QPushButton('Select Color', self)
        color_button.clicked.connect(self.select_color)
        button_layout.addWidget(color_button)

        label_button = QPushButton('Set Label', self)
        label_button.clicked.connect(self.set_label)
        button_layout.addWidget(label_button)

        undo_button = QPushButton('Undo', self)
        undo_button.clicked.connect(self.image_label.undo)
        button_layout.addWidget(undo_button)

        save_button = QPushButton('Save ROIs', self)
        save_button.clicked.connect(self.save_rois)
        button_layout.addWidget(save_button)

        button_layout.addStretch(1)  # Push buttons to the top

        # Add layouts to main layout
        main_layout.addLayout(image_layout, 4)
        main_layout.addLayout(button_layout, 1)

    def load_image(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Open Image File", "", "Images (*.png *.jpg *.bmp *.tiff)", options=options)
        if file_path:
            try:
                self.image_label.load_image(file_path)
            except ValueError as e:
                QMessageBox.critical(self, "Error", str(e))

    def select_color(self):
        color = QColorDialog.getColor()
        if color.isValid():
            self.image_label.set_current_color(color)

    def set_label(self):
        label, ok = QInputDialog.getText(self, 'Set ROI Label', 'Enter label for ROI:')
        if ok and label:
            self.image_label.set_current_label(label)

    def save_rois(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getSaveFileName(
            self, "Save ROIs", "", "JSON Files (*.json);;All Files (*)", options=options)
        if file_path:
            self.image_label.save_rois(file_path)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    drawer = ROIDrawer()
    drawer.show()
    sys.exit(app.exec_())