#define DEBUG_PRINT_ENABLE 0
#include <stdint.h>
#include "debug_print.h"
#include <quadflashlib.h>
#include <xcore/chanend.h>
#include <xcore/channel.h>
#include "hmac.h"
#include "xua_conf.h"
#include "swlock.h"

// Global flash lock defined in lfs_services.c — protects all flash hardware access
extern swlock_t flash_lock;

#define WRITE_ENABLE_COMMAND      (0x06)
#define WRITE_DISABLE_COMMAND     (0x04)

#define ERASE_CHIP_COMMAND        (0xC7)


#define READ_STATUS_REG_COMMAND   (0x05)
#define READ_ID_COMMAND           (0x9F)

#define READ_UID_COMMAND          (0x4B)

#define WRITE_ENABLE_STATUS_COMMAND (0x50)

#define WRITE_STATUS_REG1_COMMAND (0x01)
#define READ_STATUS_REG2_COMMAND  (0x35)
#define WRITE_STATUS_REG2_COMMAND (0x31)
#define READ_STATUS_REG3_COMMAND  (0x15)
#define WRITE_STATUS_REG3_COMMAND (0x11)

#define RDSTATUS_COMMAND          (0x35)
#define WRSTATUS_COMMAND           (0x31)

#define PRSCUR_COMMAND            (0x42)
#define ERSCUR_COMMAND            (0x44)
#define RDSCUR_COMMAND            (0x48)
#define KEY_BLOCK_LEN  100
#define HMAC_GEN_LEN  20


extern unsigned short XUA_Endpoint0_getProductId();
extern unsigned short XUA_Endpoint0_getVendorId();

unsigned char hid_reply[64] = {0};

void printArrayHex1(uint8_t array[], int size, int elementsPerLine) {
    for (int i = 0; i < size; i++) {
        debug_printf("0x%x ", array[i]);

        // 在达到每行元素个数时换行
        if ((i + 1) % elementsPerLine == 0) {
            debug_printf("\n");
        }
    }
    debug_printf("\n");
}

#define settw(a,b) {__asm__ __volatile__("settw res[%0], %1": : "r" (a) , "r" (b));}
#define setc(a,b) {__asm__  __volatile__("setc res[%0], %1": : "r" (a) , "r" (b));}
#define setclk(a,b) {__asm__ __volatile__("setclk res[%0], %1": : "r" (a) , "r" (b));}
#define portin(a,b) {__asm__  __volatile__("in %0, res[%1]": "=r" (b) : "r" (a));}
#define portout(a,b) {__asm__  __volatile__("out res[%0], %1": : "r" (a) , "r" (b));}

#ifdef DFU_FLASH_DEVICE

#if (XUA_QUAD_SPI_FLASH)
/* Using specified flash device rather than all supported in tools */
fl_QuadDeviceSpec opt_flash_devices[] = {DFU_FLASH_DEVICE};
#else
/* Using specified flash device rather than all supported in tools */
fl_DeviceSpec opt_flash_devices[] = {DFU_FLASH_DEVICE};
#endif
#endif

/*
typedef struct {
      out port qspiCS;
      out port qspiSCLK;
      out buffered port:32 qspiSIO;
      clock qspiClkblk;
} fl_QSPIPorts;
*/
fl_QSPIPorts p_opt_tile0 =
{
    XS1_PORT_1B,
    XS1_PORT_1C,
    XS1_PORT_4B,
    XS1_CLKBLK_3
};

fl_QSPIPorts p_opt_tile1 =
{
    XS1_PORT_1F,
    XS1_PORT_1G,
    XS1_PORT_4D,
    XS1_CLKBLK_3
};

#define KEY_ADDRESS                0xFFFD0

/* Store Flag to fixed address */
static void SetKeyFlag(unsigned x)
{
    asm volatile("stw %0, %1[0]" :: "r"(x), "r"(KEY_ADDRESS));
}

/* Load flag from fixed address */
unsigned GetKeyFlag()
{
    unsigned x;
    asm volatile("ldw %0, %1[0]" : "=r"(x) : "r"(KEY_ADDRESS));
    return x;
}

