import serial
import time

# --- Cấu hình ---
PORT = 'COM6'
BAUD_RATE = 115200
TIMEOUT = 2
MAX_ADDR = 0x100000 
STEP_SIZE = 64
OUTPUT_FILE = "W25Q80_Main_v1.bin"

def dump_flash_v3_3():
    try:
        ser = serial.Serial(PORT, BAUD_RATE, timeout=TIMEOUT)
        print(f"Kết nối {PORT}...")
        time.sleep(2)
        ser.reset_input_buffer()
        
        current_addr = 0
        print(f"--- Bắt đầu Dump v1 (Dựa trên Header Header) ---")
        
        with open(OUTPUT_FILE, "wb") as f:
            while current_addr < MAX_ADDR:
                # 1. Gửi lệnh
                cmd = f"mem r {current_addr}\r\n"
                ser.write(cmd.encode())
                
                byte_list = []
                header_found = False
                start_t = time.time()
                
                # 2. Đọc phản hồi
                while len(byte_list) < STEP_SIZE and (time.time() - start_t < 2.5):
                    line = ser.readline().decode('utf-8', errors='ignore').strip()
                    
                    if not line:
                        continue
                    
                    # Nếu thấy dòng Header, đánh dấu bắt đầu lấy dữ liệu từ dòng SAU
                    if "Read flash address" in line:
                        header_found = True
                        continue 
                    
                    # Chỉ xử lý các dòng nằm SAU dòng Header
                    if header_found:
                        # Tách các chuỗi byte (ví dụ: "FF", "08", "03"...)
                        parts = line.split()
                        for p in parts:
                            # Kiểm tra xem có phải là byte hex hợp lệ (2 ký tự)
                            if len(p) == 2:
                                try:
                                    val = int(p, 16)
                                    byte_list.append(val)
                                except ValueError:
                                    continue # Bỏ qua nếu gặp "OK" hoặc rác
                
                # 3. Kiểm tra và ghi file
                if len(byte_list) >= STEP_SIZE:
                    # Lấy đúng 64 byte đầu tiên tìm thấy sau Header
                    f.write(bytearray(byte_list[:STEP_SIZE]))
                    current_addr += STEP_SIZE
                    
                    # In tiến độ
                    if current_addr % 512 == 0:
                        percent = (current_addr / MAX_ADDR) * 100
                        print(f"\rTiến độ: {percent:6.2f}% | Addr: {current_addr}", end="")
                else:
                    print(f"\n[Lỗi] Không gom đủ 64 byte tại {current_addr} (Chỉ có {len(byte_list)}B). Thử lại...")
                    ser.reset_input_buffer()
                    time.sleep(0.5)
                    # Không tăng current_addr để vòng lặp while thử lại địa chỉ này

        print(f"\n\n--- HOÀN THÀNH ---")
        ser.close()

    except Exception as e:
        print(f"\nLỗi hệ thống: {e}")

if __name__ == "__main__":
    dump_flash_v3_3()