/** 
 * @file    mcu.c
 * @brief   MCU module
 * @author  Kenny Tran
 * @date    Aug 16 2024
 * @version 0.0.1
 *
 * @section LICENSES
 * Copyright (C) Morgan Advanced Programmable Systems, Inc.
 */


/*--------------------------------------------------------------------------------------------------
 * INCLUDE FILES
 *------------------------------------------------------------------------------------------------*/
#include <stdint.h>
#include <stdbool.h>
#include <string.h>

#include <zephyr/kernel.h>
#include <zephyr/pm/device.h>
#include <zephyr/sys/poweroff.h>
#include <zephyr/sys/reboot.h>

#include <hal/nrf_power.h>

#include "../lib/errno.h"
#include "../lib/crc.h"
#include "../lib/print.h"

#include "mcu.h"


/*--------------------------------------------------------------------------------------------------
 * TYPE DEFINITIONS & ENUMERATIONS
 *------------------------------------------------------------------------------------------------*/


/*--------------------------------------------------------------------------------------------------
 * MACROS & CONSTANTS
 *------------------------------------------------------------------------------------------------*/
#define _MCU_RAM_RETENTION_CRC_LEN  (4)

/* nRF52 RAM (really, RAM AHB slaves) are partitioned as:
 * * Up to 8 blocks of two 4 KiBy byte "small" sections
 * * A 9th block of with 32 KiBy "large" sections
 *
 * At time of writing the maximum number of large sections is 6, all
 * within the first large block.  Theoretically there could be more
 * sections in the 9th block, and possibly more blocks.
 */
/* Inclusive address of RAM start */
#define _MCU_SRAM_BEGIN (uintptr_t)DT_REG_ADDR(DT_NODELABEL(sram0))

/* Exclusive address of RAM end */
#define _MCU_SRAM_END (_MCU_SRAM_BEGIN + (uintptr_t)DT_REG_SIZE(DT_NODELABEL(sram0)))

/* Size of a controllable RAM section in the small blocks */
#define _MCU_SMALL_SECTION_SIZE (4096)

/* Number of controllable RAM sections in each of the lower blocks */
#define _MCU_SMALL_SECTIONS_PER_BLOCK (2)

/* Span of a small block */
#define _MCU_SMALL_BLOCK_SIZE (_MCU_SMALL_SECTIONS_PER_BLOCK * _MCU_SMALL_SECTION_SIZE)

/* Number of small blocks */
#define _MCU_SMALL_BLOCK_COUNT (8)

/* Span of the SRAM area covered by small sections */
#define _MCU_SMALL_SECTION_SPAN (_MCU_SMALL_BLOCK_COUNT * _MCU_SMALL_BLOCK_SIZE)

/* Inclusive address of the RAM range covered by large sections */
#define _MCU_LARGE_SECTION_BEGIN (_MCU_SRAM_BEGIN + _MCU_SMALL_SECTION_SPAN)

/* Size of a controllable RAM section in large blocks */
#define _MCU_LARGE_SECTION_SIZE (32768)


/*--------------------------------------------------------------------------------------------------
 * PRIVATE CONSTANTS & VARIABLES
 *------------------------------------------------------------------------------------------------*/
__noinit static uint8_t _mcu_ram_data[MCU_RAM_RETENTION_LEN + _MCU_RAM_RETENTION_CRC_LEN];
// another option: static uint8_t _mcu_ram_data[MCU_RAM_RETENTION_LEN + _MCU_RAM_RETENTION_CRC_LEN] __attribute__((section(".noinit")));
         static bool    _mcu_init = false;

/*--------------------------------------------------------------------------------------------------
 * PRIVATE FUNCTION PROTOTYPES
 *------------------------------------------------------------------------------------------------*/
/* Set or clear RAM retention in SYSTEM_OFF for the provided object.
 *
 * @note This only works for nRF52 with the POWER module.  The other
 * Nordic chips use a different low-level API, which is not currently
 * used by this function.
 *
 * @param ptr pointer to the start of the retainable object
 *
 * @param len length of the retainable object
 *
 * @param enable true to enable retention, false to clear retention
 */
