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

#include "FreeRTOS.h"
#include "task.h"
#include "lab_freertos.h"
#include "lab_io.h"
#include "task01_scheduler.h"
#include "task01_scheduler_model.h"

#define MCAUSE_MACHINE_TIMER_INTERRUPT ( 0x80000007UL )

volatile uint32_t g_last_checkpoint;
volatile uintptr_t g_checkpoint_subject;

volatile uint32_t g_switch_task[ SCHED_TRACE_CAPACITY ];
volatile TickType_t g_switch_tick[ SCHED_TRACE_CAPACITY ];
volatile uint32_t g_switch_count;
volatile uint32_t g_last_peer_id = UINT32_MAX;
volatile uint32_t g_peer_changes;
volatile uint32_t g_aba_task[ 3 ];
volatile TickType_t g_aba_tick[ 3 ];
volatile uint32_t g_aba_seen;

volatile uint32_t g_preemption_order[ 3 ];
volatile uint32_t g_preemption_order_count;
volatile uint32_t g_priority_raise_started;
volatile uint32_t g_peer_after_raise;
volatile uint32_t g_stop_peers;
volatile uint32_t g_peers_finished;
volatile uint32_t g_verifier_started;

volatile uint32_t g_tick_hook_count;
volatile TickType_t g_first_tick_count;
volatile uintptr_t g_first_tick_sp;
volatile uintptr_t g_first_tick_mepc;
volatile uint32_t g_first_tick_mcause;

volatile uint32_t g_scheduler_pass;

PeerContext g_peer_a = {
    .id = SCHED_PEER_A,
    .configured_priority = SCHED_PEER_PRIORITY,
    .handle = NULL,
    .phase = SCHED_PHASE_PREPARED
};

PeerContext g_peer_b = {
    .id = SCHED_PEER_B,
    .configured_priority = SCHED_PEER_PRIORITY,
    .handle = NULL,
    .phase = SCHED_PHASE_PREPARED
};

HighProbeContext g_high_probe = {
    .initial_priority = SCHED_PROBE_INITIAL_PRIORITY,
    .target_priority = SCHED_PROBE_TARGET_PRIORITY,
    .handle = NULL,
    .phase = SCHED_PHASE_PREPARED
};

VerifierContext g_verifier = {
    .configured_priority = SCHED_VERIFY_PRIORITY,
    .handle = NULL,
    .phase = SCHED_PHASE_PREPARED
};

static uintptr_t read_sp( void )
{
    uintptr_t value;

    __asm__ volatile( "mv %0, sp" : "=r"( value ) );
    return value;
}

static uintptr_t read_mepc( void )
{
    uintptr_t value;

    __asm__ volatile( "csrr %0, mepc" : "=r"( value ) );
    return value;
}

static uint32_t read_mcause( void )
{
    uint32_t value;

    __asm__ volatile( "csrr %0, mcause" : "=r"( value ) );
    return value;
}

__attribute__( ( noinline, used ) )
void task01_scheduler_checkpoint_committed( uint32_t point,
                                             const void * subject )
{
    ( void ) point;
    ( void ) subject;
    __asm__ volatile( "" ::: "memory" );
}

__attribute__( ( noinline, used ) )
void task01_scheduler_checkpoint( uint32_t point, const void * subject )
{
    g_last_checkpoint = point;
    g_checkpoint_subject = ( uintptr_t ) subject;
    __asm__ volatile( "" ::: "memory" );
    task01_scheduler_checkpoint_committed( point, subject );
}

__attribute__( ( noinline, used ) )
void task01_tick_checkpoint_committed( void )
{
    __asm__ volatile( "" ::: "memory" );
}

void vApplicationTickHook( void )
{
    ++g_tick_hook_count;

    if( g_tick_hook_count == 1U )
    {
        g_first_tick_count = xTaskGetTickCountFromISR();
        g_first_tick_sp = read_sp();
        g_first_tick_mepc = read_mepc();
        g_first_tick_mcause = read_mcause();
        g_last_checkpoint = 6U;
        g_checkpoint_subject = ( uintptr_t ) &g_tick_hook_count;
        __asm__ volatile( "" ::: "memory" );
        task01_tick_checkpoint_committed();
    }
}

static void capture_peer_runtime( PeerContext * peer )
{
    TaskStatus_t info;
    uintptr_t first;
    uintptr_t last;

    vTaskGetInfo( NULL, &info, pdFALSE, eRunning );
    first = ( uintptr_t ) info.pxStackBase;
    last = ( uintptr_t ) info.pxEndOfStack;

    peer->observed_priority = uxTaskPriorityGet( NULL );
    peer->tcb_address = ( uintptr_t ) info.xHandle;
    peer->entry_sp = read_sp();
    peer->stack_low = first < last ? first : last;
    peer->stack_high = first < last ? last : first;
}

