// MTBoss: MT6261 firmware management application
// with zero DA required
// Created by Luxferre in 2026, released into the public domain

package main

import (
  "bytes"
  "crypto/md5"
  "encoding/binary"
  "encoding/hex"
  "flag"
  "fmt"
  "os"
  "strings"
  "time"

  "go.bug.st/serial"
)

// ==================== REGISTERS & BITS ====================
const (
  SFI_BASE       = 0xA0140000
  SFI_MAC_CTL    = SFI_BASE + 0x0000
  SFI_DIRECT_CTL = SFI_BASE + 0x0004
  SFI_MISC_CTL   = SFI_BASE + 0x0008
  SFI_MAC_OUTL   = SFI_BASE + 0x0010
  SFI_MAC_INL    = SFI_BASE + 0x0014
  SFI_MISC_CTL3  = SFI_BASE + 0x0044
  SFI_GPRAM      = SFI_BASE + 0x0800

  SFI_WIP        = 1 << 0
  SFI_WIP_READY  = 1 << 1
  SFI_TRIG       = 1 << 2
  SFI_MAC_EN     = 1 << 3
  SFI_MAC_SEL    = 1 << 28

  BOOT_ENG_BASE  = 0xA0510000

  FLASH_SIZE_4MB  = 0x400000 // 4MB (4,194,304 bytes)
  FLASH_PAGE_SIZE = 0x100    // 256 bytes
  FLASH_SECTOR_SZ = 0x1000   // 4KB

  SR_WIP = 0x01
  SR_WEL = 0x02
  SR_BP0 = 0x04
  SR_BP1 = 0x08
  SR_BP2 = 0x10
  SR_BP3 = 0x20
)

// ==================== HELPER PARSER ====================
func parseSizeOrOffset(valStr string) (uint64, error) {
  valStr = strings.TrimSpace(valStr)
  if valStr == "" {
    return 0, fmt.Errorf("empty string")
  }

  multiplier := uint64(1)
  upper := strings.ToUpper(valStr)

  if strings.HasSuffix(upper, "MIB") || strings.HasSuffix(upper, "MB") {
    multiplier = 1024 * 1024
    valStr = valStr[:len(valStr)-2]
    if strings.HasSuffix(strings.ToUpper(valStr), "I") {
      valStr = valStr[:len(valStr)-1]
    }
  } else if strings.HasSuffix(upper, "M") {
    multiplier = 1024 * 1024
    valStr = valStr[:len(valStr)-1]
  } else if strings.HasSuffix(upper, "KIB") || strings.HasSuffix(upper, "KB") {
    multiplier = 1024
    valStr = valStr[:len(valStr)-2]
    if strings.HasSuffix(strings.ToUpper(valStr), "I") {
      valStr = valStr[:len(valStr)-1]
    }
  } else if strings.HasSuffix(upper, "K") {
    multiplier = 1024
    valStr = valStr[:len(valStr)-1]
  }

  valStr = strings.TrimSpace(valStr)
  var val uint64
  var err error
  if strings.HasPrefix(valStr, "0x") || strings.HasPrefix(valStr, "0X") {
    _, err = fmt.Sscanf(valStr, "0x%x", &val)
  } else {
    _, err = fmt.Sscanf(valStr, "%d", &val)
  }
  if err != nil {
    return 0, err
  }
  return val * multiplier, nil
}

// ==================== MTBOSS STRUCT ====================
type MTBoss struct {
  port     serial.Port
  portName string
  chipID   uint16
}

func NewMTBoss(portName string) (*MTBoss, error) {
  flasher := &MTBoss{portName: portName}
  if err := flasher.connectPort(); err != nil {
    return nil, err
  }
  return flasher, nil
}

