#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stddef.h>
#include <stdarg.h>
#include <string.h>
#include <windows.h>
#include "scsi_structs.h"

#define STUB_SIZE       4080
#define FW_SIZE         524288

int version = 176;

FILE *log_fp = NULL;

void lprintf(const char *fmt, ...) {
    va_list args;
    va_start(args, fmt); vprintf(fmt, args); va_end(args);
    if (log_fp) {
        va_start(args, fmt); vfprintf(log_fp, fmt, args); va_end(args);
        fflush(log_fp);
    }
}

void hex_dump(const char* label, uint8_t* buf, uint32_t len) {
    lprintf("    [DEBUG] %s (%u bytes):\n      ", label, len);
    uint32_t limit = (len > 32) ? 32 : len;
    for (uint32_t i = 0; i < limit; i++) {
        lprintf("%02X ", buf[i]);
        if ((i + 1) % 16 == 0 && i < limit - 1) lprintf("\n      ");
    }
    if (len > 32) lprintf("... [truncated]");
    lprintf("\n");
}

int must_scsi(HANDLE fd, uint8_t* cdb, uint8_t cdb_len, uint8_t* data, uint32_t data_len, int data_in, const char* label, int ignore_errors) {
    SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER pkt = {0};
    DWORD b;
    pkt.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    pkt.Spt.CdbLength = cdb_len;
    pkt.Spt.DataIn = data_in;
    pkt.Spt.DataTransferLength = data_len;
    pkt.Spt.DataBuffer = data;
    pkt.Spt.TimeOutValue = 60; 
    pkt.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    pkt.Spt.SenseInfoLength = 32;
    memcpy(pkt.Spt.Cdb, cdb, cdb_len);

    hex_dump("CDB_SENT", cdb, cdb_len);
    if (data_len > 0 && data_in == SCSI_IOCTL_DATA_OUT) hex_dump("PAYLOAD", data, data_len);

    BOOL res = DeviceIoControl(fd, IOCTL_SCSI_PASS_THROUGH_DIRECT, &pkt, sizeof(pkt), &pkt, sizeof(pkt), &b, NULL);

    if (!res || pkt.Spt.ScsiStatus != 0) {
        // --- SURGICAL EXCEPTION: Hitachi Checksum Verification ---
        // If Status is 02 (Check Condition) but Sense is 00/00/00 (No Error)
        if (pkt.Spt.ScsiStatus == 0x02 && pkt.SenseBuf[2] == 0x00 && pkt.SenseBuf[12] == 0x00 && pkt.SenseBuf[13] == 0x00) {
            lprintf("    [ERROR] %s Unsuccessful: Checksum Failed (Status 02 | Sense 00/00/00)\n", label);
            exit(1);
        }

        if (ignore_errors) {
            lprintf("    [OK] %s (Ignored Expected Error)\n", label);
            return 0;
        } else {
            lprintf("    [ERROR] %s Status %02X | Sense %02X/%02X/%02X\n", label, pkt.Spt.ScsiStatus, pkt.SenseBuf[2]&0xF, pkt.SenseBuf[12], pkt.SenseBuf[13]);
            exit(1);
        }
    } else {
        lprintf("    [OK] %s Success.\n", label);
        return 1;
    }
}


// Strips illegal Windows filename characters and trims trailing spaces
void sanitize_string(char* str) {
    char* p = str;
    while (*p) {
        // Replace illegal path characters with underscores
        if (*p == '<' || *p == '>' || *p == ':' || *p == '"' || 
            *p == '/' || *p == '\\' || *p == '|' || *p == '?' || *p == '*') {
            *p = '_';
        }
        p++;
    }
    // Trim trailing spaces (very common in SCSI inquiry data)
    while (p > str && *(p - 1) == ' ') {
        *(p - 1) = '\0';
        p--;
    }
}

