#include <stddef.h>
#include <stdint.h>
#include <type_traits>

extern "C"
{
#include "FreeRTOS.h"
#include "lab_freertos.h"
#include "lab_io.h"
#include "task.h"
}

extern "C" void task_storage_debug_checkpoint( uint32_t point,
                                                 const void * object );
#define FREERTOS_STORAGE_CHECKPOINT( point, object ) \
    task_storage_debug_checkpoint( ( point ), ( object ) )
#include "freertos/task_storage.hpp"

volatile uint32_t g_last_checkpoint;
volatile uintptr_t g_checkpoint_object;
volatile size_t g_heap_baseline;
volatile size_t g_heap_after_dynamic_start;
volatile size_t g_heap_after_static_start;
volatile size_t g_heap_after_verifier_create;
volatile size_t g_heap_before_idle_cleanup;
volatile size_t g_heap_after_idle_cleanup;
volatile size_t g_dynamic_heap_cost;
volatile size_t g_idle_cleanup_recovered;
volatile uintptr_t g_dynamic_tcb;
volatile uintptr_t g_dynamic_stack_low;
volatile uintptr_t g_dynamic_stack_high;
volatile uintptr_t g_static_tcb;
volatile uintptr_t g_static_stack_low;
volatile uintptr_t g_static_stack_high;
volatile uintptr_t g_static_tcb_storage;
volatile uintptr_t g_static_stack_storage;
volatile size_t g_static_stack_words;
volatile size_t g_static_stack_bytes;
volatile uint32_t g_dynamic_result;
volatile uint32_t g_static_result;
volatile uint32_t g_workers_completed;
volatile uint32_t g_static_storage_survived;
volatile uint32_t g_task_storage_pass;

extern "C" __attribute__( ( noinline, used ) )
void task_storage_debug_checkpoint( uint32_t point, const void * object )
{
    g_last_checkpoint = point;
    g_checkpoint_object = reinterpret_cast< uintptr_t >( object );
    __asm__ volatile( "" ::: "memory" );
}

