#include "HydraTaskScheduler.h"
#include "HydraThread.h"
#include "../task/HydraTask.h"


void HydraTaskScheduler::Initialise(int thread_count)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    HydraID id = 0;
    for (int i = 0; i < MAX_TASKS; i++)
    {
        m_available_barrier_ids.push(id);
        id++;
    }
    m_barriers.resize(MAX_TASKS);

    for (int i = 0; i < thread_count; i++)
    {
        HydraThread *pthread = new HydraThread();
        pthread->ThreadName = "THREAD_" + std::to_string(i);       
        m_general_thread.push_back(pthread);
    }

    HydraThread *render_thread = new HydraThread();
    render_thread->ThreadName = "RENDER";
    render_thread->ThreadAffinity = HydraThreadAffinity::Render;
    m_general_thread.push_back(render_thread);
    

    HydraThread *physics_thread = new HydraThread();
    physics_thread->ThreadName = "PHYSICS";
    physics_thread->ThreadAffinity = HydraThreadAffinity::Physics;
    m_general_thread.push_back(physics_thread);
    

    HydraThread *main_thread = new HydraThread();
    main_thread->ThreadAffinity = HydraThreadAffinity::General;
    main_thread->ThreadID = std::this_thread::get_id();
    main_thread->ThreadName = "MAIN";
    m_general_thread.push_back(main_thread);

    m_tasks.resize(MAX_TASKS);
    m_task_ids_waiting_for_children.resize(MAX_TASKS);

    for (size_t i = 0; i < MAX_TASKS; i++)
    {
        m_avaialble_task_ids.push_back(i);
    }
    render_thread->Start();
    physics_thread->Start();
    for (int i = 0; i < thread_count; i++)
    {
        m_general_thread[i]->Start();
    }
    
}

void HydraTaskScheduler::Shutdown()
{
    std::lock_guard<std::mutex> lock(m_mutex);
    for (int i = 0; i < m_general_thread.size(); i++)
    {
        m_general_thread[i]->Stop();
    }
}

HydraID HydraTaskScheduler::_SaveTask(HydraTask *task)
{

    if (m_avaialble_task_ids.size() > 0)
    {
        HydraID id = m_avaialble_task_ids.front();
        m_avaialble_task_ids.pop_front();

        if(id > m_tasks.size())
        {
            DebugLog("INVALID TASK");
        }
        task->m_taskID = id;

        m_tasks[id] = task;

        return id;
    }
    else
    {
        DebugLog("Create Task Fail");
    }

    return INVALID_HYDRA_ID;
}

void HydraTaskScheduler::_DeleteTask(HydraID taskID)
{
    if ((taskID <= MAX_TASKS) && (taskID != INVALID_HYDRA_ID))
    {
        if (m_tasks[taskID]->m_notifyThread)
        {
            _NotifyThread(m_tasks[taskID]->m_notifyID);
        }

        delete m_tasks[taskID];
        m_tasks[taskID] = nullptr;
        m_avaialble_task_ids.push_back(taskID);
    }
}

void HydraTaskScheduler::ProcessTask()
{
    HydraThread *thread = nullptr;
    for (int i = 0; i < m_general_thread.size(); i++)
    {
        if (m_general_thread[i]->ThreadID == std::this_thread::get_id())
        {
            thread = m_general_thread[i];
            break;
        }
    }
    if (thread != nullptr)
    {
        HydraID taskID = INVALID_HYDRA_ID;
        {
            std::lock_guard<std::mutex> lock(m_mutex);
            taskID = _GetTaskReadyToRun(thread->ThreadAffinity);
        }

        if (taskID != INVALID_HYDRA_ID)
        {
            HydraTask *task = GetTask(taskID);
            if (task != nullptr)
            {
                task->_RunTask();
            }
        }
    }
}

void HydraTaskScheduler::_NotifyThread(std::thread::id thread_id)
{
    HydraThread *thread = nullptr;

    for (int i = 0; i < m_general_thread.size(); i++)
    {
        if (m_general_thread[i]->ThreadID == thread_id)
        {
            thread = m_general_thread[i];
            break;
        }
    }
    thread->WaitCondition.notify_one();
}

HydraTask *const HydraTaskScheduler::_GetTask(HydraID taskID)
{

    if (taskID <= MAX_TASKS)
    {
        if (m_tasks[taskID] != nullptr)
        {
            return m_tasks[taskID];
        }
    }
    return nullptr;
}
HydraTask *const HydraTaskScheduler::GetTask(HydraID taskID)
{    
    return _GetTask(taskID);
}

void HydraTaskScheduler::StartTask(HydraID taskID)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    _StartTask(taskID);
}