int main(int argc, char* argv[]) {
    lprintf("============================================================\n");
    lprintf("HITACHI GCC-4241N FLASHER (V%d)\n", version);
    lprintf("============================================================\n");

    if (argc < 4) {
        lprintf("Usage: flash_gcc.exe <DriveLetter> <4080-byte_stub.bin> <512kb_firmware.bin>\n");
        return -1;
    }

    log_fp = fopen("flasher_debug.log", "w");
    if (!log_fp) lprintf("[!] Could not open log file, continuing with console output only.\n");

    char devPath[20]; 
    sprintf(devPath, "\\\\.\\%c:", argv[1][0]);
    HANDLE hDrive = CreateFileA(devPath, GENERIC_READ|GENERIC_WRITE, FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
    if (hDrive == INVALID_HANDLE_VALUE) {
        DWORD error = GetLastError();
        char* message = NULL;
        FormatMessageA(
            FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
            NULL, error, 0, (LPSTR)&message, 0, NULL);
        
        lprintf("[!] Error opening drive handle: %s (Code: 0x%08X)\n", 
                message ? message : "Unknown error", error);
        LocalFree(message);
        
        if (error == ERROR_ACCESS_DENIED) {
            lprintf("    [Hint] Run as Administrator.\n");
        } else if (error == ERROR_SHARING_VIOLATION) {
            lprintf("    [Hint] Drive is in use by another process.\n");
        } else if (error == ERROR_FILE_NOT_FOUND) {
            lprintf("    [Hint] Invalid drive letter.\n");
        }
        return 1;
    }   

    lprintf("[*] Loading STORAGE_DEVICE_DESCRIPTOR...\n");
    char vendorStr[16] = "Generic";
    char productStr[21] = "Unknown";
    if (hDrive != INVALID_HANDLE_VALUE)
    {
        BYTE buffer[1024] = {0};
        STORAGE_PROPERTY_QUERY query = {0};
        query.PropertyId = StorageDeviceProperty;
        query.QueryType = PropertyStandardQuery;
        DWORD bytes;

        lprintf("Drive Info:\n");
        if (DeviceIoControl(hDrive, IOCTL_STORAGE_QUERY_PROPERTY, &query, sizeof(query),
                            buffer, sizeof(buffer), &bytes, NULL))
        {
            PSTORAGE_DEVICE_DESCRIPTOR desc = (PSTORAGE_DEVICE_DESCRIPTOR)buffer;
            printf("  Version: %lu\n", desc->Version);
            printf("  DeviceType: %u\n", desc->DeviceType);
            if (desc->VendorIdOffset) {
                printf("  Vendor ID: %s\n", (char*)desc + desc->VendorIdOffset);
                strcpy(vendorStr, (char *)(buffer + desc->VendorIdOffset));
                sanitize_string(vendorStr);
            }
            if (desc->ProductIdOffset) {
                printf("  Product ID: %s\n", (char*)desc + desc->ProductIdOffset);
                strcpy(productStr, (char *)(buffer + desc->ProductIdOffset));
                sanitize_string(productStr);
            }
            if (desc->ProductRevisionOffset) {
                printf("  Product Revision: %s\n", (char*)desc + desc->ProductRevisionOffset);
            }
            if (desc->SerialNumberOffset) {
                printf("  Serial Number: %s\n", (char*)desc + desc->SerialNumberOffset);
            }
        } else {
            DWORD error = GetLastError();
            LPSTR message = NULL;
            FormatMessageA(
                FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                NULL, error, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                (LPSTR)&message, 0, NULL);
            printf("[!] DeviceIoControl failed: %s (%lu)\n", message ? message : "Unknown error", error);
            if (message) LocalFree(message);
            CloseHandle(hDrive);
            if (log_fp) fclose(log_fp);
            return 1;
        }
    }
    if (strstr(vendorStr, "HL-DT-ST") == NULL) {
        lprintf("[!] Drive Incompatible.");
        CloseHandle(hDrive);
        if (log_fp) fclose(log_fp);
        return 1;
    }

    // Load prep stub (Headless 4080-bytes)
    lprintf("[*] Loading stub...\n");
    FILE *fs = fopen(argv[2], "rb"); 
    if (!fs) {
        printf("[-] Failed to open stub file: %s\n", argv[2]);
        return -1;
    }
    uint8_t *stub = VirtualAlloc(NULL, STUB_SIZE, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); 
    if (!stub) {
        printf("[-] Failed to allocate memory for stub\n");
        fclose(fs);
        return -1;
    }
    size_t stub_read = fread(stub, 1, STUB_SIZE, fs); 
    fclose(fs);
    if (stub_read != STUB_SIZE) {
        printf("[-] Failed to read full stub: expected %u, got %zu bytes\n", STUB_SIZE, stub_read);
        return -1;
    }
    lprintf("[+] Stub Checksum: 0x%02X 0x%02X\n", stub[STUB_SIZE - 2], stub[STUB_SIZE - 1]);

    // Load firmware (Headless 524288-bytes)
    lprintf("[*] Loading firmware...\n");
    FILE *ff = fopen(argv[3], "rb"); 
    if (!ff) {
        printf("[-] Failed to open firmware file: %s\n", argv[3]);
        return -1;
    }
    uint8_t *fw = VirtualAlloc(NULL, FW_SIZE, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); 
    if (!fw) {
        printf("[-] Failed to allocate memory for firmware\n");
        fclose(ff);
        return -1;
    }
    size_t fw_read = fread(fw, 1, FW_SIZE, ff); 
    fclose(ff);
    if (fw_read != FW_SIZE) {
        printf("[-] Failed to read full firmware: expected %u, got %zu bytes\n", FW_SIZE, fw_read);
        return -1;
    }
    lprintf("[+] Firware Checksum: 0x%02X 0x%02X\n", fw[FW_SIZE - 2], fw[FW_SIZE - 1]);

    // Stage 0: TUR (Will silently ignore error if tray is empty)
    lprintf("[>] EXECUTING: Stage0_TUR\n");
    must_scsi(hDrive, (uint8_t[]){0x00, 0, 0, 0, 0, 0}, 6, NULL, 0, SCSI_IOCTL_DATA_UNSPECIFIED, "Stage0_TUR", 1);

    // Stage 1: Unlock
    lprintf("[>] EXECUTING: Stage1_Unlock\n");
    must_scsi(hDrive, (uint8_t[]){0xE7, 0x48, 0x4C, 0x2D, 0x44, 0x54, 0x2D, 0x53, 0x54, 0, 0, 0}, 12, NULL, 0, SCSI_IOCTL_DATA_UNSPECIFIED, "Stage1_Unlock", 0);

    // Stage 2: Setup (Contract for 4080 bytes / 0x0FF0)
    lprintf("[>] EXECUTING: Stage2_Setup\n");
    uint8_t setup_payload[16] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0x06, 0x48, 0x4C, 0, 0, 0x0F, 0xF0};
    must_scsi(hDrive, (uint8_t[]){0x55, 0x10, 0, 0, 0, 0, 0, 0, 0x10, 0, 0, 0}, 12, setup_payload, 16, SCSI_IOCTL_DATA_OUT, "Stage2_Setup", 0);

    // The mandatory Hitachi Magic Header
    uint8_t c_hdr[8] = {0, 0, 0, 0, 0x48, 0x4C, 0, 0};
    uint8_t cdb_chunk[12] = {0x55, 0x10, 0, 0, 0, 0, 0, 0x08, 0x00, 0, 0, 0}; // Send 2048 bytes

    // Stage 3: Chunk 1 (Header + First 2040 bytes)
    lprintf("[>] EXECUTING: Stage3_Chunk1\n");
    uint8_t c1[2048]; 
    memcpy(c1, c_hdr, 8); 
    memcpy(c1 + 8, stub, 2040); 
    must_scsi(hDrive, cdb_chunk, 12, c1, 2048, SCSI_IOCTL_DATA_OUT, "Stage3_Chunk1", 0);

    // Stage 4: Chunk 2 (Header + Remaining 2040 bytes)
    lprintf("[>] EXECUTING: Stage4_Chunk2\n");
    uint8_t c2[2048]; 
    memcpy(c2, c_hdr, 8); 
    memcpy(c2 + 8, stub + 2040, 2040); 
    must_scsi(hDrive, cdb_chunk, 12, c2, 2048, SCSI_IOCTL_DATA_OUT, "Stage4_Chunk2", 0);

    VirtualFree(stub, 0, MEM_RELEASE);

    // --- Stage 5: Streaming ---
    lprintf("[*] STAGE 5: Streaming Firmware to SDRAM...\n");
    for (uint32_t i = 0x00; i < 256; i++) { // Changed from i < 0xFF to process all 256 blocks
        // Countdown will scale from 0xFF down to 0x00
        uint8_t countdown = 0xFF - i; 
        
        uint8_t fw_cdb[12] = {0x55, 0x10, countdown, 0x48, 0x4C, 0x07, 0x01, 0x08, 0x00, 0x00, 0x00, 0x00};
        
        must_scsi(hDrive, fw_cdb, 12, fw + (i * 2048), 2048, SCSI_IOCTL_DATA_OUT, "FW_Stream", 0);
        
        if (i % 32 == 0) lprintf("    [+] Sent %d/256 chunks...\n", i);

        Sleep(100);
    }
    VirtualFree(fw, 0, MEM_RELEASE);

    lprintf("\n[+] Firmware Upload Complete & Verified by Drive!\n");

    // Stage 6: Final Poll
    lprintf("[*] Polling Final State...\n");
    // Use VirtualAlloc to ensure the DataBuffer pointer is strictly aligned for SPTI Direct
    uint8_t *status = VirtualAlloc(NULL, 32, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    must_scsi(hDrive, (uint8_t[]){0x5A, 0x10, 0, 0x48, 0x4C, 0x08, 0, 0, 0x12, 0, 0, 0}, 12, status, 18, SCSI_IOCTL_DATA_IN, "FinalPoll", 0);
    VirtualFree(status, 0, MEM_RELEASE);

    // Stage 7: Ignition (Triggers injected MMU flip at 0x0859)
    lprintf("[*] FIRING IGNITION...\n");
    // Allocate a page-aligned buffer to catch the Mask ROM's final parting byte
    uint8_t *ig_status = VirtualAlloc(NULL, 32, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    if (must_scsi(hDrive, (uint8_t[]){0x5A, 0x10, 0, 0x48, 0x4C, 0x09, 0, 0, 0x12, 0, 0, 0}, 12, ig_status, 18, SCSI_IOCTL_DATA_IN, "Ignition", 0)) {
        lprintf("[+] Success. The drive should now be erasing and flashing the ROM.\n");
        lprintf("[*] DO NOT UNPLUG THE DRIVE!\n");
        lprintf("[*] The LED should turn OFF when the burn is complete.\n");
        lprintf("[*] Waiting 30 seconds for hardware write to complete...\n\n");

        for (int t = 30; t > 0; t--) {
            lprintf("\r    [TIMER] %02d seconds remaining...", t);
            Sleep(1000);
        }
        lprintf("\r    [TIMER] 00 seconds remaining... DONE.          \n");
        lprintf("\n[SUCCESS] The drive is now in an intentional hardware halt state.\n");
        lprintf("[!] Power cycle the drive to boot the new firmware.\n");
    }
    VirtualFree(ig_status, 0, MEM_RELEASE);
    free(stub); free(fw); 
    CloseHandle(hDrive);
    if (log_fp) fclose(log_fp);
    return 0;
}