/* return 1 for opened ports successfully */
int flash_opt_enable_ports(fl_QSPIPorts *p_qflash)
{
    int result = 0;

#ifdef DFU_FLASH_DEVICE
    result = fl_connectToDevice(p_qflash, opt_flash_devices, sizeof(opt_flash_devices) / sizeof(fl_QuadDeviceSpec));
#else
    /* Use default flash list */
    result = fl_connect(p_qflash);
#endif
    if (!result)
    {
        /* All okay.. */
        debug_printf("flash_opt_enable_ports success\n");
        return 1;
    }
    else
    {
        debug_printf("flash_opt_enable_ports failed\n");
        return 0;
    }
}

int flash_opt_disable_ports()
{
    fl_disconnect();
    return 1;
}

void flash_cmd_opt(unsigned int cmd,
                  unsigned char input[], unsigned int num_in,
                  unsigned char output[], unsigned int num_out)
{
    fl_command(cmd,input,num_in,output, num_out);
}

void flash_opt_erase()
{
    uint8_t cmd[4] = {0, 0x10, 0x0, 0};
    uint8_t data[4] = {0, 0x10, 0x0, 0};
    flash_cmd_opt(WRITE_ENABLE_COMMAND, cmd, 0, data, 0);
    delay_milliseconds(1);
    flash_cmd_opt(ERSCUR_COMMAND, cmd, 3, data, 0);
    delay_milliseconds(50);
}

void flash_opt_read_did(uint8_t data[], unsigned char len)
{
        uint8_t cmd[3] = {0x3, 0x10, 0x0};

        flash_cmd_opt(READ_ID_COMMAND, cmd, 0, data, len);
       // printArrayHex1(data, len, 20);
}


void flash_opt_unlock(void)
{
    uint8_t cmd[1] = {0};
    uint8_t data[1] = {0};
    uint8_t sr1, sr2, sr3;
    uint8_t did[3];

    swlock_acquire(&flash_lock);

    /* Reason: 使用 tile0 端口连接 Flash，保持原有逻辑 */
    if (flash_opt_enable_ports(&p_opt_tile0) == 0) {
        debug_printf("flash_opt_unlock failed\n");
        swlock_release(&flash_lock);
        return;
    }

    /* Reason: 增加 Flash 厂商ID检查，仅针对 Winbond (0xEF) 执行特定解锁逻辑 */
    flash_opt_read_did(did, 3);
    if (did[0] != 0xEF)
    {
        debug_printf("Not Winbond flash (DID: %02x), skip unlock\n", did[0]);
        flash_opt_disable_ports();
        return;
    }

    debug_printf("Unlock Winbond Flash...\n");

    /* Reason: 解锁 Status Register 1。
       参考 flash_api_ext.c 逻辑，清除 Bit 2-6 (BP0-BP4, TB, SEC)
       Mask: 0x7C (0111 1100) */
    flash_cmd_opt(READ_STATUS_REG_COMMAND, cmd, 0, data, 1);
    sr1 = data[0];
    if (sr1 & 0x7C)
    {
        debug_printf("Unlock SR1: %02x -> %02x\n", sr1, sr1 & ~0x7C);
        flash_cmd_opt(WRITE_ENABLE_STATUS_COMMAND, cmd, 0, data, 0);
        data[0] = sr1 & ~0x7C;
        flash_cmd_opt(WRITE_STATUS_REG1_COMMAND, data, 1, cmd, 0);
        delay_milliseconds(10); /* 等待写入完成 */
    }

    /* Reason: 解锁 Status Register 2。
       参考 flash_api_ext.c 逻辑，清除 Bit 6 (CMP)
       Mask: 0x40 (0100 0000) */
    flash_cmd_opt(READ_STATUS_REG2_COMMAND, cmd, 0, data, 1);
    sr2 = data[0];
    if ((sr2 & 0x40) || ((sr2 & 0x2) == 0))
    {
        debug_printf("Unlock SR2: %02x -> %02x\n", sr2, (sr2 & ~0x40) | 0x2);
        flash_cmd_opt(WRITE_ENABLE_STATUS_COMMAND, cmd, 0, data, 0);
        data[0] = (sr2 & ~0x40) | 0x2;
        flash_cmd_opt(WRITE_STATUS_REG2_COMMAND, data, 1, cmd, 0);
        delay_milliseconds(10);
    }

    /* Reason: 解锁 Status Register 3。
       参考 flash_api_ext.c 逻辑，清除 Bit 2 (WPS)
       Mask: 0x04 (0000 0100) */
    flash_cmd_opt(READ_STATUS_REG3_COMMAND, cmd, 0, data, 1);
    sr3 = data[0];
    if (sr3 & 0x04)
    {
        debug_printf("Unlock SR3: %02x -> %02x\n", sr3, sr3 & ~0x04);
        flash_cmd_opt(WRITE_ENABLE_STATUS_COMMAND, cmd, 0, data, 0);
        data[0] = sr3 & ~0x04;
        flash_cmd_opt(WRITE_STATUS_REG3_COMMAND, data, 1, cmd, 0);
        delay_milliseconds(10);
    }

    flash_opt_disable_ports();
    swlock_release(&flash_lock);
    debug_printf("Unlock Winbond Flash done\n");
}