static void _mcu_ram_retain (const void *ptr, size_t len, bool enable);


/*--------------------------------------------------------------------------------------------------
 * PUBLIC FUNCTION DEFINITIONS
 *------------------------------------------------------------------------------------------------*/
ERRNO_e MCU_init (void)
{
    if (_mcu_init)
    {
        return ERRNO_ALREADY_INITIALIZED;
    }
    else if (CRC_init () < ERRNO_SUCCESS)
    {
        return ERRNO_CRC_FAILURE;
    }

    _mcu_init = true;

    return ERRNO_SUCCESS;
}

ERRNO_e MCU_reset (bool cold)
{
    if (!_mcu_init)
    {
        return ERRNO_NOT_INITIALIZED;
    }

    if (cold)
    {
        sys_reboot (SYS_REBOOT_COLD);
    }
    else
    {
        sys_reboot (SYS_REBOOT_WARM);
    }

    return ERRNO_SUCCESS;
}

ERRNO_e MCU_sleep (void)
{
    if (!_mcu_init)
    {
        return ERRNO_NOT_INITIALIZED;
    }
    
    sys_poweroff ();

    return ERRNO_SUCCESS;
}

ERRNO_e MCU_ram_retain (uint8_t *p_data, uint16_t len, bool enable)
{
    uint32_t checksum;

    if (!_mcu_init)
    {
        return ERRNO_NOT_INITIALIZED;
    }
    else if (len > MCU_RAM_RETENTION_LEN)
    {
        return ERRNO_LIMIT_REACHED;
    }

    if (enable)
    {
        /* Calculate CRC checksum of data. */
        if (CRC_32_calculate (p_data, len, CRC_32_CHECKSUM_INIT, &checksum))
        {
            return ERRNO_CRC_FAILURE;
        }

        /* Copy data to retention buffer. */
        memcpy (_mcu_ram_data, p_data, len);
        /* Append CRC checksum. */
        memcpy (&_mcu_ram_data[len], &checksum, _MCU_RAM_RETENTION_CRC_LEN);
    }

    _mcu_ram_retain (_mcu_ram_data, sizeof (_mcu_ram_data), enable);

    return ERRNO_SUCCESS;
}

ERRNO_e MCU_ram_fetch (uint8_t *p_dest, uint16_t len)
{
    uint32_t checksum;

    if (!_mcu_init)
    {
        return ERRNO_NOT_INITIALIZED;
    }
    else if (len > MCU_RAM_RETENTION_LEN)
    {
        return ERRNO_LIMIT_REACHED;
    }

    /* Calculate CRC checksum of data. */
    if (CRC_32_calculate (_mcu_ram_data, len, CRC_32_CHECKSUM_INIT, &checksum))
    {
        return ERRNO_CRC_FAILURE;
    }

    if (checksum != *(uint32_t *) &_mcu_ram_data[len])
    {
        memset (p_dest, 0, len);
        return ERRNO_CRC_FAILURE;
    }

    /* Copy data to retention buffer. */
    memcpy (p_dest, _mcu_ram_data, len);

    return ERRNO_SUCCESS;
}