func (f *MTBoss) connectPort() error {
  var targetPort string
  if f.portName == "" || f.portName == "auto" {
    ports, err := serial.GetPortsList()
    if err == nil {
      for _, p := range ports {
        if strings.Contains(p, "ttyUSB") || strings.Contains(p, "ttyACM") || strings.Contains(p, "COM") {
          targetPort = p
          break
        }
      }
    }
    if targetPort == "" {
      targetPort = "/dev/ttyUSB0"
    }
  } else {
    targetPort = f.portName
  }

  fmt.Printf("Connecting to serial port %s at 115200 baud (RTS/CTS enabled)...\n", targetPort)
  mode := &serial.Mode{
    BaudRate: 115200,
    DataBits: 8,
    Parity:   serial.NoParity,
    StopBits: serial.OneStopBit,
  }

  start := time.Now()
  for time.Since(start) < 10*time.Minute {
    p, err := serial.Open(targetPort, mode)
    if err == nil {
      // Enable RTSCTS hardware flow control
      _ = p.SetRTS(true)
      _ = p.SetDTR(true)
      f.port = p
      f.portName = targetPort
      return nil
    }
    // Also scan for any new port if set to auto
    if f.portName == "auto" {
      ports, _ := serial.GetPortsList()
      for _, pName := range ports {
        if strings.Contains(pName, "ttyUSB") || strings.Contains(pName, "ttyACM") {
          p, err := serial.Open(pName, mode)
          if err == nil {
            _ = p.SetRTS(true)
            _ = p.SetDTR(true)
            f.port = p
            f.portName = pName
            return nil
          }
        }
      }
    }
    time.Sleep(50 * time.Millisecond)
  }
  return fmt.Errorf("timeout waiting for serial port %s", targetPort)
}

func (f *MTBoss) Close() {
  if f.port != nil {
    _ = f.port.Close()
  }
}

// ==================== BASIC SERIAL I/O ====================
func (f *MTBoss) sendRaw(data []byte) error {
  _, err := f.port.Write(data)
  return err
}

func (f *MTBoss) readExact(length int) ([]byte, error) {
  buf := make([]byte, length)
  read := 0
  for read < length {
    n, err := f.port.Read(buf[read:])
    if err != nil {
      return nil, err
    }
    if n == 0 {
      time.Sleep(2 * time.Millisecond)
    }
    read += n
  }
  return buf, nil
}

// ==================== BROM REGISTER OPERATIONS ====================
func (f *MTBoss) readReg16(addr uint32) (uint16, error) {
  cmd := make([]byte, 9)
  cmd[0] = 0xA2
  binary.BigEndian.PutUint32(cmd[1:5], addr)
  binary.BigEndian.PutUint32(cmd[5:9], 1)

  if err := f.sendRaw(cmd); err != nil {
    return 0, err
  }
  if _, err := f.readExact(9); err != nil { // Echo
    return 0, err
  }
  resp, err := f.readExact(2)
  if err != nil {
    return 0, err
  }
  return binary.BigEndian.Uint16(resp), nil
}

func (f *MTBoss) writeReg16(addr uint32, val uint16) error {
  cmd := make([]byte, 9)
  cmd[0] = 0xD2
  binary.BigEndian.PutUint32(cmd[1:5], addr)
  binary.BigEndian.PutUint32(cmd[5:9], 1)

  if err := f.sendRaw(cmd); err != nil {
    return err
  }
  if _, err := f.readExact(9); err != nil {
    return err
  }
  if ack, err := f.readExact(2); err != nil || !bytes.Equal(ack, []byte{0x00, 0x01}) {
    return fmt.Errorf("writeReg16 cmd ACK error: %v", ack)
  }

  valBuf := make([]byte, 2)
  binary.BigEndian.PutUint16(valBuf, val)
  if err := f.sendRaw(valBuf); err != nil {
    return err
  }
  if _, err := f.readExact(2); err != nil {
    return err
  }
  if ack, err := f.readExact(2); err != nil || !bytes.Equal(ack, []byte{0x00, 0x01}) {
    return fmt.Errorf("writeReg16 data ACK error: %v", ack)
  }
  return nil
}