void HydraTaskScheduler::_StartTask(HydraID taskID, bool notifyThread)
{

    if ((taskID <= MAX_TASKS) && (taskID != INVALID_HYDRA_ID))
    {
        HydraTask *task = m_tasks[taskID];
        task->m_notifyThread = true;
        task->m_notifyID = std::this_thread::get_id();
        if (task != nullptr)
        {
            if (task->ThreadAffinity == HydraThreadAffinity::General)
            {
                m_general_tasks_ready_to_run.push_back(taskID);
            }
            else if (task->ThreadAffinity == HydraThreadAffinity::Render)
            {
                m_render_tasks_ready_to_run.push_back(taskID);
            }
            else if (task->ThreadAffinity == HydraThreadAffinity::Physics)
            {
                m_physics_tasks_ready_to_run.push_back(taskID);
            }
        }
    }
}

/// @brief Tell the task scheduler that the thread is about to start processing. This allows the task to perform any setup if required.
/// @param taskID 
void HydraTaskScheduler::NotifyStart(HydraID taskID)
{
    std::lock_guard<std::mutex> lock(m_mutex);
}

void HydraTaskScheduler::NotifyComplete(HydraID taskID)
{
    m_mutex.lock();
    if (taskID != INVALID_HYDRA_ID)
    {
        if (taskID < MAX_TASKS)
        {

            if (m_tasks[taskID] != nullptr)
            {
                HydraTask *task = m_tasks[taskID];
                if (task->m_child_tasks != 0)
                {
                    m_task_ids_waiting_for_children[taskID] = taskID;
                    m_mutex.unlock();
                    return;
                }
                else
                {
                    m_mutex.unlock();
                    task->Complete();                    
                    m_mutex.lock();
                }

                HydraID parentTaskID = task->m_dependant_task;
                if (task->m_barrierID != INVALID_HYDRA_ID)
                {
                    _DecrementBarrier(task->m_taskID, task->m_barrierID);
                }
                _DeleteTask(taskID);
                if (parentTaskID != INVALID_HYDRA_ID)
                {
                    _UpdateDependantTaskComplete(parentTaskID);
                }
            }
        }
    }
    m_mutex.unlock();
}

void HydraTaskScheduler::_UpdateDependantTaskComplete(HydraID taskID)
{
    if ((taskID <= MAX_TASKS) && (taskID != INVALID_HYDRA_ID))
    {
        if (m_tasks[taskID] != nullptr)
        {
            m_tasks[taskID]->m_child_tasks--;
            if (m_tasks[taskID]->m_child_tasks == 0)
            {
                m_tasks[taskID]->Notify();
                if (m_tasks[taskID]->m_dependant_task != INVALID_HYDRA_ID)
                {
                    _UpdateDependantTaskComplete(m_tasks[taskID]->m_dependant_task);
                }
                m_task_ids_waiting_for_children[taskID] = INVALID_HYDRA_ID;
                if (m_tasks[taskID]->m_barrierID != INVALID_HYDRA_ID)
                {
                    _DecrementBarrier(m_tasks[taskID]->m_taskID, m_tasks[taskID]->m_barrierID);
                }
                _DeleteTask(taskID);
            }
        }
    }
}

HydraID HydraTaskScheduler::_GetRenderTaskReadyToRun()
{

    if (m_render_tasks_ready_to_run.size() > 0)
    {
        HydraID taskID = m_render_tasks_ready_to_run.front();
        m_render_tasks_ready_to_run.pop_front();
        return taskID;
    }
    return INVALID_HYDRA_ID;
}

HydraID HydraTaskScheduler::_GetPhysicsTaskReadyToRun()
{

    if (m_physics_tasks_ready_to_run.size() > 0)
    {
        HydraID taskID = m_physics_tasks_ready_to_run.front();
        m_physics_tasks_ready_to_run.pop_front();
        return taskID;
    }
    return INVALID_HYDRA_ID;
}

HydraID HydraTaskScheduler::_GetTaskReadyToRun(HydraThreadAffinity thread_affinity)
{
    HydraID taskID = INVALID_HYDRA_ID;
    if (thread_affinity == HydraThreadAffinity::Render)
    {
        if (m_render_tasks_ready_to_run.size() > 0)
        {
            taskID = m_render_tasks_ready_to_run.front();
            m_render_tasks_ready_to_run.pop_front();
        }
    }
    else if (thread_affinity == HydraThreadAffinity::Physics)
    {
        if (m_physics_tasks_ready_to_run.size() > 0)
        {
            taskID = m_physics_tasks_ready_to_run.front();
            m_physics_tasks_ready_to_run.pop_front();
        }
    }
    else
    {
        if (taskID == INVALID_HYDRA_ID)
        {
            if (m_general_tasks_ready_to_run.size() > 0)
            {
                taskID = m_general_tasks_ready_to_run.front();
                m_general_tasks_ready_to_run.pop_front();
            }
        }
    }
    return taskID;
}

