/********************************* Cooperative ********************************
This file is part of the PDI stack.

This is free software. You can redistribute it and/or modify it but without any
warranty.

Author          : Suraj I.
Created Date    : 1st June 2025
******************************************************************************/

#include "Cooperative.h"
#include "../DeviceControlInterface.h"

static XtensaContext __base_ctx; 
static bool __base_ctx_saved = false;

/**
 * Start entry function here
 */
void cooperative_trampoline(void *arg) {

    Cooperative* f = static_cast<Cooperative*>(arg);

    if (f){

        f->entry(f->arg);
        // Only mark finished if entry actually returns 
        if (f->state == CooperativeState::Running) { 
            f->state = CooperativeState::Finished; 
        }        
    }

    CooperativeScheduler::exit(); // never returns
}

/**
 * Cooperative destructor.
 */
Cooperative::~Cooperative(){
    pdiutil::safe_delete_array(stack_raw);
    stack = nullptr;
}

/**
 * CooperativeScheduler constructor.
 */
CooperativeScheduler::CooperativeScheduler(){
}

/**
 * CooperativeScheduler destructor.
 */
CooperativeScheduler::~CooperativeScheduler(){

    for (auto p : sleepers) { pdiutil::safe_delete(p.f); }
    for (auto p : ready) { pdiutil::safe_delete(p); }
    sleepers.clear();
    ready.clear();
}

/**
 * Schedule the task.
 */
int CooperativeScheduler::schedule_task(task_t* task, uint32_t stacksize){

    // Init new Cooperative
    // if(task->_task_exec) { delete task->_task_exec; }
    if(task->_task_mode != TASK_MODE_COOPERATIVE) return -1;
    Cooperative* f = pdiutil::safe_new<Cooperative>();
    if (!f) return -2;
    task->_task_exec = f;

    // Allocate with padding to guarantee a 16-byte aligned top-of-stack
    const uint32_t raw_bytes = stacksize + 64;
    pdiutil::safe_delete_array(f->stack_raw);
    f->stack_raw = pdiutil::safe_new_array<uint8_t>(raw_bytes);
    if (!f->stack_raw) {
        pdiutil::safe_delete(f);
        task->_task_exec = nullptr;
        return -2;
    }
    memset(f->stack_raw, 0, raw_bytes);

    // 16-byte aligned top
    uintptr_t top = reinterpret_cast<uintptr_t>(f->stack_raw + raw_bytes);
    top &= ~static_cast<uintptr_t>(0xF);

    // Record stack and SP
    f->stack      = reinterpret_cast<uint8_t*>(top);
    f->stack_size = stacksize;
    f->state      = CooperativeState::Ready;

    // Entry wrapper
    f->entry = [](void* arg){ 
        Cooperative* f = static_cast<Cooperative*>(arg); 
        if(f){
            // Serial.printf("[cooperative_entry_task] cooperative=%p\n", f);
            task_t* t = __task_scheduler.get_task(f->task_id);
            if(t && t->_task) t->_task();
        }
    };
    f->arg = static_cast<void*>(f);
    f->task_id = task->_task_id;

    // Initialize context for first restore
    // a1 (SP), sp mirror
    f->ctx.a[1] = static_cast<uint32_t>(top - 16);
    f->ctx.sp   = static_cast<uint32_t>(top - 16);

    // a2 = argument (Cooperative*)
    f->ctx.a[2] = reinterpret_cast<uint32_t>(f);

    // a0 = return address → exit (if trampoline ever returns)
    f->ctx.a[0] = reinterpret_cast<uint32_t>(&CooperativeScheduler::exit);

    // PS (start simple; UM|EXCM is typical — if it traps, try PS_EXCM only)
    uint32_t ps_now; asm volatile("rsr.ps %0" : "=a"(ps_now));
    f->ctx.ps = ps_now;
    // f->ctx.ps   = PS_UM | PS_EXCM;

    // PC = trampoline
    f->ctx.pc   = reinterpret_cast<uint32_t>(&cooperative_trampoline);

    // Clear the rest for hygiene (optional)
    for (int i = 3; i <= 15; ++i) f->ctx.a[i] = 0;

    // Add to ready list    
    CRITICAL_SECTION_ENTER
    add_to_ready(f);
    CRITICAL_SECTION_EXIT

    return 0;
}

/**
 * mute the current task.
 */
void CooperativeScheduler::mute(){

    CRITICAL_SECTION_ENTER 

    Cooperative* f = current;
    if (!f) {
        CRITICAL_SECTION_EXIT
        return;
    }

    if (f->state == CooperativeState::Running) {
        f->state = CooperativeState::Mute;
    }

    xtensa_save_context(&f->ctx);
    if(f->state == CooperativeState::Mute)
        CooperativeScheduler::exit();

    CRITICAL_SECTION_EXIT
}

/**
 * yield the current task.
 */
void CooperativeScheduler::yield(){

    CRITICAL_SECTION_ENTER 

    Cooperative* f = current;
    if (!f) {
        CRITICAL_SECTION_EXIT
        return;
    }

    if (f->state == CooperativeState::Running) {
        add_to_ready(f);
    }

    xtensa_save_context(&f->ctx);
    if(f->state == CooperativeState::Ready)
        CooperativeScheduler::exit();

    CRITICAL_SECTION_EXIT
}

/**
 * sleep current task.
 */
void CooperativeScheduler::sleep(uint32_t ms){

    CRITICAL_SECTION_ENTER 

    Cooperative* f = current;
    if (!f) {
        CRITICAL_SECTION_EXIT
        return;
    }

    f->state = CooperativeState::Sleeping;
    sleepers.push_back({ __i_dvc_ctrl.millis_now() + ms, current }); 

    xtensa_save_context(&f->ctx);
    if(f->state == CooperativeState::Sleeping)
        CooperativeScheduler::exit();

    CRITICAL_SECTION_EXIT
}

