#include "rtthread.h"
#include "ipc/completion.h"
#include "ymodem.h"

#include "fal.h"
#define IAP_FLASH_APP_SEC "app"
#define IAP_FLASH_BOOT_SEC "boot"

struct rt_completion _wait;

#define IS_AF(c)  ((c >= 'A') && (c <= 'F'))
#define IS_af(c)  ((c >= 'a') && (c <= 'f'))
#define IS_09(c)  ((c >= '0') && (c <= '9'))
#define ISVALIDHEX(c)  IS_AF(c) || IS_af(c) || IS_09(c)
#define ISVALIDDEC(c)  IS_09(c)
#define CONVERTDEC(c)  (c - '0')

#define CONVERTHEX_alpha(c)  (IS_AF(c) ? (c - 'A'+10) : (c - 'a'+10))
#define CONVERTHEX(c)   (IS_09(c) ? (c - '0') : CONVERTHEX_alpha(c))

/**
  * @brief  Convert a string to an integer
  * @param  inputstr: The string to be converted
  * @param  intnum: The intger value
  * @retval 1: Correct
  *         0: Error
  */
rt_inline rt_uint32_t Str2Int(uint8_t *inputstr, int32_t *intnum)
{
  rt_uint32_t i = 0, res = 0;
  rt_uint32_t val = 0;

  if (inputstr[0] == '0' && (inputstr[1] == 'x' || inputstr[1] == 'X'))
  {
    if (inputstr[2] == '\0')
    {
      return 0;
    }
    for (i = 2; i < 11; i++)
    {
      if (inputstr[i] == '\0')
      {
        *intnum = val;
        /* return 1; */
        res = 1;
        break;
      }
      if (ISVALIDHEX(inputstr[i]))
      {
        val = (val << 4) + CONVERTHEX(inputstr[i]);
      }
      else
      {
        /* return 0, Invalid input */
        res = 0;
        break;
      }
    }
    /* over 8 digit hex --invalid */
    if (i >= 11)
    {
      res = 0;
    }
  }
  else /* max 10-digit decimal input */
  {
    for (i = 0;i < 11;i++)
    {
      if (inputstr[i] == '\0')
      {
        *intnum = val;
        /* return 1 */
        res = 1;
        break;
      }
      else if ((inputstr[i] == 'k' || inputstr[i] == 'K') && (i > 0))
      {
        val = val << 10;
        *intnum = val;
        res = 1;
        break;
      }
      else if ((inputstr[i] == 'm' || inputstr[i] == 'M') && (i > 0))
      {
        val = val << 20;
        *intnum = val;
        res = 1;
        break;
      }
      else if (ISVALIDDEC(inputstr[i]))
      {
        val = val * 10 + CONVERTDEC(inputstr[i]);
      }
      else
      {
        /* return 0, Invalid input */
        res = 0;
        break;
      }
    }
    /* Over 10 digit decimal --invalid */
    if (i >= 11)
    {
      res = 0;
    }
  }

  return res;
}

rt_err_t port_rx_ind(rt_device_t dev, rt_size_t size)
{
    rt_completion_done(&_wait);
    return RT_EOK;
}

#define RT_BUF_SIZE (1024)

/*
 * 固件格式:
 *
 * 作用 & 说明
 * - 记录固件对应硬件 & 版本信息
 * - 简便起见，不包含校验信息
 *
 * index   len      desc
 *   0      4       'XYFC'
 *   4      4       固件信息长度
 *   8      N       固件信息 ('类型'-'版本')
 *   8+N    M       固件数据
 */
static const char FIXED_HEAD[] = {'X', 'Y', 'F', 'C'};
extern const char IAP_BIN_PREFIX[];

static int _check_head(const rt_uint8_t *buf, rt_size_t len)
{
    if (len < 8)
    {
        return -1;
    }

    for (int i = 0; i < sizeof(FIXED_HEAD); i++)
    {
        if (FIXED_HEAD[i] != buf[i])
        {
            return -2;
        }
    }

    rt_uint32_t real = *((const rt_uint32_t *)&buf[4]);
    rt_uint32_t want = rt_strlen(IAP_BIN_PREFIX);
    if (real <= want || len - 8 < real)
    {
        return -3;
    }

    for(int i = 0; i < want; i++)
    {
        if (IAP_BIN_PREFIX[i] != buf[8 + i])
        {
            return -4;
        }
    }

    return real + 8;
}