/* @brief The calling thread will execute the specified task and will not 
return until the 
@param taskID*/ 
void HydraTaskScheduler::WaitForTask(HydraID taskID)
{
    HydraTask *task = nullptr;
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        task = GetTask(taskID);

        if (task->m_barrierID == INVALID_HYDRA_ID)
        {
            HydraID barrierID = _CreateBarrier();
            _AddTaskToBarrier(taskID, barrierID);
        }

        _StartTask(taskID);
    }
    //HydraTask *task = GetTask(taskID);
    if(task != nullptr)
    {
        WaitForBarrier(task->m_barrierID);
    }
    else
    {
        DebugLog("Invlalid Task");
    }
}
HydraID HydraTaskScheduler::_CreateBarrier()
{
    if (m_available_barrier_ids.size() > 0)
    {
        HydraID id = m_available_barrier_ids.front();
        m_available_barrier_ids.pop();
        m_barriers[id] = new HydraTaskBarrier();
        m_barriers[id]->barrier_id = id;
        m_barriers[id]->reference_count = 0;
        m_barriers[id]->task_ids.clear();
        m_barriers[id]->waiting = false;
        return id;
    }
    else
    {
        DebugLog("Barrier Fail");
    }
    return INVALID_HYDRA_ID;
}

HydraID HydraTaskScheduler::CreateBarrier()
{
    return _CreateBarrier();
}
/// @brief Include the task identified by taskID in the barrier identified by barrierID
/// @param taskID 
/// @param barrierID 
void HydraTaskScheduler::_AddTaskToBarrier(HydraID taskID, HydraID barrierID)
{

    if ((barrierID < m_barriers.size()) && (taskID < m_tasks.size()))
    {
        if (m_barriers[barrierID]->waiting)
        {
            DebugLog("Barrier already locked");
        }
        else
        {
            m_barriers[barrierID]->task_ids.push_back(barrierID);
            m_barriers[barrierID]->reference_count++;
            HydraTask *task = _GetTask(taskID);
            task->m_barrierID = barrierID;
        }
    }
}

/// @brief Include the task identified by taskID in the barrier identified by barrierID
/// @param taskID 
/// @param barrierID 
void HydraTaskScheduler::AddTaskToBarrier(HydraID taskID, HydraID barrierID)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    _AddTaskToBarrier(taskID, barrierID);
}

/// @brief The calling thread will not return from this method until all of the tasks in the barrier are complete. 
/// @brief The thread will continue to process tasks from the task pool
/// @param barrierID 
void HydraTaskScheduler::WaitForBarrier(HydraID barrierID)
{
    if (barrierID < m_barriers.size())
    {           
        if(m_barriers[barrierID] != nullptr)
        {   
            bool quit = false;       
            while (!quit)
            {
                {
                    std::lock_guard<std::mutex> lock(m_mutex);    
                    if(m_barriers[barrierID] == nullptr)
                    {
                        return;
                    }
                    if(m_barriers[barrierID]->reference_count == 0)
                    {
                        _CompleteBarrier(barrierID);
                        return;
                    }
                }
                ProcessTask();   
            }

          
        }
    }
}

void HydraTaskScheduler::_CompleteBarrier(uint32_t barrierID)
{
//    std::lock_guard<std::mutex> lock(m_mutex);
    m_barriers[barrierID]->reference_count = 0;
    m_barriers[barrierID]->task_ids.clear();
    m_barriers[barrierID]->waiting = false;
    delete m_barriers[barrierID];
    m_barriers[barrierID] = nullptr;
    m_available_barrier_ids.push(barrierID);
}
/// @brief Decremenet the barrier reference. If a process is waiting for the barrier and the reference reaches 0, the process will continue.
/// @param taskID 
/// @param barrierID 
void HydraTaskScheduler::DecrementBarrier(HydraID taskID, HydraID barrierID)
{
    std::lock_guard<std::mutex> lock(m_mutex);
    _DecrementBarrier(taskID, barrierID);
}

/// @brief Decremenet the barrier reference. If a process is waiting for the barrier and the reference reaches 0, the process will continue.
/// @param taskID 
/// @param barrierID 
void HydraTaskScheduler::_DecrementBarrier(HydraID taskID, HydraID barrierID)
{
    
    if ((barrierID < m_barriers.size()) && (taskID < m_tasks.size()))
    {
        m_barriers[barrierID]->reference_count--;     
    }
}