/**
  **********************************************************************************************************************
  * @file    cot_os.c
  * @brief   该文件提供查询协作式多任务系统功能
  * @author  const_zpc  any question please send mail to const_zpc@163.com
  * @version V1.0.0
  * @date    2024-02-04
  *
  * @details  功能详细说明：
  *           + 任务调度初始化
  *           + 任务调度功能启动
  *           + 该任务调度并非抢占式
  *
  **********************************************************************************************************************
  * 源码路径：https://gitee.com/cot_package/cot_os.git 具体问题及建议可在该网址填写 Issue
  *
  * 使用方式:
  *    1. 使用前初始化函数 cotOs_Init
  *    2. 通过 cotOs_Creat 添加任务函数
  *    3. 主函数调用 cotOs_Start 启动任务调度, 函数不会退出
  *    4. 使用 cotOs_Wait 切换到下一个任务执行
  *
  **********************************************************************************************************************
  */

/* Includes ----------------------------------------------------------------------------------------------------------*/

#pragma GCC push_options
#pragma GCC optimize("O0")
#include <stdbool.h>
#include "cot_os_config.h"
#include <setjmp.h>
// #pragma GCC pop_options
#include "cot_os.h"

/* Private typedef ---------------------------------------------------------------------------------------------------*/

#if (COT_OS_MAX_TASK > 32)
#error "task max num can't over 32"
#endif

typedef struct stTCB
{
    uint8_t *pBakStack;
    uint8_t state : 2;
    uint8_t joinWaitCnt :6;
    CotOSCondition_t cond;
    CotOSStackType_e eStackType;
    uint32_t nextRunTime;
    int param;
    jmp_buf env;
    size_t bakStackSize;
    size_t stackTopAddr;
    cotOsTask_f pfnOsTaskEnter;
    CotOSCondition_t *pCondition;
    struct stTCB *pNext;
} TCB_t;

/**
  * @brief 中断任务信息结构体定义
  */
typedef struct
{
    cotOsGetSysTime_f pfnGetTimerMs;                /*!< 获取毫秒级别的时间回调函数 */

    TCB_t *pCurTCB;

    TCB_t *pTCBList;

    uint32_t tcbMask;

    jmp_buf env;

    jmp_buf envResume;

    uint8_t *pFreeBakStack;
} OsInfo_t;

/* Private define ----------------------------------------------------------------------------------------------------*/

#define COMMON_TASK_INTI                0
#define COMMON_TASK_RUN_NEXT            1
#define COMMON_TASK_RUN                 2

#define MAIN_TASK_INTI                  0
#define MAIN_TASK_EXIT                  1

#define BACKUP_TASK_INTI                0
#define BACKUP_TASK_RUN                 1

#define RESUME_TASK_INTI                0
#define RESUME_TASK_RUN                 1

#define TASK_STATUS_READY               0
#define TASK_STATUS_RUNNING             1
#define TASK_STATUS_SUSPEND             2
#define TASK_STATUS_DELETED             3

/* Private macro -----------------------------------------------------------------------------------------------------*/
/* Private variables -------------------------------------------------------------------------------------------------*/

static OsInfo_t sg_OsInfo;
static TCB_t sg_TCB[COT_OS_MAX_TASK];

static uint8_t sg_sharedTaskStack[COT_OS_SHARED_STACK_BAK_SIZE];
static uint32_t sg_copyStack[512];

/* Private function prototypes ---------------------------------------------------------------------------------------*/

static TCB_t *CreatTCB(OsInfo_t *pObj);
static void DestroyTCB(OsInfo_t *pObj);
static void AddToTCBTaskList(OsInfo_t *pObj, TCB_t *pNewTCB);
static void DeleteFromTCBTaskList(OsInfo_t *pObj);
static void BackupCurTaskStack(void);
static void ResumeCurTaskStack(void);
static void JumpNextTask(OsInfo_t *pObj);

/* Private function --------------------------------------------------------------------------------------------------*/
static void tcbMemcpy(uint8_t *pdest, uint8_t *psrc, size_t length)
{
    while (length--)
    {
        *pdest = *psrc;
        pdest++;
        psrc++;
    }
}

