#include "ff.h"
#include "stdio.h"
#include "string.h"
#include "main.h" // Ensure huart3 is defined here

//#define NEED_FORMAT_SD
#ifdef NEED_FORMAT_SD

extern UART_HandleTypeDef huart3;

// Configuration
#define TEST_FILE_NAME    "ST_DATA.BIN"
#define TEST_BUF_SIZE     (64 * 1024)       // 64KB Buffer (Good for SDIO DMA)
#define TOTAL_TEST_SIZE   (16 * 1024 * 1024) // 16MB Total Data

// Align buffer to 4 bytes for DMA compatibility
uint8_t g_test_buffer[TEST_BUF_SIZE] __attribute__((aligned(4)));
char uart_buf[128]; // Buffer for formatting UART messages

// Helper function to print via UART3
void UART3_Print(const char* str) {
    HAL_UART_Transmit(&huart3, (uint8_t*)str, strlen(str), 100);
}

void Run_SDNAND_SpeedTest(void) {
    FIL file;
    FRESULT res;
    UINT bw, br;
    uint32_t startTime, endTime;
    float speed;

    UART3_Print("\r\n=== SD NAND Speed Test Start ===\r\n");
    
    // Initialize buffer with dummy data
    for (uint32_t i = 0; i < TEST_BUF_SIZE; i++) g_test_buffer[i] = (uint8_t)(i % 256);

    // --- WRITE TEST ---
    res = f_open(&file, TEST_FILE_NAME, FA_CREATE_ALWAYS | FA_WRITE);
    if (res != FR_OK) {
        sprintf(uart_buf, "Open Fail (Write). Error: %d\r\n", res);
        UART3_Print(uart_buf);
        return;
    }

    UART3_Print("Testing Write Speed... Please wait.\r\n");
    startTime = HAL_GetTick();
    for (uint32_t i = 0; i < TOTAL_TEST_SIZE / TEST_BUF_SIZE; i++) {
        res = f_write(&file, g_test_buffer, TEST_BUF_SIZE, &bw);
        if (res != FR_OK) break;
    }
    f_sync(&file); // Flush to physical NAND
    endTime = HAL_GetTick();

    speed = ((float)TOTAL_TEST_SIZE / 1024.0 / 1024.0) / ((float)(endTime - startTime) / 1000.0);
    sprintf(uart_buf, "WRITE: %.2f MB/s (%lu ms)\r\n", speed, endTime - startTime);
    UART3_Print(uart_buf);
    f_close(&file);

    // --- READ TEST ---
    res = f_open(&file, TEST_FILE_NAME, FA_READ);
    if (res != FR_OK) {
        UART3_Print("Open Fail (Read)\r\n");
        return;
    }

    UART3_Print("Testing Read Speed... Please wait.\r\n");
    startTime = HAL_GetTick();
    for (uint32_t i = 0; i < TOTAL_TEST_SIZE / TEST_BUF_SIZE; i++) {
        res = f_read(&file, g_test_buffer, TEST_BUF_SIZE, &br);
        if (res != FR_OK) break;
    }
    endTime = HAL_GetTick();

    speed = ((float)TOTAL_TEST_SIZE / 1024.0 / 1024.0) / ((float)(endTime - startTime) / 1000.0);
    sprintf(uart_buf, "READ:  %.2f MB/s (%lu ms)\r\n", speed, endTime - startTime);
    UART3_Print(uart_buf);

    f_close(&file);
    f_unlink(TEST_FILE_NAME); // Clean up
    UART3_Print("=== Test Finished ===\r\n");
}

extern SD_HandleTypeDef hsd;
extern DMA_HandleTypeDef hdma_sdio_rx;
extern DMA_HandleTypeDef hdma_sdio_tx;

void Raw_Hardware_Test(void) {
    uint32_t start, end;
    HAL_StatusTypeDef status;

    UART3_Print("\r\n--- Starting Raw Hardware Test (No File System) ---\r\n");

    // Test Raw DMA Write
    start = HAL_GetTick();
    status = HAL_SD_WriteBlocks_DMA(&hsd, g_test_buffer, 1000, TEST_BUF_SIZE / 512);
    if (status == HAL_OK) {
        // Wait for DMA to finish
        while (HAL_SD_GetCardState(&hsd) != HAL_SD_CARD_TRANSFER) {}
        end = HAL_GetTick();
        sprintf(uart_buf, "Raw DMA Write: %lu ms\r\n", end - start);
        UART3_Print(uart_buf);
    } else {
        UART3_Print("Raw Write Failed!\r\n");
    }
}

