// Copyright 2011-2023 XMOS LIMITED.
// This Software is subject to the terms of the XMOS Public Licence: Version 1.
#define DEBUG_PRINT_ENABLE 0
#include <xs1.h>
#include <flash.h>
#include <flashlib.h>
#include <string.h>
#include <xclib.h>
#include <stdio.h>
#include "debug_print.h"

#include "xua.h"

#if HID_DFU_EN

#ifndef DFU_FLASH_MAX_UPGRADE_SIZE
#define DFU_FLASH_MAX_UPGRADE_SIZE (2200 * 1024)
#endif

#define DFU_FLASH_ERROR()

static int dfu_flash_device_open = 0;
static fl_BootImageInfo dfu_factory_image;
static fl_BootImageInfo dfu_upgrade_image;

static int dfu_upgrade_image_valid = 0;
static int dfu_current_flash_subpage_index = 0;
static unsigned char dfu_current_flash_page_data[256];

/* Returns non-zero for error */
int dfu_flash_cmd_init(void)
{
    debug_printf("dfu_flash_cmd_init: dfu_flash_cmd_init\n");
    fl_BootImageInfo image;

    if (!dfu_flash_device_open)
    {
        if (dfu_flash_cmd_enable_ports())
            dfu_flash_device_open = 1;
        debug_printf("dfu_flash_cmd_init: dfu_flash_device_open = 1\n");
    }

    if (!dfu_flash_device_open)
    {
        debug_printf("dfu_flash_cmd_init: dfu_flash_device_open = 0\n");
        return 1;
    }

#if (!XUA_QUAD_SPI_FLASH)
    // Disable flash protection
    fl_setProtection(0);
#endif

    if (fl_getFactoryImage(&image) != 0)
    {
        return 1;
    }

    dfu_factory_image = image;

    if (fl_getNextBootImage(&image) == 0)
    {
        dfu_upgrade_image_valid = 1;
        dfu_upgrade_image = image;
        debug_printf("dfu_flash_cmd_init: dfu_upgrade_image_valid = 1\n");
    }

    debug_printf("dfu_flash_cmd_init: return 0\n");
     return 0;
}

int dfu_flash_cmd_deinit(void)
{
    if (!dfu_flash_device_open)
        return 0;

    dfu_flash_cmd_disable_ports();
    dfu_flash_device_open = 0;
    return 0;
}

int dfu_flash_cmd_read_page(unsigned char *data)
{
    if (!dfu_upgrade_image_valid)
    {
        *(unsigned int *)data = 1;
        return 4;
    }

    if (*(unsigned int *)data == 0)
    {
        fl_startImageRead(&dfu_upgrade_image);
    }

    dfu_current_flash_subpage_index = 0;

    if (fl_readImagePage(dfu_current_flash_page_data) == 0)
    {
        *(unsigned int *)data = 0;
     }
    else
    {
        *(unsigned int *)data = 1;
    }
    return 4;
}

int dfu_flash_cmd_read_page_data(unsigned char *data)
{
    unsigned char *page_data_ptr = &dfu_current_flash_page_data[dfu_current_flash_subpage_index * 64];
    memcpy(data, page_data_ptr, 64);

    dfu_current_flash_subpage_index++;

    return 64;
}

static int begin_write()
{
    int result;
    // TODO this will take a long time. To minimise the amount of time spent
    // paused on this operation it would be preferable to move to this to a
    // seperate command, e.g. start_write.
    do
    {
        result = fl_startImageAdd(&dfu_factory_image, DFU_FLASH_MAX_UPGRADE_SIZE, 0);
    } while (result > 0);

    if (result < 0)
    {
        DFU_FLASH_ERROR();
    }
    return result;
}

static int pages_written = 0;

int dfu_flash_cmd_write_page(unsigned char *data)
{
    /* 改动原因：协议仅约定第一个字节为 flag(0=第一页,1=后续页,2=结束)；
     * 若用 *(unsigned int *)data 读取 4 字节，而调用处只设置 data[0]=0 且未清零 data[1..3]，
     * 则 flag 可能非 0，导致不进入 case 0，begin_write() 被跳过。改为仅用首字节。 */
    unsigned int flag = (unsigned int)data[0];
    if (flag == 0)
    debug_printf("dfu_flash_cmd_write_page: flag = %d\n", flag);

    if (dfu_upgrade_image_valid)
    {
        debug_printf("dfu_flash_cmd_write_page: dfu_upgrade_image_valid = 1\n");
        return 0;
    }

    switch (flag)
    {
        case 0:
            // First page.
            debug_printf("dfu_flash_cmd_write_page: begin_write\n");
            int result = begin_write();
            if (result < 0)
                return result;
            pages_written = 0;
            // fallthrough
        case 1:
            // Do nothing.
            break;
        case 2:
            // Termination.
            if (fl_endWriteImage() != 0)
                DFU_FLASH_ERROR();

            // Sanity check
            fl_BootImageInfo image = dfu_factory_image;
            if (fl_getNextBootImage(&image) != 0)
                DFU_FLASH_ERROR();
            break;
    }
    dfu_current_flash_subpage_index = 0;

    return 0;
}

int dfu_flash_cmd_write_page_data(unsigned char *data)
{
    unsigned char *page_data_ptr = &dfu_current_flash_page_data[dfu_current_flash_subpage_index * 64];

    if (dfu_upgrade_image_valid)
    {
        return 0;
    }

    if (dfu_current_flash_subpage_index >= 4)
    {
        return 0;
    }

    memcpy(page_data_ptr, data, 64);

    dfu_current_flash_subpage_index++;

    if (dfu_current_flash_subpage_index == 4)
    {
        if (fl_writeImagePage(dfu_current_flash_page_data) != 0)
            DFU_FLASH_ERROR();
        pages_written++;
    }

    return 0;
}


int dfu_flash_cmd_erase_all(void)
{
    fl_BootImageInfo tmp_image = dfu_upgrade_image;

    if (dfu_upgrade_image_valid)
    {
        if (fl_deleteImage(&dfu_upgrade_image) != 0)
        {
            DFU_FLASH_ERROR();
        }

        // Keep deleting all upgrade images
        // TODO Perhaps using replace would be nicer...
        while(1)
        {
            if (fl_getNextBootImage(&tmp_image) == 0)
            {
                if (fl_deleteImage(&tmp_image) != 0)
                {
                    DFU_FLASH_ERROR();
                }
            }
            else
            {
                break;
            }
        }

    dfu_upgrade_image_valid = 0;
    }
    /* 改动原因：强制升级流程中，erase 后必须允许后续第一页写执行 begin_write()；
     * 若仅在上面 if 块内清零，某些路径（如删除失败或 while 未正确退出）可能导致
     * dfu_upgrade_image_valid 仍为 1，write_page(flag=0) 会直接 return 0 而不调用
     * begin_write()，导致 begin_write 未执行。此处无条件清零，保证 erase 后下一
     * 次 write_page(0) 必定执行 begin_write。 */
    dfu_upgrade_image_valid = 0;
    debug_printf("dfu_flash_cmd_erase_all: set dfu_upgrade_image_valid = 0\n");
    return 0;
}
#endif