static TCB_t *CreatTCB(OsInfo_t *pObj)
{
    for (int i = 0; i < COT_OS_MAX_TASK; i++)
    {
        if (!((pObj->tcbMask >> i) & 0x01))
        {
            pObj->tcbMask |= (0x1 << i);
            return &sg_TCB[i];
        }
    }

    return NULL;
}

static void DestroyTCB(OsInfo_t *pObj)
{
    for (int i = 0; i < COT_OS_MAX_TASK; i++)
    {
        if (&sg_TCB[i] == pObj->pCurTCB)
        {
            pObj->tcbMask &= ~(0x1 << i);
            break;
        }
    }
}

static void AddToTCBTaskList(OsInfo_t *pObj, TCB_t *pNewTCB)
{
    if (pObj->pTCBList == NULL)
    {
        pObj->pTCBList = pNewTCB;
        pObj->pTCBList->pNext = pObj->pTCBList;
    }
    else
    {
        TCB_t *pTCB = pObj->pTCBList;

        while (pTCB->pNext != NULL && pTCB->pNext != pObj->pTCBList)
        {
            pTCB = pTCB->pNext;
        }

        pNewTCB->pNext = pTCB->pNext;
        pTCB->pNext = pNewTCB;
    }
}

static void DeleteFromTCBTaskList(OsInfo_t *pObj)
{
    TCB_t *pTCB = pObj->pTCBList;

    if (pTCB->pNext == pObj->pTCBList)
    {
        pObj->pTCBList = NULL;
    }
    else
    {
        while (pTCB->pNext != pObj->pTCBList)
        {
            if (pTCB->pNext == pObj->pCurTCB)
            {
                break;
            }

            pTCB = pTCB->pNext;
        }

        pTCB->pNext = pObj->pCurTCB->pNext;

        if (pObj->pCurTCB == pObj->pTCBList)
        {
            pObj->pTCBList = pObj->pCurTCB->pNext;
        }
    }
}

static uint8_t GetTaskNum(OsInfo_t *pObj)
{
    uint8_t num = 0;
    TCB_t *pTCB = sg_OsInfo.pTCBList;

    if (pTCB == NULL)
    {
        return 0;
    }

    do {
        if (pTCB->state != TASK_STATUS_DELETED)
        {
            num++;
        }

        pTCB = pTCB->pNext;
    } while (pTCB != sg_OsInfo.pTCBList);

    return num;
}

/**
  * @brief      初始化协程OS系统
  * 
  * @attention  回调函数不可为NULL
  * @param      pfOsTimer 获取毫秒级别的时间回调函数
  */
void cotOs_Init(cotOsGetSysTime_f pfnOsTimer)
{
    sg_OsInfo.pfnGetTimerMs = pfnOsTimer;
    sg_OsInfo.pTCBList = NULL;
    sg_OsInfo.pCurTCB = NULL;
    sg_OsInfo.tcbMask = 0;
    sg_OsInfo.pFreeBakStack = sg_sharedTaskStack;
}

static void MallocBakStack(OsInfo_t *pObj, size_t length)
{
    if ((&sg_sharedTaskStack[COT_OS_SHARED_STACK_BAK_SIZE] - pObj->pFreeBakStack) < length)
    {
        while (1){}  // 剩余共享栈任务的共享备份栈内存大小不足够申请备份
    }

    sg_OsInfo.pCurTCB->pBakStack = sg_OsInfo.pFreeBakStack;
    sg_OsInfo.pCurTCB->bakStackSize = length;
    sg_OsInfo.pFreeBakStack += length;
}

static void FreeBakStack(OsInfo_t *pObj)
{
    if (pObj->pCurTCB->pBakStack != NULL)
    {
        tcbMemcpy(pObj->pCurTCB->pBakStack, pObj->pCurTCB->pBakStack + pObj->pCurTCB->bakStackSize, 
                    pObj->pFreeBakStack - (pObj->pCurTCB->pBakStack + pObj->pCurTCB->bakStackSize));

        TCB_t *pTCB = pObj->pTCBList;

        do {
            if (pTCB != NULL && pTCB->pBakStack > pObj->pCurTCB->pBakStack)
            {
                pTCB->pBakStack -= pObj->pCurTCB->bakStackSize;
            }

            pTCB = pTCB->pNext;
        } while (pTCB != sg_OsInfo.pTCBList);

        pObj->pFreeBakStack -= pObj->pCurTCB->bakStackSize;
        pObj->pCurTCB->pBakStack = NULL;
        pObj->pCurTCB->bakStackSize = 0;
    }
}