struct custom_ctx {
    struct rym_ctx parent;
    int32_t file_size;
    int32_t w_index;
    int32_t recevied_head;
    const struct fal_partition *flash;
    rt_uint8_t ck_buffer[RT_BUF_SIZE];
};

static enum rym_code _rym_recv_begin(
    struct rym_ctx *ctx,
    rt_uint8_t *buf,
    rt_size_t len)
{
    // return RYM_ERR_ACK;
    // return RYM_CODE_CAN;
    struct custom_ctx *cctx = (struct custom_ctx *) ctx;
    RT_ASSERT(cctx);

    rt_uint8_t tmp_size[16] = {0};
    while (*buf != 0) 
    {
        buf ++;
    }
    buf ++;

    rt_bool_t end = RT_FALSE;
    for (rt_uint8_t i = 0; i < 16; i++)
    {
        if (*buf == 0)
        {
            tmp_size[i] = 0;
            end = RT_TRUE;
            break;
        }

        tmp_size[i] = *buf++;
    }
    if (!end)
    {
        return RYM_CODE_CAN;
    }
    if (!Str2Int(tmp_size, &cctx->file_size))
    {
        return RYM_CODE_CAN;
    }

    if (cctx->file_size >= cctx->flash->len)
    {
        return RYM_CODE_CAN;
    }

    cctx->w_index = 0;
    cctx->recevied_head = 0;

    return RYM_CODE_ACK;
}

static enum rym_code _rym_recv_data(
    struct rym_ctx *ctx,
    rt_uint8_t *buf,
    rt_size_t len)
{
    struct custom_ctx *cctx = (struct custom_ctx *) ctx;
    RT_ASSERT(cctx);
    RT_ASSERT(len <= RT_BUF_SIZE);

    int offset = 0;
    if (cctx->recevied_head == 0)
    {
        offset = _check_head(buf, len);
        if (offset < 0)
        {
            return RYM_ERR_FILE;
        }

        if (fal_partition_erase(cctx->flash, 0, cctx->file_size) < 0)
        {
            return RYM_CODE_CAN;
        }

        cctx->recevied_head = 1;
    }

    if (cctx->w_index >= cctx->file_size)
    {
        return RYM_CODE_ACK;
    }

    if (cctx->w_index + len > cctx->file_size)
    {
        len = cctx->file_size - cctx->w_index;
    }

    if (fal_partition_write(cctx->flash, cctx->w_index, &buf[offset], len - offset) <= 0)
    {
        return RYM_CODE_CAN;
    }

    if (fal_partition_read(cctx->flash, cctx->w_index, cctx->ck_buffer, len - offset) <= 0)
    {
        return RYM_CODE_CAN;
    }

    for (int i = 0; i < len - offset; i++)
    {
        if (cctx->ck_buffer[i] != buf[i + offset])
            return RYM_CODE_CAN;
    }
    cctx->w_index += len - offset;

    return RYM_CODE_ACK;
}

static enum rym_code _rym_recv_end(
    struct rym_ctx *ctx,
    rt_uint8_t *buf,
    rt_size_t len)
{
    return RYM_CODE_ACK;
}

enum iap_v5_cmd
{
    IAP_V5_CMD_ONE_REPLY = 1,
    IAP_V5_CMD_C_REPLY,
    IAP_V5_CMD_UPDATE_SUCCEED,
    IAP_V5_CMD_UPDATE_FAILED,
    IAP_V5_CMD_BIN_OVERSIZE,
    IAP_V5_CMD_CHECK_FAILED,
    IAP_V5_CMD_USER_CANCEL,
    IAP_V5_CMD_RECEIVE_FAILED,
    IAP_V5_CMD_POWER_ON_AGAIN,
};

/**
 * @brief 响应协议
 *
 *  0xB5,0x5B,{cmd},0xBB, 0x00
 */
#define IAP_V5_CMD_SIZE (5)

