# dmx.py

from PyQt5.QtCore import QObject, QThread, pyqtSignal
import socket
import struct
import numpy as np
import subprocess
import time

class DMXRecorder(QObject):
    def __init__(self):
        super().__init__()
        self.thread = None
        self.recording_thread = None

    def start_recording(self):
        self.thread = QThread()
        self.recording_thread = RecordingThread()
        self.recording_thread.moveToThread(self.thread)

        self.thread.started.connect(self.recording_thread.run)
        self.recording_thread.finished.connect(self.thread.quit)
        self.recording_thread.finished.connect(self.recording_thread.deleteLater)
        self.thread.finished.connect(self.thread.deleteLater)

        self.thread.start()

    def stop_recording(self):
        if self.recording_thread:
            self.recording_thread.stop()
        if self.thread:
            self.thread.quit()
            self.thread.wait()

class RecordingThread(QObject):
    finished = pyqtSignal()

    def __init__(self):
        super().__init__()
        self.running = False
        self.socket = None
        self.ffmpeg_process = None

    def run(self):
        self.running = True

        # Set up UDP socket to listen for Art-Net DMX data
        UDP_IP = "0.0.0.0"
        UDP_PORT = 6454

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind((UDP_IP, UDP_PORT))

        # Set socket to non-blocking
        self.socket.setblocking(False)

        # Prepare for recording
        expected_universes = set(range(4))  # Universes 0 to 3 (adjust as needed)
        num_universes = len(expected_universes)
        universe_size = 512

        # Define frame dimensions
        frame_width = 64  # Adjust as needed (must divide evenly into total data size)
        total_pixels = num_universes * universe_size
        frame_height = total_pixels // frame_width

        # Frame rate and timing
        frame_rate = 30  # frames per second
        frame_interval = 1.0 / frame_rate
        last_frame_time = time.time()

        # Set up ffmpeg process
        ffmpeg_cmd = [
            'ffmpeg',
            '-y',
            '-f', 'rawvideo',
            '-pix_fmt', 'gray',
            '-s', f'{frame_width}x{frame_height}',
            '-r', str(frame_rate),
            '-i', '-',  # Input from stdin
            '-c:v', 'ffv1',
            'dmx_video.mkv'
        ]

        self.ffmpeg_process = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE)

        universes_data = {}
        last_received = {}

        try:
            while self.running:
                current_time = time.time()

                # Receive data
                try:
                    data, addr = self.socket.recvfrom(1024)
                    if data.startswith(b'Art-Net'):
                        op_code = struct.unpack('<H', data[8:10])[0]
                        # OpCode for ArtDMX is 0x5000
                        if op_code == 0x5000:
                            # Parse universe
                            universe = struct.unpack('<H', data[14:16])[0]
                            # Get length
                            length = struct.unpack('>H', data[16:18])[0]
                            dmx_data = data[18:18+length]
                            # Store the data for this universe
                            universes_data[universe] = dmx_data
                            last_received[universe] = current_time
                except BlockingIOError:
                    # No data received
                    pass
                except Exception as e:
                    print(f"Error receiving data: {e}")

                # Check if it's time to write a frame
                if current_time - last_frame_time >= frame_interval:
                    # Assemble frame data
                    frame_data = b''
                    for universe in sorted(expected_universes):
                        dmx_data = universes_data.get(universe, b'\x00'*512)
                        if len(dmx_data) < 512:
                            dmx_data += b'\x00' * (512 - len(dmx_data))
                        frame_data += dmx_data
                    # Convert frame_data to numpy array
                    frame_array = np.frombuffer(frame_data, dtype=np.uint8)
                    # Reshape to image dimensions
                    frame_array = frame_array.reshape((frame_height, frame_width))
                    # Write frame to ffmpeg stdin
                    self.ffmpeg_process.stdin.write(frame_array.tobytes())
                    last_frame_time = current_time

                time.sleep(0.001)  # Sleep briefly to avoid maxing out CPU
        finally:
            # Clean up
            self.ffmpeg_process.stdin.close()
            self.ffmpeg_process.wait()
            self.socket.close()
            self.finished.emit()

    def stop(self):
        self.running = False