void Run_SDNAND_SpeedTest_V2(void) {
    FIL file;
    FRESULT res;
    UINT bw, br;
    uint32_t startTime, endTime, duration;
    float speed;

    UART3_Print("\r\n=== SD NAND Speed Test V2 ===\r\n");

    // 预填充数据
    for (uint32_t i = 0; i < TEST_BUF_SIZE; i++) g_test_buffer[i] = (uint8_t)(i % 256);

    // --- WRITE TEST ---
    res = f_open(&file, TEST_FILE_NAME, FA_CREATE_ALWAYS | FA_WRITE);
    if (res != FR_OK) {
        sprintf(uart_buf, "Open Fail (Write). Error: %d\r\n", res);
        UART3_Print(uart_buf);
        return;
    }

    UART3_Print("Writing 16MB...\r\n");
    startTime = HAL_GetTick();

    for (uint32_t i = 0; i < TOTAL_TEST_SIZE / TEST_BUF_SIZE; i++) {
        res = f_write(&file, g_test_buffer, TEST_BUF_SIZE, &bw);
        if (res != FR_OK || bw != TEST_BUF_SIZE) {
            sprintf(uart_buf, "Write Fail at block %lu, Res: %d\r\n", i, res);
            UART3_Print(uart_buf);
            f_close(&file);
            return;
        }
    }

    f_close(&file); // 必须先关闭文件，确保 FAT 目录项更新
    endTime = HAL_GetTick();
    duration = (endTime - startTime > 0) ? (endTime - startTime) : 1; // 防止除以0

    speed = ((float)TOTAL_TEST_SIZE / 1024.0 / 1024.0) / ((float)duration / 1000.0);
    sprintf(uart_buf, "WRITE Result: %.2f MB/s (%lu ms)\r\n", speed, duration);
    UART3_Print(uart_buf);

    // --- READ TEST ---
    HAL_Delay(100); // 稍微等待文件系统稳定
    res = f_open(&file, TEST_FILE_NAME, FA_READ);
    if (res != FR_OK) {
        sprintf(uart_buf, "Open Fail (Read). Error: %d\r\n", res);
        UART3_Print(uart_buf);
        return;
    }

    UART3_Print("Reading 16MB...\r\n");
    startTime = HAL_GetTick();
    for (uint32_t i = 0; i < TOTAL_TEST_SIZE / TEST_BUF_SIZE; i++) {
        res = f_read(&file, g_test_buffer, TEST_BUF_SIZE, &br);
        if (res != FR_OK || br != TEST_BUF_SIZE) {
            sprintf(uart_buf, "Read Fail at block %lu, Res: %d\r\n", i, res);
            UART3_Print(uart_buf);
            break;
        }
    }
    endTime = HAL_GetTick();
    duration = (endTime - startTime > 0) ? (endTime - startTime) : 1;

    speed = ((float)TOTAL_TEST_SIZE / 1024.0 / 1024.0) / ((float)duration / 1000.0);
    sprintf(uart_buf, "READ Result:  %.2f MB/s (%lu ms)\r\n", speed, duration);
    UART3_Print(uart_buf);

    f_close(&file);
    UART3_Print("=== Test Finished ===\r\n");
}

void SDNAND_ForceFormat_and_Mount(void) {
    FATFS fs;
    FRESULT res;
    BYTE work[_MAX_SS];
    char msg[128];

    UART3_Print("\r\n--- Initializing SD NAND for Optimization ---\r\n");

    // 1. 先尝试挂载 (立即挂载模式)
    res = f_mount(&fs, "", 1);
    if (res != FR_OK && res != FR_NO_FILESYSTEM) {
        sprintf(msg, "Mount failed before format: %d\r\n", res);
        UART3_Print(msg);
        return;
    }

    // 2. 执行格式化 (强制 512 簇)
    UART3_Print("Formatting... this may take a few seconds.\r\n");
    res = f_mkfs("", FM_FAT32, 32768, work, sizeof(work));
    if (res != FR_OK) {
        sprintf(msg, "Format failed: %d\r\n", res);
        UART3_Print(msg);
        return;
    }
    UART3_Print("Format completed successfully.\r\n");

    // 3. 卸载以清除旧缓存
    f_mount(NULL, "", 0);

    // 4. 重新挂载
    res = f_mount(&fs, "", 1);
    if (res == FR_OK) {
        UART3_Print("SD NAND Remounted with 4KB cluster size.\r\n");
    } else {
        sprintf(msg, "Final mount failed: %d\r\n", res);
        UART3_Print(msg);
    }
}

#endif