// taskswitch.c -- Cooperative task switching library for DJGPP
//
// Uses inline assembly for context switching (i386 and x86_64).  The
// scheduler uses credit-based weighted round-robin so all tasks run,
// but higher-priority tasks run proportionally more often.
//
// Task storage is a stb_ds dynamic array that grows as needed.
// Terminated task slots are recycled by tsCreate().

#define STB_DS_IMPLEMENTATION
#include "thirdparty/stb_ds.h"

#include "taskswitch.h"
#include <stdlib.h>
#include <string.h>

// ============================================================================
// Internal types
// ============================================================================

#if defined(__x86_64__)
// Saved CPU context for x86_64 (field order matches asm byte offsets)
typedef struct {
    uintptr_t rbx; // offset  0
    uintptr_t r12; // offset  8
    uintptr_t r13; // offset 16
    uintptr_t r14; // offset 24
    uintptr_t r15; // offset 32
    uintptr_t rbp; // offset 40
    uintptr_t rsp; // offset 48
    uintptr_t rip; // offset 56
} TaskContextT;
#else
// Saved CPU context for i386 (field order matches asm byte offsets)
typedef struct {
    uintptr_t ebx; // offset  0
    uintptr_t esi; // offset  4
    uintptr_t edi; // offset  8
    uintptr_t ebp; // offset 12
    uintptr_t esp; // offset 16
    uintptr_t eip; // offset 20
} TaskContextT;
#endif

// Task control block
typedef struct {
    char          name[TS_NAME_MAX];
    TaskContextT  context;
    uint8_t      *stack;
    uint32_t      stackSize;
    TaskStateE    state;
    int32_t       priority;
    int32_t       credits;
    TaskEntryT    entry;
    void         *arg;
    bool          isMain;
    bool          allocated; // true if slot is in use, false if free for reuse
} TaskBlockT;

// ============================================================================
// Module state
// ============================================================================

static TaskBlockT *tasks       = NULL; // stb_ds dynamic array
static uint32_t    currentIdx  = 0;
static bool        initialized = false;

// ============================================================================
// Forward declarations
// ============================================================================

// Static helpers
static void     contextSwitch(TaskContextT *save, TaskContextT *restore);
static int32_t  findFreeSlot(void);
static uint32_t scheduleNext(void);
static void     taskTrampoline(void);

// Public API prototypes are provided by taskswitch.h via #include.
// Explicit prototypes repeated here per project convention:
uint32_t   tsActiveCount(void);
int32_t    tsCreate(const char *name, TaskEntryT entry, void *arg, uint32_t stackSize, int32_t priority);
uint32_t   tsCurrentId(void);
void       tsExit(void);
const char *tsGetName(uint32_t taskId);
int32_t    tsGetPriority(uint32_t taskId);
TaskStateE tsGetState(uint32_t taskId);
int32_t    tsInit(void);
int32_t    tsKill(uint32_t taskId);
int32_t    tsPause(uint32_t taskId);
void       tsRecoverToMain(void);
int32_t    tsResume(uint32_t taskId);
int32_t    tsSetPriority(uint32_t taskId, int32_t priority);
void       tsShutdown(void);
void       tsYield(void);

// ============================================================================
// Static functions (alphabetical)
// ============================================================================