func (f *MTBoss) readReg32(addr uint32) (uint32, error) {
  cmd := make([]byte, 9)
  cmd[0] = 0xD1
  binary.BigEndian.PutUint32(cmd[1:5], addr)
  binary.BigEndian.PutUint32(cmd[5:9], 1)

  if err := f.sendRaw(cmd); err != nil {
    return 0, err
  }
  if _, err := f.readExact(9); err != nil {
    return 0, err
  }
  resp, err := f.readExact(8) // status[2] + data[4] + status[2]
  if err != nil {
    return 0, err
  }
  return binary.LittleEndian.Uint32(resp[2:6]), nil
}

func (f *MTBoss) writeReg32(addr uint32, val uint32) error {
  cmd := make([]byte, 9)
  cmd[0] = 0xD4
  binary.BigEndian.PutUint32(cmd[1:5], addr)
  binary.BigEndian.PutUint32(cmd[5:9], 1)

  if err := f.sendRaw(cmd); err != nil {
    return err
  }
  if _, err := f.readExact(9); err != nil {
    return err
  }
  if ack, err := f.readExact(2); err != nil || !bytes.Equal(ack, []byte{0x00, 0x01}) {
    return fmt.Errorf("writeReg32 cmd ACK error: %v", ack)
  }

  valBuf := make([]byte, 4)
  binary.BigEndian.PutUint32(valBuf, val)
  if err := f.sendRaw(valBuf); err != nil {
    return err
  }
  if _, err := f.readExact(4); err != nil {
    return err
  }
  if ack, err := f.readExact(2); err != nil || !bytes.Equal(ack, []byte{0x00, 0x01}) {
    return fmt.Errorf("writeReg32 data ACK error: %v", ack)
  }
  return nil
}

// ==================== BOOTROM HANDSHAKE ====================
func (f *MTBoss) WaitAndConnect(timeoutSec int) error {
  fmt.Println("\n============================================================")
  fmt.Println("Waiting for MTBoss (MT6261) BootROM Connection")
  fmt.Println("============================================================")
  fmt.Println("1. Power OFF phone completely")
  fmt.Println("2. Press and hold BOOT key (or Power button)")
  fmt.Println("3. Plug in USB cable now...")

  start := time.Now()
  dots := 0
  for time.Since(start) < time.Duration(timeoutSec)*time.Second {
    _ = f.sendRaw([]byte{0xA0})
    time.Sleep(5 * time.Millisecond)

    b := make([]byte, 1)
    n, _ := f.port.Read(b)
    if n > 0 && b[0] == 0x5F {
      // Second sync packet
      _ = f.sendRaw([]byte{0x0A, 0x50, 0x05})
      ack3, err := f.readExact(3)
      if err == nil && bytes.Equal(ack3, []byte{0xF5, 0xAF, 0xFA}) {
        fmt.Println("\nConnected to MT6261 BootROM!")

        chip, err := f.readReg16(0x80000008)
        if err == nil {
          f.chipID = chip
          fmt.Printf("Chip ID: 0x%04X\n", chip)
        }

        // Disable watchdogs & enable USB download
        _ = f.writeReg16(0xA0030000, 0x2200) // System watchdog
        _ = f.writeReg16(0xA0700A28, 0x8000) // USB download mode
        _ = f.writeReg16(0xA0700A24, 0x0002) // Battery watchdog
        _ = f.writeReg32(BOOT_ENG_BASE, 2)   // Boot engine memory map mode 2
        fmt.Println("BootROM initialized & watchdogs disabled.")
        return nil
      }
    }

    dots++
    if dots%20 == 0 {
      fmt.Print(".")
    }
  }
  return fmt.Errorf("timeout waiting for BootROM sync")
}