void flash_opt_write2(uint8_t data[], unsigned char len)
{
        uint8_t d[1];
        uint8_t cmd[256] =  {0, 0x10, 0x0};

        for (int i = 0; i < len; i++)
        {
            cmd[i + 3] = data[i];
        }

        flash_opt_erase();
        flash_cmd_opt(WRITE_ENABLE_COMMAND, cmd, 0, d, 0);
        flash_cmd_opt(PRSCUR_COMMAND, cmd, 3 + len, d, 0);
        delay_milliseconds(50);
}

void flash_opt_read(uint8_t data[], unsigned char len)
{
        uint8_t cmd[3] = {0x00, 0x10, 0x00};

        flash_cmd_opt(RDSCUR_COMMAND, cmd, 3, data, len);
        //printf("read sec key:\n");
        //printArrayHex1(data, len, 20);
}

void flash_opt_read_uid(uint8_t data[], unsigned char len)
{
        uint8_t cmd[3] = {0x3, 0x10, 0x0};

        flash_cmd_opt(READ_UID_COMMAND, cmd, 0, data, len);
       // printArrayHex1(data, len, 20);
}



unsigned compare_buff_diff(const uint8_t* actual, const uint8_t* expected, unsigned len)
{
  for (int i = 0; i < len; ++i)
  {
    if (actual[i] != expected[i])
    {
        debug_printf("compare_buff_diff: %d %02x %02x\n", i, actual[i], expected[i]);
        return 0;
    }
  }
  return 1;
}


void cal_hmac(uint8_t uid[], uint8_t did[], uint8_t hmac_bin[])
{
    {
        uint8_t key_bin[HMAC_GEN_LEN] = {0x76, 0x37, 0x0b, 0xb8, 0xbd, 0xe5, 0x54, 0x82, 0x8f, 0x13, 0x29, 0xa6, 0x26, 0x07, 0x2d, 0xe1, 0x11, 0x6c, 0xfb, 0x07};
        uint8_t msg_bin[HMAC_GEN_LEN];
        for (int i = 0; i < 3; i++)
        {
            msg_bin[i] = did[i];
        }

        for (int i = 3; i < 20; i++)
        {
            msg_bin[i] = uid[i];
        }

#if 0
        printf("msg_bin: \n");
        printArrayHex1(msg_bin, HMAC_GEN_LEN, 20);
#endif
        hmac_sha1(key_bin, HMAC_GEN_LEN, msg_bin, HMAC_GEN_LEN, hmac_bin);

#if 0
#endif
    }
}

uint8_t opt_key_read(uint8_t key[], unsigned offset)
{
     uint8_t read_bin[KEY_BLOCK_LEN + 1];
     uint8_t ret = 0;

     swlock_acquire(&flash_lock);
     if (flash_opt_enable_ports(&p_opt_tile0) == 0)
     {
        swlock_release(&flash_lock);
        return 0;
     }

     flash_opt_read(read_bin, KEY_BLOCK_LEN + 1);
     for (int i = 0; i < 20; i++)
     {
         key[i] = read_bin[i + offset * 20 + 1];
     }
     flash_opt_disable_ports();
     swlock_release(&flash_lock);
     return 1;
}

