#include "stm32g4xx.h"
#include <stddef.h>

/* TCB */
typedef struct {
    uint32_t *sp;
} TCB;

#define MAX_TASKS  3
#define STACK_SIZE 256  /* words (1KB per task) */

static uint32_t task_stacks[MAX_TASKS][STACK_SIZE];
static TCB      tcbs[MAX_TASKS];
static uint8_t  num_tasks;
static uint8_t  curr_idx;

/* asm reads these */
TCB *os_curr_task;
TCB *os_next_task;

/* Fake initial stack frame so PendSV can "restore" a new task */
void os_task_init(TCB *tcb, void (*handler)(void),
                  uint32_t *stack_base, uint32_t stack_size)
{
    uint32_t *sp = stack_base + stack_size;
    *(--sp) = 0x01000000;
    *(--sp) = (uint32_t)handler;
    *(--sp) = 0xDEADDEAD;
    *(--sp) = 0; *(--sp) = 0; *(--sp) = 0; *(--sp) = 0; *(--sp) = 0;
    for (int i = 0; i < 8; i++) *(--sp) = 0;
    tcb->sp = sp;
}

/* Register a task */
void os_task_create(void (*handler)(void))
{
    uint8_t id = num_tasks++;
    os_task_init(&tcbs[id], handler, task_stacks[id], STACK_SIZE);
}

/* Round-robin scheduler */
void os_sched(void)
{
    curr_idx = (curr_idx + 1) % num_tasks;
    os_next_task = &tcbs[curr_idx];
    if (os_next_task != os_curr_task)
        SCB->ICSR |= SCB_ICSR_PENDSVSET_Msk;  /* pend PendSV */
}

/* SysTick calls scheduler each tick */
void SysTick_Handler(void)
{
    os_sched();
}

/* Start the OS */
void os_start(void)
{
    __disable_irq();

    curr_idx     = 0;
    /* tells asm: first switch */
    os_curr_task = NULL;
    os_next_task = &tcbs[0];

    /* PendSV = lowest priority (context switch never preempts other IRQs) */
    NVIC_SetPriority(PendSV_IRQn, 0xFF);

    /* SysTick every 10ms (assuming 16MHz HSI default) */
    SysTick_Config(16000000 / 100);

    __enable_irq();

    /* Trigger first context switch — PendSV will load task0's stack into PSP
       and return with EXC_RETURN=0xFFFFFFFD, which switches to PSP automatically */
    SCB->ICSR |= SCB_ICSR_PENDSVSET_Msk;

    /* Wait here on MSP until PendSV fires and jumps to task0 */
    while (1);
}

/* Two test tasks: toggle PA5 at different rates */
static void delay(volatile uint32_t count)
{
    while (count--);
}

void task0(void)
{
    while (1) {
        GPIOA->ODR ^= (1U << 5);   /* toggle PA5 */
        delay(1000000);
    }
}

void task1(void)
{
    while (1) {
        GPIOA->ODR ^= (1U << 5);
        delay(4000000);              /* slower toggle */
    }
}

/* main */
int main(void)
{
    RCC->AHB2ENR |= RCC_AHB2ENR_GPIOAEN;
    __DSB();

    GPIOA->MODER &= ~GPIO_MODER_MODE5_Msk;
    GPIOA->MODER |= (1U << GPIO_MODER_MODE5_Pos);

    os_task_create(task0);
    os_task_create(task1);
    os_start();
}