void MCU_cli_cmd (int32_t argc, char **argv)
{
    if (argc == 1)
    {
        if (!strcmp (argv[0], "--init") || !strcmp (argv[0], "-i"))
        {
            PRINT (INFO, "Ret code : %d\r\n", MCU_init ());
        }
        else if (!strcmp (argv[0], "--sleep") || !strcmp (argv[0], "-s"))
        {
	        PRINT (INFO, "Ret code : %d\r\n", MCU_sleep ());
        }
        else
        {
            goto USAGE;
        }
    }
    else if (argc == 2)
    {
        if (!strcmp (argv[0], "--reset") || !strcmp (argv[0], "-r"))
        {
            if (!strcmp (argv[1], "cold"))
            {
                PRINT (INFO, "Ret code : %d\r\n", MCU_reset (true));
            }
            else if (!strcmp (argv[1], "warm"))
            {
                PRINT (INFO, "Ret code : %d\r\n", MCU_reset (false));
            }
            else
            {
                goto USAGE;
            }
        }
        else
        {
            goto USAGE;
        }
    }
    else if (argc == 3)
    {

        if (!strcmp (argv[0], "--ram"))
        {
            if (!strcmp (argv[1], "fetch"))
            {
                uint8_t len = (uint8_t) strtoul (argv[2], NULL, 0);
                uint8_t buf[len + 1];

                buf[len] = 0; // NULL termination

                PRINT (INFO, "Ret code : %d\r\n", MCU_ram_fetch (buf, len));

                PRINT (INFO, "Ram data : %s\r\n", buf);
            }
            else if (!strcmp (argv[1], "write"))
            {
                PRINT (INFO, "Ret code : %d\r\n", MCU_ram_retain (argv[2], strlen (argv[2]), true));
            }
            else
            {
                goto USAGE;
            }
        }
        else
        {
            goto USAGE;
        }
    }
    else
    {
USAGE:
        PRINT (INFO, "Usage: mcu [OPTION...]\r\n");
        PRINT (INFO, "  -i | --init          Initializes MCU\r\n");
        PRINT (INFO, "  -r | --reset         Resets MCU\r\n");
        PRINT (INFO, "    <%%s>               [cold, warm]\r\n");
        PRINT (INFO, "  -s | --sleep         Puts MCU to sleep\r\n");
    }
}


/*--------------------------------------------------------------------------------------------------
 * PRIVATE FUNCTION DEFINITIONS
 *------------------------------------------------------------------------------------------------*/
static void _mcu_ram_retain (const void *ptr, size_t len, bool enable)
{
	uintptr_t addr     = (uintptr_t)ptr;
	uintptr_t addr_end = addr + len;

	/* Error if the provided range is empty or doesn't lie
	 * entirely within the SRAM address space.
	 */
	if ((len == 0U) || (addr < _MCU_SRAM_BEGIN) || (addr > (_MCU_SRAM_END - len)))
    {
        PRINT (ERROR, "Range is empty or doesn't lie entirely within SRAM space\r\n");
		return;
	}

	/* Iterate over each section covered by the range, setting the
	 * corresponding RAM OFF retention bit in the parent block.
	 */
	do {
		uintptr_t block_base         = _MCU_SRAM_BEGIN;
		uint32_t  section_size       = _MCU_SMALL_SECTION_SIZE;
		uint32_t  sections_per_block = _MCU_SMALL_SECTIONS_PER_BLOCK;
		bool      is_large           = (addr >= _MCU_LARGE_SECTION_BEGIN);
		uint8_t   block              = 0;
        uint32_t  section;
        uint32_t  section_mask;

		if (is_large)
        {
			block        = 8;
			block_base   = _MCU_LARGE_SECTION_BEGIN;
			section_size = _MCU_LARGE_SECTION_SIZE;

			/* RAM[x] supports only 16 sections, each its own bit
			 * for POWER (0..15) and RETENTION (16..31).  We don't
			 * know directly how many sections are present, so
			 * assume they all are; the true limit will be
			 * determined by the SRAM size.
			 */
			sections_per_block = 16;
		}

		section = (addr - block_base) / section_size;

		if (section >= sections_per_block)
        {
			block   += section / sections_per_block;
			section %= sections_per_block;
		}

        section_mask = (POWER_RAM_POWERSET_S0RETENTION_On << (section + POWER_RAM_POWERSET_S0RETENTION_Pos));

		if (enable)
        {
			nrf_power_rampower_mask_on (NRF_POWER, block, section_mask);
		}
        else
        {
			nrf_power_rampower_mask_off (NRF_POWER, block, section_mask);
		}

		/* Move to the first address in the next section. */
		addr += section_size - (addr % section_size);
	} while (addr < addr_end);
}