uint8_t key_verify_program(unsigned verify_write, uint8_t expected_bin[], unsigned offset)
{
     uint8_t read_bin[KEY_BLOCK_LEN + 1];
     uint8_t comp_bin[KEY_BLOCK_LEN + 1];
     uint8_t uid[20];
     uint8_t did[4];
     uint8_t ret = 0;
     uint8_t hmac_bin[20];

     flash_opt_read(read_bin, KEY_BLOCK_LEN + 1);
     flash_opt_read_uid(uid, 20);
     flash_opt_read_did(did, 3);


     if (verify_write)
     {
        if (offset == 1)
        {
            cal_hmac(uid, did, hmac_bin);
            if (compare_buff_diff(hmac_bin, expected_bin, 20) == 0)
            {
                return 0;
            }
        }

        for (int i = 0; i < 20; i++)
        {
            read_bin[i + offset * 20 + 1] = expected_bin[i];
        }

         flash_opt_write2(&read_bin[1], KEY_BLOCK_LEN);
         delay_milliseconds(10);

         flash_opt_read(comp_bin, KEY_BLOCK_LEN + 1);
         if (compare_buff_diff(&comp_bin[1], &read_bin[1], KEY_BLOCK_LEN) == 1)
         {
            debug_printf("compare OK\n");
             ret = 1;
         }
     }
     else
     {

        flash_opt_read(read_bin, KEY_BLOCK_LEN + 1);
        if (compare_buff_diff(expected_bin, &read_bin[offset * 20 + 1], 20) == 1)
        {
            debug_printf("key_verified\n");
            ret = 1;
        }
     }

     return ret;
}

void get_key_ret(uint8_t *buffer)
{
    //printf("get_key_ret2\n");
    for (int i = 0; i < 64; i ++)
    {
        buffer[i] = hid_reply[i];
        hid_reply[i] = 0;
    }

    //printArrayHex1(buffer, 20, 20);
}

static chanend_t g_flash_opt_c;

void flash_opt_set_chan(chanend_t c) {
    g_flash_opt_c = c;
    //debug_printf("flash_opt_set_chan\n");
}

void flash_opt_exchange_buffer(uint8_t in[64], uint8_t out[64]) {
    //debug_printf("flash_opt_exchange_buffer\n");
    chan_out_buf_byte (g_flash_opt_c , in, 40) ;
    //debug_printf("flash_opt_exchange_buffer1\n");
    chan_in_buf_byte(g_flash_opt_c, out, 40);
}


uint8_t key_validate(uint8_t offset)
{
    uint8_t ret = 0;
    uint8_t read_bin[KEY_BLOCK_LEN + 1];
    uint8_t uid[20];
    uint8_t did[4];
    uint8_t hmac_bin[20];

    swlock_acquire(&flash_lock);
    if (flash_opt_enable_ports(&p_opt_tile0) == 0)
    {
        swlock_release(&flash_lock);
        return 0;
    }

    flash_opt_read(read_bin, KEY_BLOCK_LEN + 1);
    flash_opt_read_uid(uid, 20);
    flash_opt_read_did(did, 3);

    cal_hmac(uid, did, hmac_bin);

    if (compare_buff_diff(hmac_bin, &read_bin[offset * 20 + 1], 20) == 1)
    {
        debug_printf("key_verified\n");
        ret = 1;
    }

    flash_opt_disable_ports();
    swlock_release(&flash_lock);
    return ret;
}


void read_uid_did(uint8_t uid[])
{
    swlock_acquire(&flash_lock);
    if (flash_opt_enable_ports(&p_opt_tile0) == 0)
    {
        swlock_release(&flash_lock);
        return;
    }
    flash_opt_read_uid(uid, 20);
    flash_opt_read_did(uid, 3);
    flash_opt_disable_ports();
    swlock_release(&flash_lock);
}