static const rt_uint8_t *_pack_cmd(enum iap_v5_cmd cmd)
{
    static rt_uint8_t _cmd_line[IAP_V5_CMD_SIZE]
        = { 0xB5, 0x5B, 0x00, 0xBB, 0x00 };

    _cmd_line[2] = cmd;
    return _cmd_line;
}

#define _response_cmd(console_dev, cmd)                                                    \
    rt_device_write(console_dev, 0, _pack_cmd(cmd), IAP_V5_CMD_SIZE)

static void jump_to_app(void);
extern const uint32_t IapAppAddr;

#define PORTS_NUM (10) // uart1 - uart10

static rt_device_t _read_ports(rt_device_t ports[], rt_uint8_t *pch)
{
    for (int i = 0; i < PORTS_NUM; i++)
    {
        while (ports[i])
        {
            if (rt_device_read(ports[i], 0, pch, 1) != 1)
            {
                break;
            }

            if (*pch == '1' || *pch == 'q')
            {
                return ports[i];
            }
        }
    }

    return RT_NULL;
}

static void _print_ports(rt_device_t ports[], const char *fmt, ...)
{
    va_list args;
    char line_buf[128] = { 0 };
    rt_size_t real;

    va_start(args, fmt);
    real = rt_vsnprintf(line_buf, sizeof(line_buf), fmt, args);
    va_end(args);

    for (int i = 0; i < PORTS_NUM; i++)
    {
        if (ports[i] == RT_NULL)
        {
            break;
        }

        rt_device_write(ports[i], 0, line_buf, real);
    }
}

static int _iap_entry(rt_bool_t is_boot);

int iap_main_entry(void)
{
    return _iap_entry(RT_TRUE);
}

#include "shell.h"