// ==================== SFI HARDWARE MAC MODE ====================
func (f *MTBoss) sfiMacCmdWrite(cmdByte byte, addr *uint32, data []byte) error {
  var payload []byte
  payload = append(payload, cmdByte)
  if addr != nil {
    addrBytes := make([]byte, 4)
    binary.BigEndian.PutUint32(addrBytes, *addr)
    payload = append(payload, addrBytes[1:]...) // 3-byte big-endian address
  }
  payload = append(payload, data...)

  totalLen := uint32(len(payload))

  // 1. Write payload to GPRAM (0xA0140800) in 32-bit LE words
  for i := uint32(0); i < totalLen; i += 4 {
    end := i + 4
    if end > totalLen {
      end = totalLen
    }
    chunk := make([]byte, 4)
    copy(chunk, payload[i:end])
    valLE := binary.LittleEndian.Uint32(chunk)
    if err := f.writeReg32(SFI_GPRAM+i, valLE); err != nil {
      return err
    }
  }

  // 2. Mask AHB Channel 2
  misc3, _ := f.readReg32(SFI_MISC_CTL3)
  _ = f.writeReg32(SFI_MISC_CTL3, misc3|(1<<9))

  // 3. Enable MAC Mode FIRST (SFI_MAC_SEL bit 28 | SFI_MAC_EN bit 3)
  macVal := uint32(SFI_MAC_SEL | SFI_MAC_EN)
  if err := f.writeReg32(SFI_MAC_CTL, macVal); err != nil {
    return err
  }

  // 4. Set OUTL and INL lengths while MAC mode is active
  _ = f.writeReg32(SFI_MAC_OUTL, totalLen)
  _ = f.writeReg32(SFI_MAC_INL, 0)

  // 5. Trigger transaction (SFI_TRIG bit 2)
  _ = f.writeReg32(SFI_MAC_CTL, macVal|SFI_TRIG)

  // 6. Poll for completion (WIP_READY bit 1 set, WIP bit 0 clear)
  for i := 0; i < 300; i++ {
    v, err := f.readReg32(SFI_MAC_CTL)
    if err == nil && (v&SFI_WIP_READY != 0) && (v&SFI_WIP == 0) {
      break
    }
    time.Sleep(1 * time.Millisecond)
  }

  // 7. Clean up MAC mode and unmask AHB
  _ = f.writeReg32(SFI_MAC_CTL, 0)
  _ = f.writeReg32(SFI_MISC_CTL3, misc3&^(1<<9))
  return nil
}

func (f *MTBoss) sfiMacCmdRead(cmdByte byte, readLen uint32) ([]byte, error) {
  // 1. Write command to GPRAM
  if err := f.writeReg32(SFI_GPRAM, uint32(cmdByte)); err != nil {
    return nil, err
  }

  // 2. Mask AHB Channel 2
  misc3, _ := f.readReg32(SFI_MISC_CTL3)
  _ = f.writeReg32(SFI_MISC_CTL3, misc3|(1<<9))

  // 3. Enable MAC Mode FIRST
  macVal := uint32(SFI_MAC_SEL | SFI_MAC_EN)
  _ = f.writeReg32(SFI_MAC_CTL, macVal)

  // 4. Set OUTL and INL lengths
  _ = f.writeReg32(SFI_MAC_OUTL, 1)
  _ = f.writeReg32(SFI_MAC_INL, readLen)

  // 5. Trigger
  _ = f.writeReg32(SFI_MAC_CTL, macVal|SFI_TRIG)

  // 6. Poll completion
  for i := 0; i < 300; i++ {
    v, err := f.readReg32(SFI_MAC_CTL)
    if err == nil && (v&SFI_WIP_READY != 0) && (v&SFI_WIP == 0) {
      break
    }
    time.Sleep(1 * time.Millisecond)
  }

  // 7. Read response from GPRAM
  totalBytes := 1 + readLen
  var resBuf bytes.Buffer
  for i := uint32(0); i < (totalBytes + 3); i += 4 {
    val, _ := f.readReg32(SFI_GPRAM + i)
    wordBytes := make([]byte, 4)
    binary.LittleEndian.PutUint32(wordBytes, val)
    resBuf.Write(wordBytes)
  }

  // 8. Clean up
  _ = f.writeReg32(SFI_MAC_CTL, 0)
  _ = f.writeReg32(SFI_MISC_CTL3, misc3&^(1<<9))

  fullData := resBuf.Bytes()
  if uint32(len(fullData)) < 1+readLen {
    return nil, fmt.Errorf("short SFI MAC read")
  }
  return fullData[1 : 1+readLen], nil
}