static void capture_probe_runtime( HighProbeContext * probe )
{
    TaskStatus_t info;
    uintptr_t first;
    uintptr_t last;

    vTaskGetInfo( NULL, &info, pdFALSE, eRunning );
    first = ( uintptr_t ) info.pxStackBase;
    last = ( uintptr_t ) info.pxEndOfStack;

    probe->observed_priority = uxTaskPriorityGet( NULL );
    probe->tcb_address = ( uintptr_t ) info.xHandle;
    probe->entry_sp = read_sp();
    probe->stack_low = first < last ? first : last;
    probe->stack_high = first < last ? last : first;
}

static int volatile_trace_has_aba( void )
{
    uint32_t index;

    for( index = 2U; index < g_switch_count; ++index )
    {
        if( g_switch_task[ index - 2U ] == SCHED_PEER_A &&
            g_switch_task[ index - 1U ] == SCHED_PEER_B &&
            g_switch_task[ index ] == SCHED_PEER_A )
        {
            return 1;
        }
    }
    return 0;
}

static BaseType_t record_peer_switch( uint32_t id, TickType_t tick )
{
    BaseType_t aba_completed_by_a = pdFALSE;

    taskENTER_CRITICAL();
    if( g_last_peer_id != id )
    {
        if( g_last_peer_id != UINT32_MAX )
        {
            ++g_peer_changes;
        }
        g_last_peer_id = id;

        if( g_switch_count < SCHED_TRACE_CAPACITY )
        {
            g_switch_task[ g_switch_count ] = id;
            g_switch_tick[ g_switch_count ] = tick;
            ++g_switch_count;
        }
    }

    if( id == SCHED_PEER_A &&
        g_aba_seen == 0U &&
        volatile_trace_has_aba() != 0 )
    {
        const uint32_t last = g_switch_count - 1U;

        g_aba_task[ 0 ] = g_switch_task[ last - 2U ];
        g_aba_task[ 1 ] = g_switch_task[ last - 1U ];
        g_aba_task[ 2 ] = g_switch_task[ last ];
        g_aba_tick[ 0 ] = g_switch_tick[ last - 2U ];
        g_aba_tick[ 1 ] = g_switch_tick[ last - 1U ];
        g_aba_tick[ 2 ] = g_switch_tick[ last ];
        g_aba_seen = 1U;
        aba_completed_by_a = pdTRUE;
    }
    taskEXIT_CRITICAL();

    return aba_completed_by_a;
}

static void append_preemption_marker( uint32_t marker )
{
    taskENTER_CRITICAL();
    if( g_preemption_order_count < 3U )
    {
        g_preemption_order[ g_preemption_order_count ] = marker;
        ++g_preemption_order_count;
    }
    taskEXIT_CRITICAL();
}

static int sp_inside_peer_stack( const PeerContext * peer )
{
    return peer->entry_sp >= peer->stack_low &&
           peer->entry_sp <= peer->stack_high;
}

static int sp_inside_probe_stack( const HighProbeContext * probe )
{
    return probe->entry_sp >= probe->stack_low &&
           probe->entry_sp <= probe->stack_high;
}

__attribute__( ( noinline, used ) )
void peer_task_entry( void * pv_parameters )
{
    PeerContext * peer = ( PeerContext * ) pv_parameters;
    TickType_t observed_tick;
    BaseType_t emit_entry = pdFALSE;

    if( peer == NULL || peer->id > SCHED_PEER_B )
    {
        lab_fail( 0xFC030101U );
    }

    capture_peer_runtime( peer );
    peer->phase = SCHED_PHASE_RUNNING;
    peer->first_tick = xTaskGetTickCount();
    observed_tick = peer->first_tick;

    taskENTER_CRITICAL();
    if( g_peer_a.iterations == 0U && g_peer_b.iterations == 0U )
    {
        emit_entry = pdTRUE;
    }
    taskEXIT_CRITICAL();

    if( emit_entry != pdFALSE )
    {
        task01_scheduler_checkpoint( 5U, peer );
    }

    for( ; ; )
    {
        const TickType_t tick = xTaskGetTickCount();

        ++peer->iterations;
        peer->checksum += ( peer->id + 1U ) *
                          ( ( peer->iterations & 0xFFU ) + 1U );

        if( tick != observed_tick || g_switch_count == 0U )
        {
            observed_tick = tick;

            if( record_peer_switch( peer->id, tick ) != pdFALSE )
            {
                task01_scheduler_checkpoint( 7U, peer );

                taskENTER_CRITICAL();
                if( g_priority_raise_started == 0U )
                {
                    g_priority_raise_started = 1U;
                    taskEXIT_CRITICAL();

                    append_preemption_marker( SCHED_ORDER_BEFORE_RAISE );
                    task01_scheduler_checkpoint( 8U, peer );
                    vTaskPrioritySet( g_high_probe.handle,
                                      SCHED_PROBE_TARGET_PRIORITY );
                    append_preemption_marker( SCHED_ORDER_AFTER_RAISE );
                    g_peer_after_raise = 1U;
                    g_stop_peers = 1U;
                    task01_scheduler_checkpoint( 10U, peer );
                }
                else
                {
                    taskEXIT_CRITICAL();
                }
            }
        }

        if( g_stop_peers != 0U )
        {
            break;
        }

        if( ( tick - peer->first_tick ) > SCHED_TIMEOUT_TICKS )
        {
            lab_fail( 0xFC030102U );
        }
    }

    peer->last_tick = xTaskGetTickCount();
    peer->phase = SCHED_PHASE_COMPLETED;
    peer->handle = NULL;

    taskENTER_CRITICAL();
    ++g_peers_finished;
    if( g_peers_finished == 2U )
    {
        taskEXIT_CRITICAL();
        task01_scheduler_checkpoint( 11U, peer );
    }
    else
    {
        taskEXIT_CRITICAL();
    }

    vTaskDelete( NULL );
    lab_fail( 0xFC030103U );
}