void program_key(uint8_t *buffer, int datalength)
{
    if (buffer[0] == 0x77 && buffer[1] == 0x5B)
    {
        uint8_t ret = 0;
        swlock_acquire(&flash_lock);
        switch (buffer[2])
        {
            case 1:
                debug_printf("verify key\n");
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    ret = key_verify_program(0, &buffer[3], 1);
                    if (ret == 1)
                    {
                        hid_reply[0] = 0x55;
                    }
                    else
                    {
                        hid_reply[0] = 0;
                    }
                    flash_opt_disable_ports();
                }
                break;
            case 2:
                debug_printf("write key\n");
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    ret = key_verify_program(1, &buffer[3], 1);

                    debug_printf("write key ret %d\n", ret);
                    if (ret == 1)
                    {
                        hid_reply[0] = 0x55;
                    }
                    else
                    {
                        hid_reply[0] = 0;
                    }
                    flash_opt_disable_ports();
                }

                break;
            case 3:
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    flash_opt_read_did(&hid_reply[1], 3);
                    debug_printf("return did\n");
                    hid_reply[0] = 0x55;
                    flash_opt_disable_ports();
                }
                break;
            case 4:
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    debug_printf("return uid\n");
                    flash_opt_read_uid(&hid_reply[1], 21);
                    hid_reply[0] = 0x55;
                    flash_opt_disable_ports();
                }
                break;
            case 5:
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    debug_printf("erase key\n");
                    flash_opt_erase();
                    hid_reply[0] = 0x55;
                    flash_opt_disable_ports();
                }
                break;
            case 6:
#if XUA_USB_EN == 1
                debug_printf("get pid vid\n");
                unsigned short pid = XUA_Endpoint0_getProductId();
                unsigned short vid = XUA_Endpoint0_getVendorId();
                hid_reply[2] = (pid >> 8) & 0xff;
                hid_reply[1] = pid & 0xff;
                hid_reply[4] = (vid >> 8) & 0xff;
                hid_reply[3] = vid & 0xff;
                hid_reply[0] = 0x55;
#endif
                break;
            case 7:
                debug_printf("write aizip key\n");
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    uint8_t ret = key_verify_program(1, &buffer[3], 0);

                    debug_printf("write key ret %d\n", ret);
                    if (ret == 1)
                    {
                        hid_reply[0] = 0x55;
                    }
                    else
                    {
                        hid_reply[0] = 0;
                    }
                    flash_opt_disable_ports();
                }
                break;
            case 12:
                debug_printf("verify aizip key\n");
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    ret = key_verify_program(0, &buffer[3], 0);
                    if (ret == 1)
                    {
                        hid_reply[0] = 0x55;
                    }
                    else
                    {
                        hid_reply[0] = 0;
                    }
                    flash_opt_disable_ports();
                }
                break;
            case 13:
            case 15:
            case 17:
            case 19:
            case 21:
            case 23:
            case 25:
            case 27:
            case 29:
            case 31:
            case 33:
                debug_printf("write key offset %d\n", ((buffer[2] - 13) / 2) + 2);
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    uint8_t ret = key_verify_program(1, &buffer[3], ((buffer[2] - 13) / 2) + 2);

                    debug_printf("write key ret %d\n", ret);
                    if (ret == 1)
                    {
                        hid_reply[0] = 0x55;
                    }
                    else
                    {
                        hid_reply[0] = 0;
                    }
                    flash_opt_disable_ports();
                }
                break;
            case 14:
            case 16:
            case 18:
            case 20:
            case 22:
            case 24:
            case 26:
            case 28:
            case 30:
            case 32:
            case 34:
                debug_printf("verify key offset %d\n", ((buffer[2] - 14) / 2) + 2);
                if (flash_opt_enable_ports(&p_opt_tile0) == 0)
                {
                    hid_reply[0] = 0;
                }
                else
                {
                    ret = key_verify_program(0, &buffer[3], ((buffer[2] - 14) / 2) + 2);
                    if (ret == 1)
                    {
                        hid_reply[0] = 0x55;
                    }
                    else
                    {
                        hid_reply[0] = 0;
                    }
                    flash_opt_disable_ports();
                }

                break;
            case 8:
            case 9:
            case 10:
            case 11:
                break;
        }
        swlock_release(&flash_lock);
    }
}