static int _iap_entry(rt_bool_t is_boot)
{
    if (IapAppAddr && is_boot)
    {
        return 0;
    }

    if ((!IapAppAddr) && is_boot)
    {
        finsh_set_prompt("msh (IAP)");
    }

    rt_err_t err = RT_EOK;
    struct custom_ctx *ctx = RT_NULL;
    rt_device_t console = rt_console_get_device();
    RT_ASSERT(console);

    rt_device_t ports[10] = {0};
    char name[RT_NAME_MAX] = {0};

    rt_completion_init(&_wait);

    rt_bool_t _wait_mode = RT_FALSE;
    rt_bool_t _continue_run = RT_FALSE;
    rt_err_t (*odev_rx_ind)(rt_device_t dev, rt_size_t size);

    for (int i = 0; i < 10 && is_boot; i++)
    {
        rt_snprintf(name, RT_NAME_MAX, "uart%d", i+1);
        ports[i] = rt_device_find(name);
        RT_ASSERT(ports[i]);

        struct serial_configure config = RT_SERIAL_CONFIG_DEFAULT;
        config.baud_rate = BAUD_RATE_460800;
        rt_device_control(ports[i], RT_DEVICE_CTRL_CONFIG, &config);

        rt_device_set_rx_indicate(ports[i], port_rx_ind);
        rt_device_open(ports[i], RT_DEVICE_OFLAG_RDWR | RT_DEVICE_FLAG_INT_RX);
    }

    if (!is_boot)
    {
        ports[0] = rt_console_get_device();
        RT_ASSERT(ports[0]);

        odev_rx_ind = ports[0]->rx_indicate;
        rt_device_set_rx_indicate(ports[0], port_rx_ind);

        struct serial_configure config = RT_SERIAL_CONFIG_DEFAULT;
        config.baud_rate = BAUD_RATE_460800;
        rt_device_control(ports[0], RT_DEVICE_CTRL_CONFIG, &config);

        rt_device_open(ports[0], RT_DEVICE_OFLAG_RDWR | RT_DEVICE_FLAG_INT_RX);
    }

    rt_device_t get_port = RT_NULL;
    rt_uint8_t wait_cnt = 3;
    rt_uint8_t ch;

    _print_ports(ports, "iap loop ...");

    while (wait_cnt > 0 || !is_boot)
    {
        get_port = _read_ports(ports, &ch);

        if (get_port != RT_NULL)
        {
            if (ch == 'q')
            {
                _continue_run = RT_TRUE;
                break;
            }

            if (ch == '1')
            {
                _wait_mode = RT_TRUE;
                _response_cmd(get_port, IAP_V5_CMD_ONE_REPLY);
                break;
            }
        }
        else
        {
            wait_cnt--;
        }
 
        rt_completion_wait(&_wait, rt_tick_from_millisecond(100));
    }

    if (get_port == RT_NULL)
    {
        get_port = console;
    }

    for (int i = 0; i < PORTS_NUM && is_boot; i++)
    {
        if (ports[i] == get_port)
        {
            continue;
        }

        rt_device_set_rx_indicate(ports[i], RT_NULL);
        rt_device_close(ports[i]);
    }

    if (get_port != console)
    {
        rt_console_set_device(get_port->parent.name);
    }

    if (_continue_run)
    {
        rt_device_close(get_port);
        if (!is_boot)
        {
            rt_device_set_rx_indicate(ports[0], odev_rx_ind);
        }
        return 0;
    }

    if (!_wait_mode)
    {
        err = RT_EOK;
        goto _exit;
    }

    rt_device_set_rx_indicate(get_port, port_rx_ind);

    rt_thread_mdelay(200);
    _response_cmd(get_port, IAP_V5_CMD_C_REPLY);

    ctx = rt_calloc(1, sizeof(*ctx));
    if (!ctx)
    {
        err = RT_ENOMEM;
        rt_kprintf("rt_malloc failed\n");
        goto _exit;
    }

    ctx->flash = fal_partition_find(IapAppAddr? IAP_FLASH_BOOT_SEC: IAP_FLASH_APP_SEC);
    if (ctx->flash == RT_NULL)
    {
        err = RT_ENOSYS;
        rt_kprintf("fal_partition_find %s failed: %d\n", IapAppAddr? IAP_FLASH_BOOT_SEC: IAP_FLASH_APP_SEC, rt_get_errno());
        rt_free(ctx);
        goto _exit;
    }

    err = rym_recv_on_device(&ctx->parent, get_port, RT_DEVICE_OFLAG_RDWR | RT_DEVICE_FLAG_INT_RX,
                             _rym_recv_begin, _rym_recv_data, _rym_recv_end, 1000);
    rt_free(ctx);

    rt_thread_mdelay(200);

    if (err == RT_EOK)
    {
        rt_device_close(get_port);
        _response_cmd(get_port, IAP_V5_CMD_UPDATE_SUCCEED);
        if (!is_boot)
        {
            rt_device_set_rx_indicate(ports[0], odev_rx_ind);
        }
        return 0;
    }

_exit:
    rt_device_close(get_port);
    if (!is_boot)
    {
        rt_device_set_rx_indicate(ports[0], odev_rx_ind);
    }

    if (err == RT_EOK && is_boot)
    {
        rt_kprintf("jump to app ...\n");
        jump_to_app();
        return 0;
    }

    _response_cmd(get_port, IAP_V5_CMD_UPDATE_FAILED);
    _response_cmd(get_port, IAP_V5_CMD_POWER_ON_AGAIN);

    return 0;
}

#include <stm32f4xx.h>
#include <stdint.h>

#define FLASH_RUN_BASEADDR 0x08020000

//定义函数指针，用于跳转到APP
typedef  void (*pFunction)(void);
static pFunction JumpToApplication;
static uint32_t JumpAddress;

static void jump_to_app(void)
{
    SysTick->CTRL = 0;                                //关键代码
  
    HAL_DeInit();                                     //可选
    HAL_NVIC_DisableIRQ(SysTick_IRQn);                //可选
    HAL_NVIC_ClearPendingIRQ(SysTick_IRQn);           //可选

    /* Jump to user application */
    __disable_irq();

    JumpAddress = *(__IO uint32_t*) (FLASH_RUN_BASEADDR + 4);

    JumpToApplication = (pFunction) JumpAddress;
    /* Initialize user application's Stack Pointer */
    __set_MSP(*(__IO uint32_t*) FLASH_RUN_BASEADDR);

    JumpToApplication();
}
MSH_CMD_EXPORT(jump_to_app, jump_to_app);

static void _cmd_iap(void)
{
    _iap_entry(RT_FALSE);
}
MSH_CMD_EXPORT_ALIAS(_cmd_iap, iap, iap bootloader / app start);