// Switch execution from the current task to another by saving and restoring
// callee-saved registers and the stack pointer.  The return address is
// captured as a local label so that when another task switches back to us,
// execution resumes right after the save point.
#if defined(__x86_64__)
// x86_64: save rbx, r12-r15, rbp, rsp, rip.
// Inputs via GCC constraints:  %rdi = save ptr,  %rsi = restore ptr.
static void __attribute__((noinline)) contextSwitch(TaskContextT *save, TaskContextT *restore) {
    __asm__ __volatile__(
        // Save current context
        "movq %%rbx,  0(%%rdi)\n\t"
        "movq %%r12,  8(%%rdi)\n\t"
        "movq %%r13, 16(%%rdi)\n\t"
        "movq %%r14, 24(%%rdi)\n\t"
        "movq %%r15, 32(%%rdi)\n\t"
        "movq %%rbp, 40(%%rdi)\n\t"
        "movq %%rsp, 48(%%rdi)\n\t"
        "leaq 1f(%%rip), %%rax\n\t"
        "movq %%rax, 56(%%rdi)\n\t"
        // Restore new context
        "movq  0(%%rsi), %%rbx\n\t"
        "movq  8(%%rsi), %%r12\n\t"
        "movq 16(%%rsi), %%r13\n\t"
        "movq 24(%%rsi), %%r14\n\t"
        "movq 32(%%rsi), %%r15\n\t"
        "movq 40(%%rsi), %%rbp\n\t"
        "movq 48(%%rsi), %%rsp\n\t"
        "movq 56(%%rsi), %%rax\n\t"
        "jmp  *%%rax\n\t"
        "1:\n\t"
        :
        : "D" (save), "S" (restore)
        : "rax", "rcx", "rdx", "r8", "r9", "r10", "r11", "memory", "cc"
    );
}
#else
// i386: save ebx, esi, edi, ebp, esp, eip.
// Inputs via GCC constraints:  %eax = save ptr,  %edx = restore ptr.
static void __attribute__((noinline)) contextSwitch(TaskContextT *save, TaskContextT *restore) {
    __asm__ __volatile__(
        // Save current context
        "movl %%ebx,  0(%%eax)\n\t"
        "movl %%esi,  4(%%eax)\n\t"
        "movl %%edi,  8(%%eax)\n\t"
        "movl %%ebp, 12(%%eax)\n\t"
        "movl %%esp, 16(%%eax)\n\t"
        "movl $1f,   20(%%eax)\n\t"
        // Restore new context
        "movl  0(%%edx), %%ebx\n\t"
        "movl  4(%%edx), %%esi\n\t"
        "movl  8(%%edx), %%edi\n\t"
        "movl 12(%%edx), %%ebp\n\t"
        "movl 16(%%edx), %%esp\n\t"
        "movl 20(%%edx), %%eax\n\t"
        "jmp  *%%eax\n\t"
        "1:\n\t"
        :
        : "a" (save), "d" (restore)
        : "ecx", "memory", "cc"
    );
}
#endif


// Find a free (terminated or unallocated) slot in the task array.
// Returns the index, or -1 if no free slot exists.
static int32_t findFreeSlot(void) {
    ptrdiff_t count = arrlen(tasks);
    for (ptrdiff_t i = 1; i < count; i++) {
        if (!tasks[i].allocated) {
            return (int32_t)i;
        }
    }
    return -1;
}


// Find the next task to run using credit-based weighted round-robin.
// Each ready task holds (priority + 1) credits.  One credit is consumed
// per scheduling turn.  When no ready task has credits left, every
// ready task is refilled.  This guarantees all tasks run while giving
// higher-priority tasks proportionally more turns.
static uint32_t scheduleNext(void) {
    uint32_t count = (uint32_t)arrlen(tasks);

    // First pass: look for a ready task with remaining credits
    for (uint32_t i = 1; i <= count; i++) {
        uint32_t idx = (currentIdx + i) % count;
        if (tasks[idx].allocated && tasks[idx].state == TaskStateReady && tasks[idx].credits > 0) {
            tasks[idx].credits--;
            return idx;
        }
    }

    // All credits exhausted -- refill every ready task
    bool anyReady = false;
    for (uint32_t i = 0; i < count; i++) {
        if (tasks[i].allocated && tasks[i].state == TaskStateReady) {
            tasks[i].credits = tasks[i].priority + 1;
            anyReady         = true;
        }
    }

    if (!anyReady) {
        return currentIdx;
    }

    // Pick the first ready task after refill
    for (uint32_t i = 1; i <= count; i++) {
        uint32_t idx = (currentIdx + i) % count;
        if (tasks[idx].allocated && tasks[idx].state == TaskStateReady && tasks[idx].credits > 0) {
            tasks[idx].credits--;
            return idx;
        }
    }

    return currentIdx;
}


// Entry point for every new task.  Calls the user-supplied function and
// then terminates the task when it returns.
static void taskTrampoline(void) {
    TaskBlockT *task = &tasks[currentIdx];
    task->entry(task->arg);
    tsExit();
}