// ==================== SPI FLASH HIGH-LEVEL ====================
func (f *MTBoss) ReadJEDECID() ([]byte, error) {
  return f.sfiMacCmdRead(0x9F, 3)
}

func (f *MTBoss) ReadStatusRegister() (byte, error) {
  data, err := f.sfiMacCmdRead(0x05, 1)
  if err != nil || len(data) == 0 {
    return 0, err
  }
  return data[0], nil
}

func (f *MTBoss) WaitUntilReady(timeoutSec int) error {
  start := time.Now()
  for time.Since(start) < time.Duration(timeoutSec)*time.Second {
    sr, err := f.ReadStatusRegister()
    if err == nil && (sr&SR_WIP == 0) {
      return nil
    }
    time.Sleep(5 * time.Millisecond)
  }
  return fmt.Errorf("timeout waiting for SPI flash ready")
}

func (f *MTBoss) UnlockWriteProtection() error {
  sr, err := f.ReadStatusRegister()
  if err == nil {
    fmt.Printf("SPI Flash Status Register: 0x%02X\n", sr)
    if sr&(SR_BP0|SR_BP1|SR_BP2|SR_BP3) != 0 {
      fmt.Println("Write protection enabled - disabling...")
      _ = f.sfiMacCmdWrite(0x50, nil, nil) // Volatile WREN
      _ = f.sfiMacCmdWrite(0x01, nil, []byte{0x00})
      _ = f.WaitUntilReady(5)
      sr, _ = f.ReadStatusRegister()
      fmt.Printf("Status Register after unlocking: 0x%02X\n", sr)
    }
  }
  return nil
}

func (f *MTBoss) EnableWrite() error {
  return f.sfiMacCmdWrite(0x06, nil, nil)
}

func (f *MTBoss) EraseSector(addr uint32) error {
  if err := f.WaitUntilReady(5); err != nil {
    return err
  }
  if err := f.EnableWrite(); err != nil {
    return err
  }
  if err := f.sfiMacCmdWrite(0x20, &addr, nil); err != nil {
    return err
  }
  return f.WaitUntilReady(5)
}

func (f *MTBoss) PageProgram(addr uint32, data []byte) error {
  const maxChunk = 64
  offset := 0
  for offset < len(data) {
    end := offset + maxChunk
    if end > len(data) {
      end = len(data)
    }
    chunk := data[offset:end]
    currAddr := addr + uint32(offset)

    if err := f.WaitUntilReady(5); err != nil {
      return err
    }
    if err := f.EnableWrite(); err != nil {
      return err
    }
    if err := f.sfiMacCmdWrite(0x02, &currAddr, chunk); err != nil {
      return err
    }
    if err := f.WaitUntilReady(5); err != nil {
      return err
    }
    offset += len(chunk)
  }
  return nil
}

// ==================== BROM MEMORY-MAPPED READING ====================
func (f *MTBoss) ReadFlash(address uint32, length uint32) ([]byte, error) {
  var result bytes.Buffer
  remaining := length
  currAddr := address
  const blkSize = 1024

  for remaining > 0 {
    rsize := remaining
    if rsize > blkSize {
      rsize = blkSize
    }
    wordsCnt := rsize >> 2

    cmd := make([]byte, 9)
    cmd[0] = 0xD1
    binary.BigEndian.PutUint32(cmd[1:5], currAddr)
    binary.BigEndian.PutUint32(cmd[5:9], wordsCnt)

    if err := f.sendRaw(cmd); err != nil {
      return nil, err
    }
    if _, err := f.readExact(9); err != nil {
      return nil, err
    }

    respLen := (int(wordsCnt) * 4) + 4
    resp, err := f.readExact(respLen)
    if err != nil {
      return nil, err
    }

    // Skip 2-byte header and 2-byte status footer
    rawWords := resp[2 : 2+int(wordsCnt)*4]
    for i := 0; i < len(rawWords); i += 4 {
      wordVal := binary.LittleEndian.Uint32(rawWords[i : i+4])
      wordBE := make([]byte, 4)
      binary.BigEndian.PutUint32(wordBE, wordVal)
      result.Write(wordBE)
    }

    currAddr += rsize
    remaining -= rsize

    progress := float64(length-remaining) / float64(length) * 100.0
    fmt.Printf("\rReading: %.1f%%", progress)
  }
  fmt.Println("\nRead complete.")
  return result.Bytes(), nil
}

