import serial
import sys
import time

def capture_flash_dump(port, baudrate, dump_file):
    """
    Captures flash information and a binary dump from the ESP serial output.
    
    Parameters:
    - port: The serial port connected to the ESP (e.g., "/dev/ttyUSB0" or "COM3").
    - baudrate: The baud rate of the ESP's serial output (e.g., 115200).
    - dump_file: The name of the file to save the flash dump to.
    """
    try:
        flash_capacity = 0

        with serial.Serial(port, baudrate, timeout=1) as ser:
            print(f"Waiting for information start marker ('START_OF_INFORM') on {port}...")
            while True:
                line = ser.readline().decode('utf-8', errors='replace').strip()
                if "START_OF_INFORM" in line:
                    print("Found information start marker. Printing flash info...")
                    break
            
            # Wait for ESP to announce flash information
            while True:
                line = ser.readline().decode('utf-8', errors='replace').strip()
                if "END_OF_INFORM" in line:
                    print("Found information end marker. Flash info captured.")
                    break
                if "Capacity" in line:
                    try:
                        flash_capacity = int(line.split(":")[1].strip())  # Extract flash capacity
                    except (ValueError, IndexError):
                        print("Warning: Failed to parse flash capacity.")
                print(line)  # Print flash info to terminal

            if flash_capacity == 0:
                print("Warning: Flash capacity is unknown. Progress will not be calculated.")
            
            print(f"Waiting for dump start marker ('START_OF_DUMP')...")

            # Wait for ESP to announce start of dump
            while True:
                line = ser.read_until(b"START_OF_DUMP")
                if b"START_OF_DUMP" in line:
                    print("Found dump start marker. Beginning dump capture...")
                    break
            
            # Open output dump file
            with open(dump_file, "wb") as dump_output:
                total_bytes = 0
                last_report_time = time.time()
                
                # Read each block and handle them
                while True:
                    data = ser.read(256)  # Read in 256-byte chunks

                    if b"END_OF_DUMP" in data:
                        end_index = data.find(b"END_OF_DUMP")
                        dump_output.write(data[:end_index])  # Write up to the end marker
                        total_bytes += end_index
                        print(f"Found end marker. Dump complete. Total bytes: {total_bytes}")
                        break
                    elif b"READ_ERROR_AT_ADDRESS" in data:
                        error_index = data.find(b"READ_ERROR_AT_ADDRESS")
                        print(f"Read error detected: {data[error_index:].decode('utf-8', errors='replace').strip()}")
                        print("Terminating capture.")
                        break
                    elif b"ABORTED_DUMP_FAILURE" in data:
                        print("Error: ESP aborted the dump due to error when starting communication with flash. Terminating capture.")
                        break
                    elif data:
                        dump_output.write(data)
                        total_bytes += len(data)
                        # Print periodic status updates
                        if time.time() - last_report_time > 1:
                            if flash_capacity > 0:
                                percent_complete = (total_bytes / flash_capacity) * 100
                                print(f"Captured {total_bytes} bytes so far (~{percent_complete:.2f}% complete)...")
                            else:
                                print(f"Captured {total_bytes} bytes so far...")
                            last_report_time = time.time()
                    else:
                        print("Timeout: No data received. Retrying...")
            
            print(f"Flash dump saved to {dump_file}. Total size: {total_bytes} bytes.")
    except serial.SerialException as e:
        print(f"Serial port error: {e}")
    except IOError as e:
        print(f"File error: {e}")
    except KeyboardInterrupt:
        print("\nCapture interrupted by user.")

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python flash_dump_capture.py <port> <baudrate> <dump_file>")
        print("Example: python flash_dump_capture.py /dev/ttyUSB0 115200 flash_dump.bin")
        sys.exit(1)

    port = sys.argv[1]
    baudrate = int(sys.argv[2])
    dump_file = sys.argv[3]
    
    capture_flash_dump(port, baudrate, dump_file)