// #pragma GCC push_options
// #pragma GCC optimize("O0")

// 在备份运行栈中备份当前任务栈的数据
static void BackupCurTaskStack(void)
{
    volatile size_t oldsp;

    COT_OS_GET_STACK(oldsp);

    if (sg_OsInfo.pCurTCB->eStackType == COT_OS_SHARED_STACK &&
        sg_OsInfo.pCurTCB->pBakStack == NULL)
    {
        MallocBakStack(&sg_OsInfo, sg_OsInfo.pCurTCB->stackTopAddr - oldsp);
    }

    COT_OS_SET_STACK(((size_t)sg_copyStack + sizeof(sg_copyStack)));

    if (BACKUP_TASK_RUN != setjmp(sg_OsInfo.envResume))
    {
        COT_OS_SET_STACK(oldsp);
        longjmp(sg_OsInfo.envResume, BACKUP_TASK_RUN);
    }
    else
    {
        /* 在新的栈中进行拷贝 */
        if (sg_OsInfo.pCurTCB->eStackType == COT_OS_SHARED_STACK)
        {
            tcbMemcpy(sg_OsInfo.pCurTCB->pBakStack, (uint8_t *)(sg_OsInfo.pCurTCB->stackTopAddr - sg_OsInfo.pCurTCB->bakStackSize), sg_OsInfo.pCurTCB->bakStackSize);
        }

        longjmp(sg_OsInfo.pCurTCB->env, COMMON_TASK_RUN_NEXT);
    }
}

// 在备份运行栈中还原当前任务栈的数据
static void ResumeCurTaskStack(void)
{
    volatile size_t oldsp;

    COT_OS_GET_STACK(oldsp);
    COT_OS_SET_STACK(((size_t)sg_copyStack + sizeof(sg_copyStack)));

    if (RESUME_TASK_RUN != setjmp(sg_OsInfo.envResume))
    {
        COT_OS_SET_STACK(oldsp);
        longjmp(sg_OsInfo.envResume, RESUME_TASK_RUN);
    }
    else
    {
        /* 在新的栈中进行拷贝并完成下一个任务的切换，避免共享栈之间切换任务时拷贝后被破坏 */
        if (sg_OsInfo.pCurTCB->eStackType == COT_OS_SHARED_STACK)
        {
            tcbMemcpy((uint8_t *)(sg_OsInfo.pCurTCB->stackTopAddr - sg_OsInfo.pCurTCB->bakStackSize), sg_OsInfo.pCurTCB->pBakStack, sg_OsInfo.pCurTCB->bakStackSize);
        }

        longjmp(sg_OsInfo.pCurTCB->env, COMMON_TASK_RUN);
    }
}

static void DestoryTask(OsInfo_t *pObj)
{
    if (pObj->pCurTCB->state == TASK_STATUS_DELETED && pObj->pCurTCB->joinWaitCnt == 0)
    {
        FreeBakStack(&sg_OsInfo);
        DeleteFromTCBTaskList(&sg_OsInfo);
        DestroyTCB(&sg_OsInfo);
    }
}

/**
  * @brief      创建任务
  * 
  * @attention  若选择共享栈，则使用`cotOs_Wait`和`cotOs_ConditionWait`仅限于在任务函数中使用，禁止在任务函数中的嵌套函数中使用
  * @attention  同一个栈空间不能同时设置为共享栈和独立栈，即共享栈任务和独立栈任务不能使用相同的栈空间
  * @note       独立栈任务必须有不同的栈空间，而共享栈任务可以有相同的栈空间
  * @param      pfnOsTaskEnter 任务函数
  * @param      eStackType     栈类型选择
  * @param      pStack         栈空间
  * @param      stackSize      栈空间大小
  * @param      arg            参数
  * @return     任务句柄, 为NULL则创建失败
  */