__attribute__( ( noinline, used ) )
void high_probe_task_entry( void * pv_parameters )
{
    HighProbeContext * probe = ( HighProbeContext * ) pv_parameters;

    if( probe == NULL )
    {
        lab_fail( 0xFC030201U );
    }

    capture_probe_runtime( probe );
    probe->phase = SCHED_PHASE_RUNNING;
    probe->run_tick = xTaskGetTickCount();
    append_preemption_marker( SCHED_ORDER_HIGH_PROBE );
    probe->ran_before_caller_returned =
        ( g_peer_after_raise == 0U ) &&
        ( g_preemption_order_count == 2U );
    probe->verifier_was_not_running = ( g_verifier_started == 0U );
    task01_scheduler_checkpoint( 9U, probe );

    probe->phase = SCHED_PHASE_COMPLETED;
    probe->handle = NULL;
    vTaskDelete( NULL );
    lab_fail( 0xFC030202U );
}

static int trace_ticks_are_nondecreasing( void )
{
    uint32_t index;

    for( index = 1U; index < g_switch_count; ++index )
    {
        if( g_switch_tick[ index ] < g_switch_tick[ index - 1U ] )
        {
            return 0;
        }
    }
    return 1;
}

__attribute__( ( noinline, used ) )
void verifier_task_entry( void * pv_parameters )
{
    VerifierContext * verifier = ( VerifierContext * ) pv_parameters;

    if( verifier == NULL )
    {
        lab_fail( 0xFC030301U );
    }

    verifier->phase = SCHED_PHASE_RUNNING;
    verifier->observed_priority = uxTaskPriorityGet( NULL );
    g_verifier_started = 1U;

    g_scheduler_pass =
        configUSE_PREEMPTION == 1 &&
        configUSE_TIME_SLICING == 1 &&
        g_peer_a.observed_priority == SCHED_PEER_PRIORITY &&
        g_peer_b.observed_priority == SCHED_PEER_PRIORITY &&
        g_high_probe.observed_initial_priority ==
            SCHED_PROBE_INITIAL_PRIORITY &&
        g_high_probe.observed_priority == SCHED_PROBE_TARGET_PRIORITY &&
        verifier->observed_priority == SCHED_VERIFY_PRIORITY &&
        g_peer_a.phase == SCHED_PHASE_COMPLETED &&
        g_peer_b.phase == SCHED_PHASE_COMPLETED &&
        g_high_probe.phase == SCHED_PHASE_COMPLETED &&
        g_peer_a.handle == NULL &&
        g_peer_b.handle == NULL &&
        g_high_probe.handle == NULL &&
        g_peers_finished == 2U &&
        g_peer_a.iterations > 0U &&
        g_peer_b.iterations > 0U &&
        g_peer_a.checksum != 0U &&
        g_peer_b.checksum != 0U &&
        g_tick_hook_count >= 2U &&
        g_first_tick_mcause == MCAUSE_MACHINE_TIMER_INTERRUPT &&
        g_first_tick_sp != 0U &&
        g_first_tick_mepc != 0U &&
        g_aba_seen == 1U &&
        g_aba_task[ 0 ] == SCHED_PEER_A &&
        g_aba_task[ 1 ] == SCHED_PEER_B &&
        g_aba_task[ 2 ] == SCHED_PEER_A &&
        g_aba_tick[ 0 ] <= g_aba_tick[ 1 ] &&
        g_aba_tick[ 1 ] <= g_aba_tick[ 2 ] &&
        trace_ticks_are_nondecreasing() != 0 &&
        g_peer_changes >= 2U &&
        g_priority_raise_started == 1U &&
        g_peer_after_raise == 1U &&
        g_high_probe.ran_before_caller_returned == 1U &&
        g_high_probe.verifier_was_not_running == 1U &&
        g_preemption_order_count == 3U &&
        g_preemption_order[ 0 ] == SCHED_ORDER_BEFORE_RAISE &&
        g_preemption_order[ 1 ] == SCHED_ORDER_HIGH_PROBE &&
        g_preemption_order[ 2 ] == SCHED_ORDER_AFTER_RAISE &&
        g_high_probe.run_tick >= g_aba_tick[ 2 ] &&
        g_peer_a.created_handle == g_peer_a.tcb_address &&
        g_peer_b.created_handle == g_peer_b.tcb_address &&
        g_high_probe.created_handle == g_high_probe.tcb_address &&
        sp_inside_peer_stack( &g_peer_a ) != 0 &&
        sp_inside_peer_stack( &g_peer_b ) != 0 &&
        sp_inside_probe_stack( &g_high_probe ) != 0;

    verifier->pass = g_scheduler_pass;
    verifier->phase = SCHED_PHASE_COMPLETED;
    verifier->handle = NULL;
    task01_scheduler_checkpoint( 12U, verifier );

    if( g_scheduler_pass == 0U )
    {
        lab_fail( 0xFC030302U );
    }

    lab_puts( "PASS FC03: tick + priorities + preemption + time slicing\n" );
    lab_put_u32( g_tick_hook_count );
    lab_put_u32( g_peer_changes );
    lab_put_u32( g_preemption_order_count );
    lab_exit( 0U );
}