// ============================================================================
// Public API (alphabetical, main-equivalent functions last if applicable)
// ============================================================================

uint32_t tsActiveCount(void) {
    if (!initialized) {
        return 0;
    }

    uint32_t active = 0;
    ptrdiff_t count = arrlen(tasks);
    for (ptrdiff_t i = 0; i < count; i++) {
        if (tasks[i].allocated && tasks[i].state != TaskStateTerminated) {
            active++;
        }
    }
    return active;
}


int32_t tsCreate(const char *name, TaskEntryT entry, void *arg, uint32_t stackSize, int32_t priority) {
    if (!initialized || !entry) {
        return TS_ERR_PARAM;
    }
    if (stackSize == 0) {
        stackSize = TS_DEFAULT_STACK_SIZE;
    }

    // Reuse a terminated/free slot, or append a new one
    int32_t id = findFreeSlot();
    if (id < 0) {
        TaskBlockT blank = {0};
        arrput(tasks, blank);
        id = (int32_t)(arrlen(tasks) - 1);
    }

    TaskBlockT *task = &tasks[id];
    memset(task, 0, sizeof(*task));

    task->stack = (uint8_t *)malloc(stackSize);
    if (!task->stack) {
        return TS_ERR_NOMEM;
    }

    if (name) {
        strncpy(task->name, name, TS_NAME_MAX - 1);
        task->name[TS_NAME_MAX - 1] = '\0';
    }

    task->stackSize = stackSize;
    task->state     = TaskStateReady;
    task->priority  = priority;
    task->credits   = priority + 1;
    task->entry     = entry;
    task->arg       = arg;
    task->isMain    = false;
    task->allocated = true;

    // Set up initial stack (grows downward, 16-byte aligned)
    uintptr_t top = (uintptr_t)(task->stack + stackSize);
    top &= ~(uintptr_t)0xF;
    top -= sizeof(uintptr_t);
    *(uintptr_t *)top = 0; // dummy return address; trampoline never returns

#if defined(__x86_64__)
    task->context.rsp = top;
    task->context.rbp = 0;
    task->context.rbx = 0;
    task->context.r12 = 0;
    task->context.r13 = 0;
    task->context.r14 = 0;
    task->context.r15 = 0;
    task->context.rip = (uintptr_t)taskTrampoline;
#else
    task->context.esp = top;
    task->context.ebp = 0;
    task->context.ebx = 0;
    task->context.esi = 0;
    task->context.edi = 0;
    task->context.eip = (uintptr_t)taskTrampoline;
#endif

    return id;
}


uint32_t tsCurrentId(void) {
    return currentIdx;
}


void tsExit(void) {
    if (!initialized || tasks[currentIdx].isMain) {
        return;
    }

    tasks[currentIdx].state = TaskStateTerminated;

    // Free the stack immediately
    free(tasks[currentIdx].stack);
    tasks[currentIdx].stack     = NULL;
    tasks[currentIdx].allocated = false;

    uint32_t next = scheduleNext();
    uint32_t prev = currentIdx;

    currentIdx        = next;
    tasks[next].state = TaskStateRunning;

    contextSwitch(&tasks[prev].context, &tasks[next].context);
    // Terminated task never resumes here
}


const char *tsGetName(uint32_t taskId) {
    if (!initialized || taskId >= (uint32_t)arrlen(tasks)) {
        return NULL;
    }
    if (!tasks[taskId].allocated) {
        return NULL;
    }
    return tasks[taskId].name;
}


int32_t tsGetPriority(uint32_t taskId) {
    if (!initialized || taskId >= (uint32_t)arrlen(tasks)) {
        return TS_ERR_PARAM;
    }
    if (!tasks[taskId].allocated) {
        return TS_ERR_PARAM;
    }
    return tasks[taskId].priority;
}


TaskStateE tsGetState(uint32_t taskId) {
    if (!initialized || taskId >= (uint32_t)arrlen(tasks)) {
        return TaskStateTerminated;
    }
    if (!tasks[taskId].allocated) {
        return TaskStateTerminated;
    }
    return tasks[taskId].state;
}


