#!/usr/bin/env python3
import serial
import struct
from PIL import Image
import numpy as np
import time
import sys
import csv

class RP2350BufferReader:
    def __init__(self, port, baudrate=115200):
        self.Ser = serial.Serial(port, baudrate, timeout=5)
        self.Width = 0
        self.Height = 0
        self.Bits_Per_Pixel = 0
        
    def Read_Display_Buffer(self, image_filename=None, csv_filename=None):
        """Request and read display buffer from RP2350"""
        
        # Send buffer read command
        self.Ser.write(b'b')
        
        # Read header
        Header = self.Ser.read_until(b'\r').decode('ascii').strip()
        if not Header.startswith('IMGBUF'):
            raise ValueError(f"Invalid header: {Header}")
            
        # Parse dimensions
        Parts = Header[6:].split(',')
        self.Width = int(Parts[0])
        self.Height = int(Parts[1])
        self.Bits_Per_Pixel = int(Parts[2])
        
        print(f"Reading {self.Width}x{self.Height} buffer ({self.Bits_Per_Pixel}bpp)")
        
        # Read binary chunked data with progress bar
        Pixel_Data_ASCII : str = ""
        Expected_Pixels = self.Width * self.Height
        Expected_Byes = Expected_Pixels * (self.Bits_Per_Pixel // 8) * 2
        Start_Time = time.time()
        
        def Update_Progress(current, total, width=50):
            Percent = (current / total) * 100
            Filled = int(width * current / total)
            Bar = '█' * Filled + '░' * (width - Filled)
            Elapsed = time.time() - Start_Time
            Rate = current / Elapsed if Elapsed > 0 else 0
            Eta = (total - current) / Rate if Rate > 0 else 0
            
            sys.stdout.write(f'\r[{Bar}] {Percent:.1f}% ({current}/{total}) 'f'Rate: {Rate:.0f} px/s ETA: {Eta:.1f}s    ')
            sys.stdout.flush()
        
        print("Progress:")
        while len(Pixel_Data_ASCII) < Expected_Byes:
            Byte = self.Ser.read_all()
            if Byte is not None:
                Pixel_Data_ASCII = Pixel_Data_ASCII + Byte.decode("utf-8")

            Update_Progress(len(Pixel_Data_ASCII), Expected_Byes)

        
        Update_Progress(len(Pixel_Data_ASCII), Expected_Byes) # Final progress update
        print()

        Byte = self.Ser.read(1)
        Byte_Value = ord(Byte.decode('utf-8'))

        if Byte_Value != 0x0D:
            raise ValueError(f"Invalid byte after pixel data: {str(int(Byte))}")

        # Save ASCII data to CSV if requested
        if csv_filename:
            self._Save_ASCII_To_CSV(Pixel_Data_ASCII, csv_filename)

        Pixel_Data_Hex = []
        for i in range(0, len(Pixel_Data_ASCII), 4):
            Int_Value = int(Pixel_Data_ASCII[i:i+4], 16)
            # if Int_Value > 0:
            #     print("0x{:04x}".format(Int_Value))
            Pixel_Data_Hex.append(Int_Value)
        
        # Convert to numpy array
        RGB565_Data = np.array(Pixel_Data_Hex, dtype=np.uint16)
        RGB888_Data = self._RGB565_to_RGB888(RGB565_Data)
        
        # Create PIL Image
        img_array = RGB888_Data.reshape((self.Height, self.Width, 3))
        image = Image.fromarray(img_array, 'RGB')
        # image = Image.new('RGB', (self.Width, self.Height))
        
        # Save if filename provided
        if image_filename:
            if not image_filename.lower().endswith(('.png', '.bmp', '.tiff')):
                image_filename += '.png'
            image.save(image_filename)
            print(f"Image saved as {image_filename}")
            
        return image
        
    def _RGB565_to_RGB888(self, rgb565_array):
        """Convert RGB565 to RGB888"""
        RGB888_Data = np.zeros((len(rgb565_array), 3), dtype=np.uint8)
        
        # Extract RGB components from RGB565
        r5 = ((rgb565_array & 0x1F00) >> 8)
        g6 = ((rgb565_array & 0xE000) >> 11) | ((rgb565_array & 0x0007) << 3)
        b5 = ((rgb565_array & 0x00F8) >> 3)
        
        # Scale to 8-bit
        RGB888_Data[:, 0] = (r5 * 255) // 31  # Red
        RGB888_Data[:, 1] = (g6 * 255) // 63  # Green  
        RGB888_Data[:, 2] = (b5 * 255) // 31  # Blue
        
        return RGB888_Data
    
    
    def _Save_ASCII_To_CSV(self, pixel_data_ascii, csv_filename):
        """Save ASCII pixel data to CSV file with display dimensions"""
        if not csv_filename.lower().endswith('.csv'):
            csv_filename += '.csv'
                 
        with open(csv_filename, 'w', newline='') as CSV_File:
            Writer = csv.writer(CSV_File)
            
            # Process data in chunks of 4 characters (one pixel)
            Pixel_Index = 0
            Row_Data = []
            
            for i in range(0, len(pixel_data_ascii), 4):
                # Get 4-character hex value for current pixel
                Pixel_Hex = pixel_data_ascii[i:i+4]
                Row_Data.append(Pixel_Hex)
                Pixel_Index += 1
                
                # When we've collected a full row worth of pixels, write to CSV
                if Pixel_Index % self.Width == 0:
                    Writer.writerow(Row_Data)
                    Row_Data = []
                    
        print(f"ASCII data saved to {csv_filename} ({self.Height} rows × {self.Width} columns)")

    def Close(self):
        """Close serial connection"""
        self.Ser.close()

# Usage example
if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python buffer_reader.py <serial_port> [output_filename]")
        print("Example: python buffer_reader.py COM3 screen_capture.png")
        sys.exit(1)
        
    port = sys.argv[1]
    image_filename = sys.argv[2] if len(sys.argv) > 2 else f"buffer_capture_{int(time.time())}.png"
    csv_filename = sys.argv[3] if len(sys.argv) > 3 else f"pixel_data_{int(time.time())}.csv"
    
    try:
        reader = RP2350BufferReader(port)
        image = reader.Read_Display_Buffer(image_filename, csv_filename)
        reader.Close()
        print("Buffer read complete!")
        
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)