// ==================== HIGH-LEVEL FLASH OPERATIONS ====================
func (f *MTBoss) WriteRegion(address uint32, data []byte, verify bool) error {
  fmt.Printf("\n============================================================\n")
  fmt.Printf("Flashing Region at 0x%08X (Size: %d bytes / 0x%X)\n", address, len(data), len(data))
  fmt.Printf("============================================================\n")

  _ = f.UnlockWriteProtection()

  startSector := address & ^uint32(FLASH_SECTOR_SZ-1)
  endAddr := address + uint32(len(data))

  fmt.Printf("Erasing sectors from 0x%08X to 0x%08X...\n", startSector, endAddr)
  for sec := startSector; sec < endAddr; sec += FLASH_SECTOR_SZ {
    fmt.Printf("  Erasing sector at 0x%08X...\n", sec)
    if err := f.EraseSector(sec); err != nil {
      return fmt.Errorf("failed erasing sector 0x%08X: %v", sec, err)
    }
  }
  fmt.Println("Sectors erased.")

  fmt.Printf("Programming %d bytes to 0x%08X...\n", len(data), address)
  offset := 0
  for offset < len(data) {
    end := offset + 256
    if end > len(data) {
      end = len(data)
    }
    chunk := data[offset:end]
    currAddr := address + uint32(offset)
    if err := f.PageProgram(currAddr, chunk); err != nil {
      return fmt.Errorf("failed programming at 0x%08X: %v", currAddr, err)
    }
    offset += len(chunk)
    fmt.Printf("\rProgramming: %.1f%%", float64(offset)/float64(len(data))*100.0)
  }
  fmt.Println("\nProgramming complete.")

  if verify {
    fmt.Printf("Verifying %d bytes at 0x%08X...\n", len(data), address)
    readData, err := f.ReadFlash(address, uint32(len(data)))
    if err != nil {
      return fmt.Errorf("verification read failed: %v", err)
    }

    refMD5 := md5.Sum(data)
    readMD5 := md5.Sum(readData)
    fmt.Printf("Reference MD5: %s\n", hex.EncodeToString(refMD5[:]))
    fmt.Printf("Readback  MD5: %s\n", hex.EncodeToString(readMD5[:]))

    if bytes.Equal(data, readData) {
      fmt.Println("\n*** VERIFICATION PASSED - EXACT BYTE-FOR-BYTE MATCH ***")
    } else {
      return fmt.Errorf("verification FAILED - binary mismatch!")
    }
  }
  return nil
}

func (f *MTBoss) Reset() {
  fmt.Println("Resetting device...")
  _ = f.writeReg16(0xA003001C, 0x1209)
}

