import sys
import cv2
import numpy as np
from PyQt6.QtWidgets import QApplication, QLabel, QPushButton, QVBoxLayout, QWidget, QFileDialog, QSlider, QHBoxLayout
from PyQt6.QtCore import Qt, QTimer, QPoint
from PyQt6.QtGui import QImage, QPixmap, QMouseEvent


class ChromaPreviewer(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Live Chroma Previewer")
        self.setGeometry(100, 100, 800, 600)

        # Video capture
        self.cap = cv2.VideoCapture(0)

        # Default chroma key values
        self.hue_center = 50  # Default to greenish hue
        self.tolerance = 10  # Default tolerance
        self.lower_bound = np.array([self.hue_center - self.tolerance, 50, 50], dtype=np.uint8)
        self.upper_bound = np.array([self.hue_center + self.tolerance, 255, 255], dtype=np.uint8)
        self.background = None
        self.bg_video = None
        self.bg_video_cap = None

        # UI Elements
        self.video_label = QLabel(self)
        self.video_label.mousePressEvent = self.pick_color  # Allow picking chroma key color
        self.load_bg_button = QPushButton("Load Background Image/Video")
        self.load_bg_button.clicked.connect(self.load_background)

        self.tolerance_slider = QSlider(Qt.Orientation.Horizontal)
        self.tolerance_slider.setRange(0, 50)
        self.tolerance_slider.setValue(self.tolerance)
        self.tolerance_slider.valueChanged.connect(self.update_tolerance)

        self.layout = QVBoxLayout()
        self.layout.addWidget(self.video_label)
        self.layout.addWidget(self.load_bg_button)

        slider_layout = QHBoxLayout()
        slider_layout.addWidget(QLabel("Tolerance"))
        slider_layout.addWidget(self.tolerance_slider)
        self.layout.addLayout(slider_layout)

        self.setLayout(self.layout)

        # Timer for updating frames
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        self.timer.start(30)

    def load_background(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select Background Image or Video", "",
                                                   "Images/Videos (*.png *.jpg *.jpeg *.mp4 *.avi *.mov)")

        if file_path:
            if file_path.lower().endswith(('.mp4', '.avi', '.mov')):
                self.bg_video_cap = cv2.VideoCapture(file_path)
                self.bg_video = True
                self.background = None  # Clear image background if video is used
            else:
                self.background = cv2.imread(file_path)
                self.background = cv2.cvtColor(self.background, cv2.COLOR_BGR2RGB)
                self.bg_video = False
                self.bg_video_cap = None

    def update_tolerance(self, value):
        self.tolerance = value
        self.update_chroma_key()

    def pick_color(self, event: QMouseEvent):
        x, y = event.pos().x(), event.pos().y()
        if hasattr(self, 'current_frame') and self.current_frame is not None:
            h, w, _ = self.current_frame.shape
            if 0 <= x < w and 0 <= y < h:
                pixel_color = self.current_frame[y, x]
                hsv_color = cv2.cvtColor(np.uint8([[pixel_color]]), cv2.COLOR_RGB2HSV)[0][0]

                self.hue_center = int(hsv_color[0])  # Update center hue
                self.update_chroma_key()

    def update_chroma_key(self):
        self.lower_bound = np.array([
            np.clip(self.hue_center - self.tolerance, 0, 179),
            50, 50
        ], dtype=np.uint8)

        self.upper_bound = np.array([
            np.clip(self.hue_center + self.tolerance, 0, 179),
            255, 255
        ], dtype=np.uint8)

    def update_frame(self):
        try:
            ret, frame = self.cap.read()
            if not ret:
                return
        except AttributeError:
            return

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        self.current_frame = frame.copy()
        hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)

        mask = cv2.inRange(hsv, self.lower_bound, self.upper_bound)
        mask_inv = cv2.bitwise_not(mask)

        fg = cv2.bitwise_and(frame, frame, mask=mask_inv)

        if self.bg_video and self.bg_video_cap:
            ret_bg, bg_frame = self.bg_video_cap.read()
            if not ret_bg:
                self.bg_video_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)  # Loop the video
                ret_bg, bg_frame = self.bg_video_cap.read()

            if ret_bg:
                bg_frame = cv2.cvtColor(bg_frame, cv2.COLOR_BGR2RGB)
                bg_resized = cv2.resize(bg_frame, (frame.shape[1], frame.shape[0]))
                bg = cv2.bitwise_and(bg_resized, bg_resized, mask=mask)
                frame = cv2.add(fg, bg)
        elif self.background is not None:
            bg_resized = cv2.resize(self.background, (frame.shape[1], frame.shape[0]))
            bg = cv2.bitwise_and(bg_resized, bg_resized, mask=mask)
            frame = cv2.add(fg, bg)
        else:
            frame = fg

        h, w, ch = frame.shape
        bytes_per_line = ch * w
        qimg = QImage(frame.data, w, h, bytes_per_line, QImage.Format.Format_RGB888)
        self.video_label.setPixmap(QPixmap.fromImage(qimg))

    def closeEvent(self, event):
        self.cap.release()
        if self.bg_video_cap:
            self.bg_video_cap.release()
        event.accept()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ChromaPreviewer()
    window.show()
    sys.exit(app.exec())