import socket
import time
import select
import sys
import micropython

class VideoStreamServer:
    def __init__(self, display, tcp_port=8081, udp_port=8082):
        self.display = display
        self.tcp_port = tcp_port
        self.udp_port = udp_port
        
        # Sockets
        self.tcp_server = None
        self.tcp_client = None
        self.udp_sock = None
        
        # State
        self.active = False
        self.last_packet_time = 0
        self.timeout_ms = 3000
        
        # Frame buffering (Pre-allocated to prevent memory allocations)
        self.buffer = bytearray(15000)
        self.view = memoryview(self.buffer)
        self.tcp_bytes_received = 0
        
        # UDP Reassembly (15 chunks of 1000 bytes each)
        self.udp_temp_buffer = bytearray(1002)
        self.udp_view = memoryview(self.udp_temp_buffer)
        self.current_frame_id = -1
        self.chunks_received = 0  # Bitmask for chunks 0-14
        
        # Performance Stats
        self.debug = False
        self.frames_drawn = 0
        self.udp_packets_received = 0
        self.udp_frames_complete = 0
        self.dropped_udp_frames = 0
        self.last_draw_ms = 0
        self.last_fps = 0.0
        
        # FPS Calculation
        self.fps_start_time = time.ticks_ms()
        self.fps_frame_count = 0

    @micropython.native
    def rlcd_to_mono(self, rlcd_buf, canvas_buf):
        # canvas_buf size: 19200 bytes (480x320)
        # Clear canvas_buf to 0
        for i in range(19200):
            canvas_buf[i] = 0
            
        X_OFF = 40
        Y_OFF = 10
        
        # Loop over 15000 bytes
        for index in range(15000):
            val = rlcd_buf[index]
            if val == 0:
                continue
                
            byte_x = index // 75
            block_y = index % 75
            x_base = 2 * byte_x + X_OFF
            y_base = 299 - 4 * block_y + Y_OFF
            
            # bit 7: local_y = 0, local_x = 0
            if val & 0x80:
                cx = x_base
                cy = y_base
                canvas_buf[cy * 60 + (cx >> 3)] |= (1 << (7 - (cx & 7)))
                
            # bit 6: local_y = 0, local_x = 1
            if val & 0x40:
                cx = x_base + 1
                cy = y_base
                canvas_buf[cy * 60 + (cx >> 3)] |= (1 << (7 - (cx & 7)))
                
            # bit 5: local_y = 1, local_x = 0
            if val & 0x20:
                cx = x_base
                cy = y_base - 1
                canvas_buf[cy * 60 + (cx >> 3)] |= (1 << (7 - (cx & 7)))
                
            # bit 4: local_y = 1, local_x = 1
            if val & 0x10:
                cx = x_base + 1
                cy = y_base - 1
                canvas_buf[cy * 60 + (cx >> 3)] |= (1 << (7 - (cx & 7)))
                
            # bit 3: local_y = 2, local_x = 0
            if val & 0x08:
                cx = x_base
                cy = y_base - 2
                canvas_buf[cy * 60 + (cx >> 3)] |= (1 << (7 - (cx & 7)))
                
            # bit 2: local_y = 2, local_x = 1
            if val & 0x04:
                cx = x_base + 1
                cy = y_base - 2
                canvas_buf[cy * 60 + (cx >> 3)] |= (1 << (7 - (cx & 7)))
                
            # bit 1: local_y = 3, local_x = 0
            if val & 0x02:
                cx = x_base
                cy = y_base - 3
                canvas_buf[cy * 60 + (cx >> 3)] |= (1 << (7 - (cx & 7)))
                
            # bit 0: local_y = 3, local_x = 1
            if val & 0x01:
                cx = x_base + 1
                cy = y_base - 3
                canvas_buf[cy * 60 + (cx >> 3)] |= (1 << (7 - (cx & 7)))

    def start(self):
        """Initialize and start the sockets."""
        # 1. Start TCP Server
        try:
            self.tcp_server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.tcp_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.tcp_server.bind(('', self.tcp_port))
            self.tcp_server.listen(1)
            self.tcp_server.setblocking(False)
            print(f"TCP Stream Server started on port {self.tcp_port}")
        except Exception as e:
            print(f"Failed to start TCP stream server: {e}")
            
        # 2. Start UDP Server
        try:
            self.udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.udp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.udp_sock.bind(('', self.udp_port))
            self.udp_sock.setblocking(False)
            print(f"UDP Stream Server started on port {self.udp_port}")
        except Exception as e:
            print(f"Failed to start UDP stream server: {e}")

    def restart_tcp_server(self):
        """Re-initializes the TCP stream server socket after a fatal error."""
        print("Restarting TCP Stream Server...")
        self.close_tcp_client()
        try:
            if self.tcp_server:
                self.tcp_server.close()
        except:
            pass
        self.tcp_server = None
        time.sleep_ms(100)
        try:
            self.tcp_server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.tcp_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.tcp_server.bind(('', self.tcp_port))
            self.tcp_server.listen(1)
            self.tcp_server.setblocking(False)
            print(f"TCP Stream Server restarted on port {self.tcp_port}")
        except Exception as e:
            print(f"Failed to restart TCP stream server: {e}")

    def restart_udp_sock(self):
        """Re-initializes the UDP stream server socket after a fatal error."""
        print("Restarting UDP Stream Server...")
        try:
            if self.udp_sock:
                self.udp_sock.close()
        except:
            pass
        self.udp_sock = None
        time.sleep_ms(100)
        try:
            self.udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.udp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.udp_sock.bind(('', self.udp_port))
            self.udp_sock.setblocking(False)
            print(f"UDP Stream Server restarted on port {self.udp_port}")
        except Exception as e:
            print(f"Failed to restart UDP stream server: {e}")

    def update(self):
        """Updates connection states and checks for incoming video data.
        Call this periodically in the main execution loop.
        """
        now = time.ticks_ms()
        
        # Calculate dynamic FPS every 1 second
        fps_elapsed = time.ticks_diff(now, self.fps_start_time)
        if fps_elapsed >= 1000:
            self.last_fps = (self.fps_frame_count * 1000.0) / fps_elapsed
            self.fps_frame_count = 0
            self.fps_start_time = now
        
        # 1. Check for inactivity timeout
        if self.active and time.ticks_diff(now, self.last_packet_time) > self.timeout_ms:
            print("Video stream timed out. Returning to dashboard.")
            self.active = False
            self.tcp_bytes_received = 0
            
        # 2. Handle incoming UDP packet (if UDP is active)
        if self.udp_sock:
            while True:
                try:
                    n = self.udp_sock.readinto(self.udp_temp_buffer)
                    if n is None or n == 0:
                        break # Socket buffer queue is empty
                        
                    self.udp_packets_received += 1
                    frame_id = self.udp_temp_buffer[0]
                    chunk_idx = self.udp_temp_buffer[1]
                    
                    if self.debug:
                        print(f"[UDP] Received packet size {n}: frame={frame_id}, chunk={chunk_idx}")
                    
                    if chunk_idx < 15:
                        # If a new frame sequence starts, reset tracking and check if old frame was incomplete
                        if frame_id != self.current_frame_id:
                            if self.chunks_received != 0:
                                self.dropped_udp_frames += 1
                            self.current_frame_id = frame_id
                            self.chunks_received = 0
                        
                        # Copy chunk data directly to frame buffer
                        start_offset = chunk_idx * 1000
                        self.buffer[start_offset : start_offset + 1000] = self.udp_view[2:1002]
                        
                        # Mark chunk index as received
                        self.chunks_received |= (1 << chunk_idx)
                        
                        # Check if all 15 chunks (indices 0 to 14) are received (0x7FFF)
                        if self.chunks_received == 0x7FFF:
                            self.udp_frames_complete += 1
                            if self.debug:
                                print(f"[UDP] Frame {frame_id} complete! Drawing to screen.")
                            self.active = True
                            self.last_packet_time = now
                            self._draw_frame()
                            self.chunks_received = 0  # Reset for next frame
                except OSError as e:
                    import errno
                    err = getattr(e, 'errno', None)
                    if err is None and e.args:
                        err = e.args[0]
                    ewouldblock = getattr(errno, 'EWOULDBLOCK', errno.EAGAIN)
                    if err in (errno.EAGAIN, ewouldblock) or err is None:
                        break # EAGAIN/EWOULDBLOCK (no data available)
                    print(f"Fatal UDP socket error in update(): {e}. Re-initializing UDP socket...")
                    self.restart_udp_sock()
                    break

        # 3. Handle TCP stream
        if self.tcp_server:
            # If no client is connected, try to accept one
            if self.tcp_client is None:
                try:
                    self.tcp_client, addr = self.tcp_server.accept()
                    self.tcp_client.setblocking(False)
                    self.tcp_bytes_received = 0
                    self.active = True
                    self.last_packet_time = now
                    print(f"TCP Stream client connected from: {addr}")
                except OSError as e:
                    import errno
                    err = getattr(e, 'errno', None)
                    if err is None and e.args:
                        err = e.args[0]
                    ewouldblock = getattr(errno, 'EWOULDBLOCK', errno.EAGAIN)
                    if err not in (errno.EAGAIN, ewouldblock) and err is not None:
                        print(f"Fatal TCP stream server socket error in accept(): {e}. Re-initializing...")
                        self.restart_tcp_server()
            
            # If client is connected, read data
            if self.tcp_client is not None:
                try:
                    remaining = 15000 - self.tcp_bytes_received
                    n = self.tcp_client.readinto(self.view[self.tcp_bytes_received : self.tcp_bytes_received + remaining])
                    
                    if n is not None:
                        if n == 0:
                            # Client closed connection
                            print("TCP Stream client disconnected.")
                            self.close_tcp_client()
                        else:
                            self.tcp_bytes_received += n
                            if self.debug:
                                print(f"[TCP] Received {n} bytes (Total: {self.tcp_bytes_received}/15000)")
                            self.last_packet_time = now
                            self.active = True
                            
                            # If we have a full frame, render it
                            if self.tcp_bytes_received == 15000:
                                if self.debug:
                                    print("[TCP] Frame complete! Drawing to screen.")
                                self.tcp_bytes_received = 0
                                self._draw_frame()
                except OSError as e:
                    # Connection reset or socket error
                    print(f"TCP Stream socket error: {e}")
                    self.close_tcp_client()

    def _draw_frame(self):
        """Sends the frame buffer directly to the SPI display bus."""
        draw_start = time.ticks_ms()
        
        if sys.platform == 'rp2':
            # Decode the 15,000-byte RLCD buffer and write it centered on the 480x320 canvas
            self.rlcd_to_mono(self.buffer, self.display.canvas_buffer)
            self.display.show()
        else:
            # Frame boundary command sequences for the RLCD display
            self.display.write_cmd(0x2A); self.display.write_data(0x12); self.display.write_data(0x2A)
            self.display.write_cmd(0x2B); self.display.write_data(0x00); self.display.write_data(0xC7)
            self.display.write_cmd(0x2C) 
            
            # Send raw 15,000 bytes directly
            self.display.write_data(self.buffer)
        
        self.last_draw_ms = time.ticks_diff(time.ticks_ms(), draw_start)
        self.frames_drawn += 1
        self.fps_frame_count += 1

    def get_stats(self):
        """Returns streaming performance counters."""
        return {
            "frames_drawn": self.frames_drawn,
            "tcp_bytes_received": self.tcp_bytes_received,
            "udp_packets_received": self.udp_packets_received,
            "udp_frames_complete": self.udp_frames_complete,
            "dropped_udp_frames": self.dropped_udp_frames,
            "last_draw_ms": self.last_draw_ms,
            "last_fps": round(self.last_fps, 1),
            "active": self.active
        }

    def close_tcp_client(self):
        """Safely disconnects the TCP client socket."""
        if self.tcp_client:
            try:
                self.tcp_client.close()
            except:
                pass
            self.tcp_client = None
        self.tcp_bytes_received = 0