cotOsTask_t cotOs_CreatTask(cotOsTask_f pfnOsTaskEnter, CotOSStackType_e eStackType, void *pStack, size_t stackSize, int arg)
{
    size_t oldsp;
    TCB_t *pNewTCB;

    if (pStack == NULL || stackSize == 0 || sg_OsInfo.pfnGetTimerMs == NULL)
    {
        return NULL;
    }

    COT_OS_GET_STACK(oldsp);

    pNewTCB = CreatTCB(&sg_OsInfo);

    if (NULL != pNewTCB)
    {
        /* 解决在共享栈中使用同一个共享栈创建任务时异常的问题 */
        if (sg_OsInfo.pCurTCB != NULL && sg_OsInfo.pCurTCB->eStackType == COT_OS_SHARED_STACK)
        {
            if (COMMON_TASK_RUN_NEXT != setjmp(sg_OsInfo.pCurTCB->env))
            {
                BackupCurTaskStack();
            }
        }

        COT_OS_SET_STACK(((size_t)pStack + stackSize));

        if (COMMON_TASK_INTI == setjmp(pNewTCB->env))
        {
            COT_OS_SET_STACK(oldsp);
            pNewTCB->pfnOsTaskEnter = pfnOsTaskEnter;
            pNewTCB->param = arg;
            pNewTCB->stackTopAddr = (size_t)pStack + stackSize;
            pNewTCB->eStackType = eStackType;
            pNewTCB->pNext = NULL;
            pNewTCB->state = TASK_STATUS_READY;
            pNewTCB->bakStackSize = 0;
            pNewTCB->pBakStack = NULL;
            pNewTCB->pCondition = NULL;
            pNewTCB->nextRunTime = 0;
            pNewTCB->joinWaitCnt = 0;
            cotOs_ConditionInit(&pNewTCB->cond);

            AddToTCBTaskList(&sg_OsInfo, pNewTCB);
        }
        else
        {
            sg_OsInfo.pCurTCB->state = TASK_STATUS_RUNNING;
            sg_OsInfo.pCurTCB->pfnOsTaskEnter(sg_OsInfo.pCurTCB->param);
            sg_OsInfo.pCurTCB->state = TASK_STATUS_DELETED;
            cotOs_ConditionNotify(&sg_OsInfo.pCurTCB->cond);
            DestoryTask(&sg_OsInfo);

            if (GetTaskNum(&sg_OsInfo) > 0)
            {
                JumpNextTask(&sg_OsInfo);
            }
            else
            {
                longjmp(sg_OsInfo.env, MAIN_TASK_EXIT);
            }
        }

        /* 解决在共享栈中使用同一个共享栈创建任务时异常的问题 */
        if (sg_OsInfo.pCurTCB != NULL && sg_OsInfo.pCurTCB->eStackType == COT_OS_SHARED_STACK)
        {
            if (COMMON_TASK_RUN != setjmp(sg_OsInfo.pCurTCB->env))
            {
                ResumeCurTaskStack();
            }
        }
    }

    return pNewTCB;
}

/**
  * @brief      启动协程OS任务
  * 
  * @attention  该函数不会退出, 因此一切准备就绪后最后调用该函数即可
  * @return     0,成功; -1,失败
  */
int cotOs_Start(void)
{
    if (sg_OsInfo.pTCBList == NULL || sg_OsInfo.pfnGetTimerMs == NULL)
    {
        return -1;
    }

    if (MAIN_TASK_EXIT != setjmp(sg_OsInfo.env))
    {
        sg_OsInfo.pCurTCB = sg_OsInfo.pTCBList;
        longjmp(sg_OsInfo.pCurTCB->env, COMMON_TASK_RUN);
    }

    return 0;
}