int32_t tsInit(void) {
    if (initialized) {
        return TS_ERR_PARAM;
    }

    // Start with the main task at slot 0
    TaskBlockT main = {0};
    strncpy(main.name, "main", TS_NAME_MAX - 1);
    main.state     = TaskStateRunning;
    main.priority  = TS_PRIORITY_NORMAL;
    main.credits   = TS_PRIORITY_NORMAL + 1;
    main.isMain    = true;
    main.stack     = NULL;
    main.allocated = true;

    arrput(tasks, main);

    currentIdx  = 0;
    initialized = true;
    return TS_OK;
}


int32_t tsKill(uint32_t taskId) {
    if (!initialized || taskId >= (uint32_t)arrlen(tasks)) {
        return TS_ERR_PARAM;
    }
    if (!tasks[taskId].allocated) {
        return TS_ERR_PARAM;
    }
    if (tasks[taskId].isMain) {
        return TS_ERR_STATE;
    }
    if (taskId == currentIdx) {
        return TS_ERR_STATE;
    }
    if (tasks[taskId].state == TaskStateTerminated) {
        return TS_ERR_STATE;
    }

    tasks[taskId].state = TaskStateTerminated;
    free(tasks[taskId].stack);
    tasks[taskId].stack     = NULL;
    tasks[taskId].allocated = false;

    return TS_OK;
}


int32_t tsPause(uint32_t taskId) {
    if (!initialized || taskId >= (uint32_t)arrlen(tasks)) {
        return TS_ERR_PARAM;
    }
    if (!tasks[taskId].allocated) {
        return TS_ERR_PARAM;
    }
    if (tasks[taskId].isMain) {
        return TS_ERR_STATE;
    }
    if (tasks[taskId].state != TaskStateRunning && tasks[taskId].state != TaskStateReady) {
        return TS_ERR_STATE;
    }

    tasks[taskId].state = TaskStatePaused;

    // If we paused ourselves, yield immediately
    if (taskId == currentIdx) {
        uint32_t next = scheduleNext();
        if (next != currentIdx) {
            uint32_t prev = currentIdx;

            currentIdx        = next;
            tasks[next].state = TaskStateRunning;

            contextSwitch(&tasks[prev].context, &tasks[next].context);
        }
    }

    return TS_OK;
}


void tsRecoverToMain(void) {
    if (!initialized) {
        return;
    }

    currentIdx        = 0;
    tasks[0].state = TaskStateRunning;
}


int32_t tsResume(uint32_t taskId) {
    if (!initialized || taskId >= (uint32_t)arrlen(tasks)) {
        return TS_ERR_PARAM;
    }
    if (!tasks[taskId].allocated) {
        return TS_ERR_PARAM;
    }
    if (tasks[taskId].state != TaskStatePaused) {
        return TS_ERR_STATE;
    }

    tasks[taskId].state   = TaskStateReady;
    tasks[taskId].credits = tasks[taskId].priority + 1;
    return TS_OK;
}


int32_t tsSetPriority(uint32_t taskId, int32_t priority) {
    if (!initialized || taskId >= (uint32_t)arrlen(tasks)) {
        return TS_ERR_PARAM;
    }
    if (!tasks[taskId].allocated) {
        return TS_ERR_PARAM;
    }
    if (tasks[taskId].state == TaskStateTerminated) {
        return TS_ERR_STATE;
    }

    tasks[taskId].priority = priority;
    tasks[taskId].credits  = priority + 1;
    return TS_OK;
}


void tsShutdown(void) {
    if (!initialized) {
        return;
    }

    ptrdiff_t count = arrlen(tasks);
    for (ptrdiff_t i = 0; i < count; i++) {
        free(tasks[i].stack);
    }

    arrfree(tasks);
    tasks       = NULL;
    currentIdx  = 0;
    initialized = false;
}


void tsYield(void) {
    if (!initialized) {
        return;
    }

    uint32_t next = scheduleNext();
    if (next == currentIdx) {
        return;
    }

    uint32_t prev = currentIdx;

    if (tasks[prev].state == TaskStateRunning) {
        tasks[prev].state = TaskStateReady;
    }

    currentIdx        = next;
    tasks[next].state = TaskStateRunning;

    contextSwitch(&tasks[prev].context, &tasks[next].context);
}