// ==================== MAIN CLI ====================
func main() {
  portFlag := flag.String("port", "auto", "Serial port device (e.g. /dev/ttyUSB0, /dev/ttyUSB1 or auto)")
  modeFlag := flag.String("mode", "flash", "Mode of operation: flash, read, erase, identify")
  fileFlag := flag.String("file", "", "Input/Output binary file path")
  startFlag := flag.String("start", "", "Flash start offset address (e.g. 0, 0x0000, 4K)")
  offsetFlag := flag.String("offset", "0x00000000", "Flash offset address (alias for -start, default 0)")
  sizeFlag := flag.String("size", "", "Operation size in bytes/hex/units (e.g. 4MB, 64KB, 0x400000)")
  lengthFlag := flag.String("length", "0x400000", "Operation size (alias for -size, default 4MiB / 0x400000)")
  verifyFlag := flag.Bool("verify", true, "Verify flash write via readback")
  timeoutFlag := flag.Int("timeout", 600, "Timeout in seconds waiting for BootROM sync")

  flag.Parse()

  // Determine start address string (prefer -start if explicitly set)
  startStr := *offsetFlag
  if *startFlag != "" {
    startStr = *startFlag
  }

  // Determine size string (prefer -size if explicitly set)
  sizeStr := *lengthFlag
  if *sizeFlag != "" {
    sizeStr = *sizeFlag
  }

  start, err := parseSizeOrOffset(startStr)
  if err != nil {
    fmt.Printf("Invalid start/offset format '%s': %v\n", startStr, err)
    os.Exit(1)
  }

  size, err := parseSizeOrOffset(sizeStr)
  if err != nil {
    fmt.Printf("Invalid size/length format '%s': %v\n", sizeStr, err)
    os.Exit(1)
  }

  flasher, err := NewMTBoss(*portFlag)
  if err != nil {
    fmt.Printf("Error opening serial port: %v\n", err)
    os.Exit(1)
  }
  defer flasher.Close()

  if err := flasher.WaitAndConnect(*timeoutFlag); err != nil {
    fmt.Printf("BootROM Connection failed: %v\n", err)
    os.Exit(1)
  }

  switch *modeFlag {
  case "identify":
    id, err := flasher.ReadJEDECID()
    if err == nil && len(id) >= 3 {
      fmt.Printf("JEDEC ID: %02X %02X %02X\n", id[0], id[1], id[2])
    } else {
      fmt.Printf("Could not read JEDEC ID: %v\n", err)
    }
    flasher.Reset()

  case "read":
    if *fileFlag == "" {
      fmt.Println("Error: --file is required for read mode")
      os.Exit(1)
    }
    fmt.Printf("Reading %d (0x%X) bytes from 0x%08X to %s...\n", size, size, uint32(start), *fileFlag)
    data, err := flasher.ReadFlash(uint32(start), uint32(size))
    if err != nil {
      fmt.Printf("Read failed: %v\n", err)
      os.Exit(1)
    }
    if err := os.WriteFile(*fileFlag, data, 0644); err != nil {
      fmt.Printf("Failed writing file: %v\n", err)
      os.Exit(1)
    }
    hash := md5.Sum(data)
    fmt.Printf("Successfully saved %s (MD5: %s)\n", *fileFlag, hex.EncodeToString(hash[:]))
    flasher.Reset()

  case "erase":
    fmt.Printf("Erasing %d (0x%X) bytes at 0x%08X...\n", size, size, uint32(start))
    startSec := uint32(start) & ^uint32(FLASH_SECTOR_SZ-1)
    endSec := uint32(start + size)
    for sec := startSec; sec < endSec; sec += FLASH_SECTOR_SZ {
      fmt.Printf("Erasing sector at 0x%08X...\n", sec)
      if err := flasher.EraseSector(sec); err != nil {
        fmt.Printf("Erase failed: %v\n", err)
        os.Exit(1)
      }
    }
    fmt.Println("Erase complete.")
    flasher.Reset()

  case "flash":
    if *fileFlag == "" {
      fmt.Println("Error: --file is required for flash mode")
      os.Exit(1)
    }
    data, err := os.ReadFile(*fileFlag)
    if err != nil {
      fmt.Printf("Failed reading file %s: %v\n", *fileFlag, err)
      os.Exit(1)
    }

    // If size was explicitly set or if size < len(data), limit payload to specified size
    if uint64(len(data)) > size {
      fmt.Printf("Limiting file payload from %d to specified size %d (0x%X) bytes\n", len(data), size, size)
      data = data[:size]
    }

    fHash := md5.Sum(data)
    fmt.Printf("Loaded %s (%d bytes, MD5: %s)\n", *fileFlag, len(data), hex.EncodeToString(fHash[:]))

    if err := flasher.WriteRegion(uint32(start), data, *verifyFlag); err != nil {
      fmt.Printf("Flashing failed: %v\n", err)
      os.Exit(1)
    }
    fmt.Println("\n============================================================")
    fmt.Println("Operation completed successfully!")
    fmt.Println("============================================================")
    flasher.Reset()

  default:
    fmt.Printf("Unknown mode: %s\n", *modeFlag)
    os.Exit(1)
  }
}