static void create_task_or_fail( TaskFunction_t entry,
                                 const char * name,
                                 configSTACK_DEPTH_TYPE stack_words,
                                 void * context,
                                 UBaseType_t priority,
                                 TaskHandle_t * handle,
                                 uint32_t failure_code )
{
    if( xTaskCreate( entry,
                     name,
                     stack_words,
                     context,
                     priority,
                     handle ) != pdPASS )
    {
        lab_fail( failure_code );
    }
}

int main( void )
{
    lab_heap_initialize();
    task01_scheduler_checkpoint( 1U, &g_peer_a );

    create_task_or_fail( high_probe_task_entry,
                         "high-probe",
                         SCHED_PROBE_STACK_WORDS,
                         &g_high_probe,
                         SCHED_PROBE_INITIAL_PRIORITY,
                         &g_high_probe.handle,
                         0xFC030401U );
    g_high_probe.created_handle = ( uintptr_t ) g_high_probe.handle;
    g_high_probe.phase = SCHED_PHASE_READY;
    task01_scheduler_checkpoint( 2U, &g_high_probe );

    create_task_or_fail( peer_task_entry,
                         "peer-a",
                         SCHED_PEER_STACK_WORDS,
                         &g_peer_a,
                         SCHED_PEER_PRIORITY,
                         &g_peer_a.handle,
                         0xFC030402U );
    g_peer_a.created_handle = ( uintptr_t ) g_peer_a.handle;
    g_peer_a.phase = SCHED_PHASE_READY;

    create_task_or_fail( peer_task_entry,
                         "peer-b",
                         SCHED_PEER_STACK_WORDS,
                         &g_peer_b,
                         SCHED_PEER_PRIORITY,
                         &g_peer_b.handle,
                         0xFC030403U );
    g_peer_b.created_handle = ( uintptr_t ) g_peer_b.handle;
    g_peer_b.phase = SCHED_PHASE_READY;

    create_task_or_fail( verifier_task_entry,
                         "verify",
                         SCHED_VERIFY_STACK_WORDS,
                         &g_verifier,
                         SCHED_VERIFY_PRIORITY,
                         &g_verifier.handle,
                         0xFC030404U );
    g_verifier.phase = SCHED_PHASE_READY;

    g_high_probe.observed_initial_priority =
        uxTaskPriorityGet( g_high_probe.handle );
    g_peer_a.observed_priority = uxTaskPriorityGet( g_peer_a.handle );
    g_peer_b.observed_priority = uxTaskPriorityGet( g_peer_b.handle );
    g_verifier.observed_priority = uxTaskPriorityGet( g_verifier.handle );
    task01_scheduler_checkpoint( 3U, &g_high_probe );

    if( g_high_probe.observed_initial_priority !=
            SCHED_PROBE_INITIAL_PRIORITY ||
        g_peer_a.observed_priority != SCHED_PEER_PRIORITY ||
        g_peer_b.observed_priority != SCHED_PEER_PRIORITY ||
        g_verifier.observed_priority != SCHED_VERIFY_PRIORITY )
    {
        lab_fail( 0xFC030405U );
    }

    task01_scheduler_checkpoint( 4U, &g_peer_a );
    vTaskStartScheduler();
    lab_fail( 0xFC030406U );
}