// 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.

#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;
} TaskBlockT;

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

static TaskBlockT *tasks        = NULL;
static uint32_t    taskCapacity = 0;
static uint32_t    taskCount    = 0;
static uint32_t    currentIdx   = 0;
static bool        initialized  = false;

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

// Static helpers
static void     contextSwitch(TaskContextT *save, TaskContextT *restore);
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(uint32_t maxTasks);
int32_t    tsPause(uint32_t taskId);
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 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) {
    // First pass: look for a ready task with remaining credits
    for (uint32_t i = 1; i <= taskCount; i++) {
        uint32_t idx = (currentIdx + i) % taskCount;
        if (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 < taskCount; i++) {
        if (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 <= taskCount; i++) {
        uint32_t idx = (currentIdx + i) % taskCount;
        if (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 count = 0;
    for (uint32_t i = 0; i < taskCount; i++) {
        if (tasks[i].state != TaskStateTerminated) {
            count++;
        }
    }
    return count;
}


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

    uint32_t    id   = taskCount;
    TaskBlockT *task = &tasks[id];

    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';
    } else {
        task->name[0] = '\0';
    }

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

    // 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

    taskCount++;
    return (int32_t)id;
}


uint32_t tsCurrentId(void) {
    return currentIdx;
}


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

    tasks[currentIdx].state = TaskStateTerminated;

    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 >= taskCount) {
        return NULL;
    }
    return tasks[taskId].name;
}


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


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


int32_t tsInit(uint32_t maxTasks) {
    if (initialized || maxTasks < 1) {
        return TS_ERR_PARAM;
    }

    tasks = (TaskBlockT *)calloc(maxTasks, sizeof(TaskBlockT));
    if (!tasks) {
        return TS_ERR_NOMEM;
    }

    taskCapacity = maxTasks;
    taskCount    = 1;
    currentIdx   = 0;

    strncpy(tasks[0].name, "main", TS_NAME_MAX - 1);
    tasks[0].state    = TaskStateRunning;
    tasks[0].priority = TS_PRIORITY_NORMAL;
    tasks[0].credits  = TS_PRIORITY_NORMAL + 1;
    tasks[0].isMain   = true;
    tasks[0].stack    = NULL;

    initialized = true;
    return TS_OK;
}


int32_t tsPause(uint32_t taskId) {
    if (!initialized || taskId >= taskCount) {
        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;
}


int32_t tsResume(uint32_t taskId) {
    if (!initialized || taskId >= taskCount) {
        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 >= taskCount) {
        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;
    }

    for (uint32_t i = 0; i < taskCount; i++) {
        free(tasks[i].stack);
    }

    free(tasks);

    tasks        = NULL;
    taskCapacity = 0;
    taskCount    = 0;
    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);
}