/**
 * Run the scheduled tasks.
 */
void CooperativeScheduler::run(){

    if (!__base_ctx_saved) { 

        __base_ctx_saved = true; 
        // Save the main loop context once every time run called
        CRITICAL_SECTION_ENTER
        xtensa_save_context(&__base_ctx); 
        CRITICAL_SECTION_EXIT
    }    

    // Wake sleepers whose time has arrived
    uint32_t now = __i_dvc_ctrl.millis_now();
    for (size_t i = 0; i < sleepers.size();) {

        auto si = sleepers[i];
        if (si.f && (int32_t)(now - si.wake_ms) >= 0) {

            CRITICAL_SECTION_ENTER 
            add_to_ready(si.f); 
            // erase by swap-pop for O(1)
            sleepers[i] = sleepers.back();
            sleepers.pop_back();
            CRITICAL_SECTION_EXIT
        } else {
            ++i;
        }
    }

    // If nothing is running, start one
    if (__base_ctx_saved && !current && !ready.empty()) {

        Cooperative* next = pick_next_ready();
        if (!next) return;

        CRITICAL_SECTION_ENTER
        current = next;
        next->state = CooperativeState::Running;
        xtensa_restore_context(&next->ctx); // never returns        
    }

    // __base_ctx_saved = false;
}

/**
 * Exit.
 */
void CooperativeScheduler::exit(){

    Cooperative* f = __i_cooperative_scheduler.current;
    if (f && f->state == CooperativeState::Finished){

        __i_cooperative_scheduler.destroy_cooperative(f);
    }

    CRITICAL_SECTION_ENTER
    __i_cooperative_scheduler.current = nullptr;
    if(__base_ctx_saved){
        __base_ctx_saved = false;
        xtensa_restore_context(&__base_ctx); // jumps back main
    }
    CRITICAL_SECTION_EXIT
}

/**
 * Destroy Cooperative
 */
void CooperativeScheduler::destroy_cooperative(Cooperative* f) {

    if (f){

        CRITICAL_SECTION_ENTER 
        remove_from_sleepers(f);
        remove_from_ready(f);
    
        __task_scheduler.remove_task(f->task_id);
        if(current == f) current = nullptr;
        pdiutil::safe_delete(f);
        CRITICAL_SECTION_EXIT
    }
}

/**
 * Add Cooperative to ready queue
 */
void CooperativeScheduler::add_to_ready(Cooperative* f) {
    if(f){
        // f->wait_ticks = 0;
        f->state = PreemptiveState::Ready;
        ready.push_back(f);
    }
}

/**
 * Remove Cooperative from ready queue
 */
void CooperativeScheduler::remove_from_ready(Cooperative* f) {

    for (size_t i = 0; i < ready.size(); ++i) {
        if (ready[i] == f) { ready[i] = ready.back(); ready.pop_back(); break; }
    }
}

/**
 * Remove Cooperative from sleeper queue
 */
void CooperativeScheduler::remove_from_sleepers(Cooperative* f) {

    for (size_t i = 0; i < sleepers.size();) {
        if (sleepers[i].f == f) { sleepers[i] = sleepers.back(); sleepers.pop_back(); }
        else { ++i; }
    }
}

/**
 * Pick the next ready Cooperative
 */
Cooperative* CooperativeScheduler::pick_next_ready() {

    Cooperative* best = nullptr;
    int bestScore = -1;

    for (int i = 0; i < ready.size(); ++i) {

        Cooperative* f = ready[i];

        task_t* t = __task_scheduler.get_task(f->task_id);
        if (!t) continue;

        // Aging: score(effective priority) = base_priority + wait_ticks
        int score = t->_task_priority + f->wait_ticks;

        if (score > bestScore) {
            bestScore = score;
            best = f;
        }

        // Increment wait_ticks for all. max boost 255
        if (f->wait_ticks < 255) f->wait_ticks++;
    }

    if (nullptr != best) {
        // Reset wait_ticks for the chosen Cooperative
        best->wait_ticks = 0;

        // Remove from ready list
        CRITICAL_SECTION_ENTER
        remove_from_ready(best);
        CRITICAL_SECTION_EXIT
    }

    return best;
}

/**
 * Yield other scheduler to manage pending work in this scheduler
 */
void CooperativeScheduler::yield_from_othersched() {

    if (!can_yield_from_othersched()) return; // reject if already active
    othersched_active = true;

    // capture caller context as base to switch back and return control back to caller
    __base_ctx_saved = false;
    run();
    // safer side if exit of this scheduler missed
    __base_ctx_saved = false;

    othersched_active = false;
}

/**
 * Sleep other scheduler to manage pending work in this scheduler
 */
void CooperativeScheduler::sleep_from_othersched(uint32_t ms) {

    if (!can_sleep_from_othersched()) return; // reject if already active
    othersched_active = true;

    uint32_t exitms = __i_dvc_ctrl.millis_now() + ms;

    while ((int32_t)(__i_dvc_ctrl.millis_now() - exitms) < 0)
    {
        // capture caller context as base to switch back and return control back to caller
        __base_ctx_saved = false;
        run();
        // yield device level functionalities
        __i_dvc_ctrl.yield();
    }
    
    // safer side if exit of this scheduler missed
    __base_ctx_saved = false;

    othersched_active = false;
}

CooperativeScheduler __i_cooperative_scheduler;
