/*
 * Copyright (c) 2019-2023 Beijing Hanwei Innovation Technology Ltd. Co. and
 * its subsidiaries and affiliates (collectly called MKSEMI).
 *
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form, except as embedded into an MKSEMI
 *    integrated circuit in a product or a software update for such product,
 *    must reproduce the above copyright notice, this list of conditions and
 *    the following disclaimer in the documentation and/or other materials
 *    provided with the distribution.
 *
 * 3. Neither the name of MKSEMI nor the names of its contributors may be used
 *    to endorse or promote products derived from this software without
 *    specific prior written permission.
 *
 * 4. This software, with or without modification, must only be used with a
 *    MKSEMI integrated circuit.
 *
 * 5. Any software provided in binary form under this license must not be
 *    reverse engineered, decompiled, modified and/or disassembled.
 *
 * THIS SOFTWARE IS PROVIDED BY MKSEMI "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY, NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL MKSEMI OR CONTRIBUTORS BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "mk_sbl.h"
#include "mk_flash.h"
#include "mk_clock.h"
#include "mk_trace.h"

#if !defined(FACT_BAK_MODE) && !defined(AB_CYCLE_MODE)
#define AB_CYCLE_MODE
#endif

#if defined(FACT_BAK_MODE)
#include "LzmaDec.h"
#define NORMAL_IMAGE_LOAD_ADDR IMAGE2_LOAD_ADDR
#define FACT_IMAGE_LOAD_ADDR IMAGE1_LOAD_ADDR
#endif

/*
        AB_CYCLE_MODE example
+---------+--------------+---------+
|  area   |     start    |  length |
+---------+--------------+---------+
|  sbl    |  0x04000000  | 12*1024 |
+---------+--------------+---------+
| active  |              |         |
|img flag |  0x04003000  |  4*1024 |
+---------+--------------+---------+
|  img1   |  0x04004000  | 128*1024|
+---------+--------------+---------+
|  img2   |  0x04024000  | 128*1024|
+---------+--------------+---------+

        FACT_BAK_MODE example
+----------+--------------+---------+
|   area   |     start    |  length |
+----------+--------------+---------+
|   sbl    |  0x04000000  | 12*1024 |
+----------+--------------+---------+
|compressed|              |         |
| factory  |  0x04004000  | 128*1024|
|   img    |              |         |
+----------+--------------+---------+
|normal img|  0x04024000  | 128*1024|
+----------+--------------+---------+

The address and length of each area can be determined based on actual conditions

*/

#if defined(FACT_BAK_MODE)

static char lzma_heap_flag = 0;
static char lzma_heap[16 * 1024];

static void *lzma_alloc(ISzAllocPtr p, size_t size)
{
    if (lzma_heap_flag == 0 && size < sizeof(lzma_heap))
    {
        return lzma_heap;
    }

    return NULL;
}

static void lzma_free(ISzAllocPtr p, void *address)
{
    if (address == lzma_heap)
    {
        lzma_heap_flag = 0;
    }
}

static ISzAlloc allocator = {lzma_alloc, lzma_free};

#endif

/*
 * Check whether there is a valid image at the specified address
 */
int sbl_check_image(uint32_t img_addr, uint32_t *img_cp_len, uint32_t *sp, uint32_t *entry)
{
    uint32_t *p = (uint32_t *)img_addr;
    struct BOOT_DESC_T *desc;

    if ((p[0] < SRAM_BASE) || (p[0] > (SRAM_BASE + SRAM_SIZE)))
    {
        return -1;
    }

    if (p[1] > (CODE_BASE + SRAM_SIZE))
    {
        return -1;
    }

    if ((p[9] > (CODE_BASE + SRAM_SIZE)) || ((p[9] & 0x03) != 0))
    {
        return -1;
    }

    desc = (struct BOOT_DESC_T *)(img_addr + p[9]);

    if (desc->header_marker != BOOT_HEADER_MARKER)
    {
        return -1;
    }

    if (img_cp_len != NULL)
    {
        *img_cp_len = desc->copy_len;
    }

    if (sp != NULL)
    {
        *sp = p[0];
    }

    if (entry != NULL)
    {
        *entry = p[1] + SRAM_BASE;
    }

    return 0;
}

/*
 * Get the address where the activated image is located
 */
uint32_t sbl_get_active_image(void)
{
#if defined(AB_CYCLE_MODE)
    uint32_t img = REG_READ(IMAGE_FLAG_ADDR);

    if ((img == IMAGE1_LOAD_ADDR) || (img == IMAGE2_LOAD_ADDR))
    {
        return img;
    }
    else
    {
        return 0;
    }
#else
    return NORMAL_IMAGE_LOAD_ADDR;
#endif
}

/*
 * Set the address of the activated image
 */