namespace
{

constexpr size_t worker_stack_words = 256U;
constexpr UBaseType_t worker_priority = 2U;
constexpr UBaseType_t verifier_priority = 1U;
constexpr uint32_t work_limit = 1000U;
constexpr uint32_t expected_result = 500500U;

void mark_completed() noexcept
{
    taskENTER_CRITICAL();
    ++g_workers_completed;
    taskEXIT_CRITICAL();
}

class DynamicWorker final : public freertos::DynamicTask< DynamicWorker >
{
    friend class freertos::DynamicTask< DynamicWorker >;

public:
    constexpr DynamicWorker() noexcept
        : DynamicTask { "dynamic", worker_stack_words, worker_priority }
    {
    }

private:
    void run() noexcept
    {
        TaskStatus_t status {};
        vTaskGetInfo( nullptr, &status, pdTRUE, eRunning );
        g_dynamic_tcb = reinterpret_cast< uintptr_t >( status.xHandle );
        g_dynamic_stack_low =
            reinterpret_cast< uintptr_t >( status.pxStackBase );
        g_dynamic_stack_high =
            reinterpret_cast< uintptr_t >( status.pxEndOfStack );

        uint32_t result = 0U;
        for( uint32_t value = 1U; value <= work_limit; ++value )
        {
            result += value;
        }
        g_dynamic_result = result;
        mark_completed();
    }
};

using StaticWorkerStorage =
    freertos::StaticTaskStorage< worker_stack_words >;

class StaticWorker final
    : public freertos::StaticTask< StaticWorker, worker_stack_words >
{
    friend class freertos::StaticTask< StaticWorker, worker_stack_words >;

public:
    constexpr explicit StaticWorker( StaticWorkerStorage & storage ) noexcept
        : StaticTask { "static", worker_priority, storage }
    {
    }

private:
    void run() noexcept
    {
        TaskStatus_t status {};
        vTaskGetInfo( nullptr, &status, pdTRUE, eRunning );
        g_static_tcb = reinterpret_cast< uintptr_t >( status.xHandle );
        g_static_stack_low =
            reinterpret_cast< uintptr_t >( status.pxStackBase );
        g_static_stack_high =
            reinterpret_cast< uintptr_t >( status.pxEndOfStack );

        uint32_t result = 0U;
        for( uint32_t value = 1U; value <= work_limit; ++value )
        {
            result += value;
        }
        g_static_result = result;
        mark_completed();
    }
};

static_assert( !std::is_copy_constructible< DynamicWorker >::value );
static_assert( !std::is_move_constructible< DynamicWorker >::value );
static_assert( !std::is_polymorphic< DynamicWorker >::value );
static_assert( !std::is_copy_constructible< StaticWorker >::value );
static_assert( !std::is_move_constructible< StaticWorker >::value );
static_assert( !std::is_polymorphic< StaticWorker >::value );
static_assert( StaticWorker::stack_bytes() ==
               worker_stack_words * sizeof( StackType_t ) );

DynamicWorker g_dynamic_worker;
StaticWorkerStorage g_static_worker_storage;
StaticWorker g_static_worker { g_static_worker_storage };
freertos::StaticTaskStorage< worker_stack_words > g_verifier_storage;

void verifier( void * )
{
    task_storage_debug_checkpoint( 9U, nullptr );
    g_heap_before_idle_cleanup = xPortGetFreeHeapSize();

    const bool lifecycle_before_idle =
        g_workers_completed == 2U &&
        g_dynamic_worker.state() ==
            freertos::StorageTaskState::completed &&
        g_static_worker.state() ==
            freertos::StorageTaskState::completed &&
        g_dynamic_worker.native_handle() == nullptr &&
        g_static_worker.native_handle() == nullptr;

    /* Let the idle task reclaim the dynamically allocated TCB and stack. */
    vTaskDelay( 2U );
    g_heap_after_idle_cleanup = xPortGetFreeHeapSize();
    g_idle_cleanup_recovered =
        g_heap_after_idle_cleanup - g_heap_before_idle_cleanup;
    task_storage_debug_checkpoint( 10U, nullptr );

    g_static_storage_survived =
        reinterpret_cast< uintptr_t >(
            g_static_worker.tcb_storage() ) == g_static_tcb_storage &&
        reinterpret_cast< uintptr_t >(
            g_static_worker.stack_storage() ) == g_static_stack_storage;

    const bool dynamic_storage_proved =
        g_dynamic_heap_cost > 0U &&
        lab_address_in_heap( reinterpret_cast< const void * >(
            g_dynamic_tcb ) ) != 0 &&
        lab_address_in_heap( reinterpret_cast< const void * >(
            g_dynamic_stack_low ) ) != 0 &&
        g_heap_before_idle_cleanup == g_heap_after_verifier_create &&
        g_idle_cleanup_recovered == g_dynamic_heap_cost &&
        g_heap_after_idle_cleanup == g_heap_baseline;
    const bool static_storage_proved =
        g_heap_after_static_start == g_heap_after_dynamic_start &&
        g_heap_after_verifier_create == g_heap_after_static_start &&
        g_static_tcb == g_static_tcb_storage &&
        g_static_stack_low == g_static_stack_storage &&
        lab_address_in_heap( reinterpret_cast< const void * >(
            g_static_tcb ) ) == 0 &&
        lab_address_in_heap( reinterpret_cast< const void * >(
            g_static_stack_low ) ) == 0 &&
        g_static_storage_survived != 0U &&
        g_static_stack_words == worker_stack_words &&
        g_static_stack_bytes ==
            worker_stack_words * sizeof( StackType_t );

    g_task_storage_pass =
        lifecycle_before_idle && dynamic_storage_proved &&
        static_storage_proved &&
        g_dynamic_result == expected_result &&
        g_static_result == expected_result;
    task_storage_debug_checkpoint( 11U, nullptr );
    if( g_task_storage_pass == 0U )
    {
        lab_fail( 0xE801U );
    }

    lab_puts( "PASS task01: dynamic/static task storage + idle cleanup\n" );
    lab_put_u32( static_cast< uint32_t >( g_dynamic_heap_cost ) );
    lab_put_u32( static_cast< uint32_t >( g_heap_before_idle_cleanup ) );
    lab_put_u32( static_cast< uint32_t >( g_heap_after_idle_cleanup ) );
    lab_put_u32( static_cast< uint32_t >( g_static_stack_words ) );
    lab_put_u32( static_cast< uint32_t >( g_static_stack_bytes ) );
    lab_exit( 0U );
}

} // namespace

int main()
{
    lab_heap_initialize();
    g_heap_baseline = xPortGetFreeHeapSize();
    task_storage_debug_checkpoint( 1U, &g_dynamic_worker );

    if( !g_dynamic_worker.start() )
    {
        lab_fail( 0xE802U );
    }
    g_heap_after_dynamic_start = xPortGetFreeHeapSize();
    g_dynamic_heap_cost = g_heap_baseline - g_heap_after_dynamic_start;
    task_storage_debug_checkpoint( 3U, &g_static_worker );

    g_static_tcb_storage = reinterpret_cast< uintptr_t >(
        g_static_worker.tcb_storage() );
    g_static_stack_storage = reinterpret_cast< uintptr_t >(
        g_static_worker.stack_storage() );
    g_static_stack_words = g_static_worker.stack_words();
    g_static_stack_bytes = g_static_worker.stack_bytes();
    if( !g_static_worker.start() )
    {
        lab_fail( 0xE803U );
    }
    g_heap_after_static_start = xPortGetFreeHeapSize();

    if( xTaskCreateStatic(
            verifier, "verify", worker_stack_words, nullptr,
            verifier_priority, g_verifier_storage.stack,
            &g_verifier_storage.tcb ) == nullptr )
    {
        lab_fail( 0xE804U );
    }
    g_heap_after_verifier_create = xPortGetFreeHeapSize();

    vTaskStartScheduler();
    lab_fail( 0xE8FFU );
}