static void JumpNextTask(OsInfo_t *pObj)
{
    if (pObj->pTCBList != NULL)
    {
        while (1)
        {
            TCB_t *pTCB = pObj->pTCBList;

            do {
                if (pTCB->state != TASK_STATUS_DELETED &&
                    ((pTCB->nextRunTime != 0 && pObj->pfnGetTimerMs() > pTCB->nextRunTime) || 
                    (pTCB->pCondition != NULL && pTCB->pCondition->flag == 1)))
                {
                    pTCB->state = TASK_STATUS_READY;
                }

                pTCB = pTCB->pNext;
            } while (pTCB != pObj->pTCBList);

            do {
                if (pTCB->pCondition != NULL)
                    pTCB->pCondition->flag = 0;

                pTCB = pTCB->pNext;
            } while (pTCB != pObj->pTCBList);

            pTCB = pObj->pCurTCB->pNext;

            do {
                if (pTCB->state == TASK_STATUS_READY)
                {
                    pObj->pCurTCB = pTCB;
                    pObj->pCurTCB->state = TASK_STATUS_RUNNING;
                    ResumeCurTaskStack();
                }

                pTCB = pTCB->pNext;
            } while (pTCB != pObj->pCurTCB->pNext);
        }
    }
}

/**
  * @brief      任务阻塞等待时间结束
  * 
  * @note       等待n毫秒时长, 则会在 >= n 后的某个时刻才会运行, 具体要看其他协程任务什么时候让出资源
  * @param      time 等待时长，单位毫秒
  */
void cotOs_Wait(uint32_t time)
{
    sg_OsInfo.pCurTCB->nextRunTime = sg_OsInfo.pfnGetTimerMs() + time;
    sg_OsInfo.pCurTCB->pCondition = NULL;
    sg_OsInfo.pCurTCB->state = TASK_STATUS_SUSPEND;

    int ret = setjmp(sg_OsInfo.pCurTCB->env);

    if (COMMON_TASK_INTI == ret)
    {
        BackupCurTaskStack();
    }
    else if (COMMON_TASK_RUN_NEXT == ret)
    {
        JumpNextTask(&sg_OsInfo);
    }
}

/**
  * @brief      等待任务退出
  * 
  * @param      task cotOs_CreatTask 返回的值
  */
void cotOs_Join(cotOsTask_t task)
{
    TCB_t *pTaskTCB = (TCB_t *)task;
    TCB_t *pTCB = sg_OsInfo.pTCBList;

    if (NULL != pTaskTCB)
    {
        do {
            if (pTCB == pTaskTCB && pTaskTCB->state != TASK_STATUS_DELETED)
            {
                pTCB->joinWaitCnt++;
                cotOs_ConditionWait(&pTaskTCB->cond);
                pTCB->joinWaitCnt--;
                DestoryTask(&sg_OsInfo);
                break;
            }

            pTCB = pTCB->pNext;
        } while (pTCB != sg_OsInfo.pTCBList);
    }
}

/**
  * @brief      条件事件初始化
  * 
  * @param      pCondition 条件事件
  */
void cotOs_ConditionInit(CotOSCondition_t *pCondition)
{
    if (pCondition == NULL)
    {
        return;
    }

    pCondition->flag = 0;
}

/**
  * @brief      任务阻塞等待条件事件发生
  * 
  * @note       等待条件事件发生, 则会在条件事件发生后的某个时刻才会运行, 具体要看其他协程任务什么时候让出资源
  * @param      pCondition 条件事件
  */
void cotOs_ConditionWait(CotOSCondition_t *pCondition)
{
    if (pCondition == NULL)
    {
        return;
    }

    sg_OsInfo.pCurTCB->nextRunTime = 0;
    sg_OsInfo.pCurTCB->pCondition = pCondition;
    sg_OsInfo.pCurTCB->state = TASK_STATUS_SUSPEND;

    int ret = setjmp(sg_OsInfo.pCurTCB->env);

    if (COMMON_TASK_INTI == ret)
    {
        BackupCurTaskStack();
    }
    else if (COMMON_TASK_RUN_NEXT == ret)
    {
        JumpNextTask(&sg_OsInfo);
    }
}

/**
  * @brief      触发条件事件
  * 
  * @param      pCondition 条件事件
  */
void cotOs_ConditionNotify(CotOSCondition_t *pCondition)
{
    if (pCondition == NULL)
    {
        return;
    }

    pCondition->flag = 1;
}

#pragma GCC pop_options