void sbl_set_active_image(uint32_t img)
{
#if defined(AB_CYCLE_MODE)
    if ((img != IMAGE1_LOAD_ADDR) && (img != IMAGE2_LOAD_ADDR))
    {
        return;
    }

    if (sbl_get_active_image() == img)
    {
        return;
    }

    uint32_t flag[64];
    flag[0] = img;

    flash_erase(FLASH_ID0, IMAGE_FLAG_ADDR, FLASH_SECTOR_SIZE);
    flash_write(FLASH_ID0, IMAGE_FLAG_ADDR, (uint8_t *)flag, FLASH_PAGE_SIZE);
#endif
}

/*
 * Get the address where the OTA image is stored
 */
uint32_t sbl_get_addr_for_ota(void)
{
#if defined(AB_CYCLE_MODE)
    uint32_t len;

    if (sbl_check_image(SRAM_BASE, &len, NULL, NULL) == 0)
    {
        if (memcmp((uint8_t *)IMAGE1_LOAD_ADDR, (uint8_t *)SRAM_BASE, len) == 0)
        {
            return IMAGE2_LOAD_ADDR;
        }

        if (memcmp((uint8_t *)IMAGE2_LOAD_ADDR, (uint8_t *)SRAM_BASE, len) == 0)
        {
            return IMAGE1_LOAD_ADDR;
        }
    }

    ASSERT(0, "A critical error has occurred in flash or RAM\r\n");

    return IMAGE1_LOAD_ADDR;
#else
    return NORMAL_IMAGE_LOAD_ADDR;
#endif
}

/*
 * Run the specified image
 */
static RAM_FUNC void sbl_run_image(uint32_t addr, uint32_t len, uint32_t sp, uint32_t entry);

#if defined(__ARMCC_VERSION) && (__ARMCC_VERSION >= 6100100)
#pragma clang optimize off
#elif defined(__GNUC__)
#pragma GCC diagnostic push
#pragma GCC optimize("O0")
#else
#pragma optimize = none
#endif

static void sbl_run_image(uint32_t addr, uint32_t len, uint32_t sp, uint32_t entry)
{
    static reset_handler_t *sbl2app;

    if (addr != SRAM_BASE)
    {
        memcpy((char *)SRAM_BASE, (char *)addr, len);
    }

    flash_handle[0].base->CMD = FLASH_CMD_DATALEN(0) | FLASH_CMD_POLL(false) | FLASH_CMD_DOUT(FLASH_DATA_OUT) | FLASH_CMD_INTLEN(0) |
                                FLASH_CMD_FIELDFORM(FLASH_CMD_ALL_SERIAL) | FLASH_CMD_FRAMEFORM(FLASH_CMD_OP_ONLY) | FLASH_CMD_OPCODE(0xB9);
    SYSCON->SYS_CMU &= ~(1U << CLOCK_FLASH_CTRL);
    sbl2app = (reset_handler_t *)entry;
    __set_MSP(sp);
    sbl2app();
}

#if defined(__ARMCC_VERSION) && (__ARMCC_VERSION >= 6100100)
#pragma clang optimize on
#elif defined(__GNUC__)
#pragma GCC diagnostic pop
#endif

/*
 * Main process of the second boot loader
 */
void sbl_main(void)
{
    uint32_t len, sp, entry, act_img;

    act_img = sbl_get_active_image();

#if defined(AB_CYCLE_MODE)
    if (act_img == IMAGE1_LOAD_ADDR || act_img == IMAGE2_LOAD_ADDR)
    {
        if (sbl_check_image(act_img, &len, &sp, &entry) == 0)
        {
            sbl_run_image(act_img, len, sp, entry);
        }
    }

    if (sbl_check_image(IMAGE1_LOAD_ADDR, &len, &sp, &entry) == 0)
    {
        sbl_run_image(IMAGE1_LOAD_ADDR, len, sp, entry);
    }
    else if (sbl_check_image(IMAGE2_LOAD_ADDR, &len, &sp, &entry) == 0)
    {
        sbl_run_image(IMAGE2_LOAD_ADDR, len, sp, entry);
    }
    else
    {
        while (1)
            ;
    }
#else
    if (act_img == NORMAL_IMAGE_LOAD_ADDR)
    {
        if (sbl_check_image(act_img, &len, &sp, &entry) == 0)
        {
            sbl_run_image(act_img, len, sp, entry);
        }
    }

    ELzmaStatus status;
    SizeT lzma_hdrlen = LZMA_PROPS_SIZE + sizeof(UInt64);
    SizeT srclen = 128 * 1024 - lzma_hdrlen;
    SizeT deslen = ROM_SIZE;
    SRes res;

    res = LzmaDecode((Byte *)SRAM_BASE, &deslen, (Byte *)FACT_IMAGE_LOAD_ADDR + lzma_hdrlen, &srclen, (Byte *)FACT_IMAGE_LOAD_ADDR,
                     LZMA_PROPS_SIZE + sizeof(UInt64), LZMA_FINISH_END, &status, &allocator);

    if (res != SZ_OK || status != LZMA_STATUS_FINISHED_WITH_MARK)
    {
        while (1)
            ;
    }

    if (sbl_check_image(SRAM_BASE, &len, &sp, &entry) == 0)
    {
        sbl_run_image(SRAM_BASE, len, sp, entry);
    }
    else
    {
        while (1)
            ;
    }
#endif
}