#include "Scheduler.hpp"

#include "MayaFlux/Journal/Archivist.hpp"

namespace MayaFlux::Vruta {

TaskScheduler::TaskScheduler(uint32_t default_sample_rate, uint32_t default_frame_rate)
    : m_clock(default_sample_rate)
    , m_cleanup_threshold(512)
    , m_registered_sample_rate(default_sample_rate)
    , m_registered_frame_rate(default_frame_rate)
{
    ensure_domain(ProcessingToken::SAMPLE_ACCURATE, default_sample_rate);
    ensure_domain(ProcessingToken::FRAME_ACCURATE, default_frame_rate);
    ensure_domain(ProcessingToken::MULTI_RATE, default_sample_rate);
    ensure_domain(ProcessingToken::ON_DEMAND, 1);
}

void TaskScheduler::add_task(const std::shared_ptr<Routine>& routine, const std::string& name, bool initialize)
{
    if (!routine) {
        MF_ERROR(Journal::Component::Vruta, Journal::Context::CoroutineScheduling, "Failed to initiate routine; routine is null. Exiting add_task.");
        return;
    }

    std::string task_name = name.empty() ? auto_generate_name(routine) : name;
    ProcessingToken token = routine->get_processing_token();

    ensure_domain(token);

    if (initialize)
        initialize_routine_state(routine, token);

    for (auto& op : m_pending_ops) {
        bool expected = false;
        if (op.active.compare_exchange_strong(expected, true,
                std::memory_order_acquire, std::memory_order_relaxed)) {
            op.entry = { routine, task_name };
            op.is_addition = true;
            m_pending_count.fetch_add(1, std::memory_order_relaxed);
            return;
        }
    }

    MF_ERROR(Journal::Component::Vruta, Journal::Context::CoroutineScheduling,
        "Pending task queue full, could not add task '{}'", task_name);
}

bool TaskScheduler::cancel_task(const std::string& name)
{
    for (auto& op : m_pending_ops) {
        bool expected = false;
        if (op.active.compare_exchange_strong(expected, true,
                std::memory_order_acquire, std::memory_order_relaxed)) {
            op.entry = { nullptr, name };
            op.is_addition = false;
            m_pending_count.fetch_add(1, std::memory_order_relaxed);
            return true;
        }
    }

    MF_ERROR(Journal::Component::Vruta, Journal::Context::CoroutineScheduling,
        "Pending task queue full, could not cancel task '{}'", name);

    return false;
}

bool TaskScheduler::cancel_task(const std::shared_ptr<Routine>& routine)
{
    if (!routine)
        return false;

    auto it = find_task_by_routine(routine);
    const std::string name = (it != m_tasks.end()) ? it->name : "";

    for (auto& op : m_pending_ops) {
        bool expected = false;
        if (op.active.compare_exchange_strong(expected, true,
                std::memory_order_acquire, std::memory_order_relaxed)) {
            op.entry = { routine, name };
            op.is_addition = false;
            m_pending_count.fetch_add(1, std::memory_order_relaxed);
            return true;
        }
    }

    MF_ERROR(Journal::Component::Vruta, Journal::Context::CoroutineScheduling,
        "Pending task queue full, could not cancel task '{}'", name);
    return false;
}

bool TaskScheduler::restart_task(const std::string& name)
{
    auto it = find_task_by_name(name);
    if (it != m_tasks.end()) {
        if (it->routine && it->routine->is_active()) {
            it->routine->restart();
        }
    }
    return false;
}

std::shared_ptr<Routine> TaskScheduler::get_task(const std::string& name) const
{
    auto it = find_task_by_name(name);
    return (it != m_tasks.end()) ? it->routine : nullptr;
}

std::vector<std::shared_ptr<Routine>> TaskScheduler::get_tasks_for_token(ProcessingToken token) const
{
    std::vector<std::shared_ptr<Routine>> result;
    for (const auto& entry : m_tasks) {
        if (entry.routine && entry.routine->get_processing_token() == token) {
            result.push_back(entry.routine);
        }
    }
    return result;
}

std::vector<std::string> TaskScheduler::get_task_names() const
{
    std::vector<std::string> names;
    for (const auto& entry : m_tasks) {
        if (entry.routine)
            names.push_back(entry.name);
    }
    return names;
}

std::vector<std::string> TaskScheduler::get_task_names(ProcessingToken token) const
{
    std::vector<std::string> names;
    for (const auto& entry : m_tasks) {
        if (entry.routine && entry.routine->get_processing_token() == token)
            names.push_back(entry.name);
    }
    return names;
}

const std::vector<TaskEntry>& TaskScheduler::get_all_tasks()
{
    drain_pending_tasks();
    return m_tasks;
}

void TaskScheduler::process_token(ProcessingToken token, uint64_t processing_units)
{
    drain_pending_tasks();

    auto processor_it = m_token_processors.find(token);
    if (processor_it != m_token_processors.end()) {
        auto tasks = get_tasks_for_token(token);
        processor_it->second(tasks, processing_units);
    } else {
        process_default(token, processing_units);
    }

    static uint64_t cleanup_counter = 0;
    if (++cleanup_counter % (static_cast<uint64_t>(m_cleanup_threshold * 2)) == 0) {
        cleanup_completed_tasks();
    }
}

void TaskScheduler::process_all_tokens()
{
    for (const auto& token : m_token_clocks) {
        process_token(token.first, 0);
    }

    static uint64_t cleanup_counter = 0;
    if (++cleanup_counter % m_cleanup_threshold == 0) {
        cleanup_completed_tasks();
    }
}

void TaskScheduler::register_token_processor(ProcessingToken token, token_processing_func_t processor)
{
    ensure_domain(token);
    m_token_processors[token] = std::move(processor);
}

const IClock& TaskScheduler::get_clock(ProcessingToken token) const
{
    auto clock_it = m_token_clocks.find(token);
    if (clock_it != m_token_clocks.end()) {
        return *clock_it->second;
    }

    auto audio_clock_it = m_token_clocks.find(ProcessingToken::SAMPLE_ACCURATE);
    if (audio_clock_it != m_token_clocks.end()) {
        return *audio_clock_it->second;
    }

    error<std::runtime_error>(Journal::Component::Vruta, Journal::Context::CoroutineScheduling, std::source_location::current(),
        "No clock found for token {}, and no sample-accurate clock available as fallback", static_cast<int>(token));
}

uint64_t TaskScheduler::seconds_to_units(double seconds, ProcessingToken token) const
{
    unsigned int rate = get_rate(token);
    return static_cast<uint64_t>(seconds * rate);
}

uint64_t TaskScheduler::seconds_to_samples(double seconds) const
{
    return static_cast<uint64_t>(seconds * get_rate(ProcessingToken::SAMPLE_ACCURATE));
}

uint64_t TaskScheduler::current_units(ProcessingToken token) const
{
    const auto& clock = get_clock(token);
    return clock.current_position();
}

unsigned int TaskScheduler::get_rate(ProcessingToken token) const
{
    auto clock_it = m_token_clocks.find(token);
    if (clock_it != m_token_clocks.end()) {
        return clock_it->second->rate();
    }

    return get_default_rate(token);
}

bool TaskScheduler::has_active_tasks(ProcessingToken token) const
{
    return std::ranges::any_of(m_tasks,
        [token](const TaskEntry& entry) {
            return entry.routine && entry.routine->is_active() && entry.routine->get_processing_token() == token;
        });
}

uint64_t TaskScheduler::get_next_task_id() const
{
    return m_next_task_id.fetch_add(1);
}

std::string TaskScheduler::auto_generate_name(const std::shared_ptr<Routine>& /*routine*/) const
{
    return "task_" + std::to_string(get_next_task_id());
}

std::vector<TaskEntry>::iterator TaskScheduler::find_task_by_name(const std::string& name)
{
    return std::ranges::find_if(m_tasks,
        [&name](const TaskEntry& entry) { return entry.name == name; });
}

std::vector<TaskEntry>::const_iterator TaskScheduler::find_task_by_name(const std::string& name) const
{
    return std::ranges::find_if(m_tasks,
        [&name](const TaskEntry& entry) { return entry.name == name; });
}

std::vector<TaskEntry>::iterator TaskScheduler::find_task_by_routine(const std::shared_ptr<Routine>& routine)
{
    return std::ranges::find_if(m_tasks,
        [&routine](const TaskEntry& entry) { return entry.routine == routine; });
}

unsigned int TaskScheduler::get_default_rate(ProcessingToken token) const
{
    switch (token) {
    case ProcessingToken::SAMPLE_ACCURATE:
        return m_registered_sample_rate;
    case ProcessingToken::FRAME_ACCURATE:
        return m_registered_frame_rate;
    case ProcessingToken::MULTI_RATE:
        return m_registered_sample_rate;
    case ProcessingToken::ON_DEMAND:
        return 1;
    case ProcessingToken::CUSTOM:
        return 1000;
    default:
        return m_registered_sample_rate;
    }
}

void TaskScheduler::ensure_domain(ProcessingToken token, unsigned int rate)
{
    auto clock_it = m_token_clocks.find(token);
    if (clock_it == m_token_clocks.end()) {
        unsigned int domain_rate = (rate > 0) ? rate : get_default_rate(token);

        switch (token) {
        case ProcessingToken::FRAME_ACCURATE:
            m_token_clocks[token] = std::make_unique<FrameClock>(domain_rate);
            break;
        case ProcessingToken::SAMPLE_ACCURATE:
        case ProcessingToken::ON_DEMAND:
        default:
            m_token_clocks[token] = std::make_unique<SampleClock>(domain_rate);
            break;
        }
    }
}

void TaskScheduler::process_default(ProcessingToken token, uint64_t processing_units)
{
    auto clock_it = m_token_clocks.find(token);
    if (clock_it == m_token_clocks.end()) {
        return;
    }

    auto tasks = get_tasks_for_token(token);
    if (tasks.empty()) {
        auto& clock = *clock_it->second;
        clock.tick(processing_units);
        return;
    }

    auto& clock = *clock_it->second;

    for (uint64_t i = 0; i < processing_units; i++) {
        uint64_t current_context = clock.current_position();

        for (auto& routine : tasks) {
            if (routine && routine->is_active()) {
                if (routine->requires_clock_sync()) {
                    if (current_context >= routine->next_execution()) {
                        routine->try_resume_with_context(current_context, DelayContext::SAMPLE_BASED);
                    }
                } else {
                    routine->try_resume_with_context(current_context, DelayContext::SAMPLE_BASED);
                }
            }
        }

        clock.tick(1);
    }
}

void TaskScheduler::cleanup_completed_tasks()
{
    std::erase_if(m_tasks, [](const TaskEntry& entry) {
        return !entry.routine || !entry.routine->is_active();
    });
}

bool TaskScheduler::initialize_routine_state(const std::shared_ptr<Routine>& routine, ProcessingToken token)
{
    if (!routine) {
        return false;
    }

    auto clock_it = m_token_clocks.find(token);
    if (clock_it == m_token_clocks.end()) {
        return false;
    }

    uint64_t current_context = clock_it->second->current_position();
    return routine->initialize_state(current_context);
}

void TaskScheduler::pause_all_tasks()
{
    for (auto& entry : m_tasks) {
        if (entry.routine && entry.routine->is_active()) {
            bool current_auto_resume = entry.routine->get_auto_resume();
            entry.routine->set_state<bool>("was_auto_resume", current_auto_resume);
            entry.routine->set_auto_resume(false);
        }
    }
}

void TaskScheduler::resume_all_tasks()
{
    for (auto& entry : m_tasks) {
        if (entry.routine && entry.routine->is_active()) {
            auto was_auto_resume = entry.routine->get_state<bool>("was_auto_resume");
            if (was_auto_resume) {
                entry.routine->set_auto_resume(*was_auto_resume);
            } else {
                entry.routine->set_auto_resume(true);
            }
        }
    }
}

void TaskScheduler::terminate_all_tasks()
{
    drain_pending_tasks();

    for (auto& entry : m_tasks) {
        if (entry.routine && entry.routine->is_active()) {
            entry.routine->set_should_terminate(true);
        }
    }

    for (auto& entry : m_tasks) {
        if (entry.routine && entry.routine->is_active()) {
            entry.routine->force_resume();
        }
    }

    constexpr int MAX_ATTEMPTS = 3;
    for (int attempt = 0; attempt < MAX_ATTEMPTS; ++attempt) {
        std::this_thread::sleep_for(std::chrono::milliseconds(5));

        bool any_active = false;
        for (auto& entry : m_tasks) {
            if (entry.routine && entry.routine->is_active()) {
                any_active = true;
                entry.routine->force_resume();
            }
        }

        if (!any_active) {
            break;
        }
    }

    bool all_done = true;
    for (const auto& entry : m_tasks) {
        if (entry.routine && entry.routine->is_active()) {
            all_done = false;
            MF_WARN(Journal::Component::Vruta, Journal::Context::CoroutineScheduling,
                "Coroutine '{}' stuck after {} attempts - forcing destruction",
                entry.name, MAX_ATTEMPTS);
        }
    }

    if (!all_done) {
        MF_WARN(Journal::Component::Vruta, Journal::Context::CoroutineScheduling,
            "Some coroutines did not complete - forcing destruction");
    } else {
        MF_LOG(Journal::Component::Vruta, Journal::Context::CoroutineScheduling,
            "All coroutines terminated successfully");
    }

    m_tasks.clear();
}

void TaskScheduler::process_buffer_cycle_tasks()
{
    drain_pending_tasks();

    m_current_buffer_cycle++;
    auto tasks = get_tasks_for_token(ProcessingToken::SAMPLE_ACCURATE);

    for (auto& task : tasks) {
        if (task && task->is_active()) {
            if (task->requires_clock_sync()) {
                if (m_current_buffer_cycle >= task->next_execution()) {
                    task->try_resume_with_context(m_current_buffer_cycle, DelayContext::BUFFER_BASED);
                }
            } else {
                task->try_resume_with_context(m_current_buffer_cycle, DelayContext::BUFFER_BASED);
            }
        }
    }
}

void TaskScheduler::drain_pending_tasks()
{
    if (m_pending_count.load(std::memory_order_relaxed) == 0)
        return;

    for (auto& op : m_pending_ops) {
        if (!op.active.load(std::memory_order_acquire))
            continue;

        if (op.is_addition) {
            auto existing_it = find_task_by_name(op.entry.name);
            if (existing_it != m_tasks.end()) {
                if (existing_it->routine && existing_it->routine->is_active())
                    existing_it->routine->set_should_terminate(true);
                m_tasks.erase(existing_it);
            }
            m_tasks.push_back(std::move(op.entry));
        } else {
            auto it = find_task_by_name(op.entry.name);
            if (it != m_tasks.end()) {
                if (it->routine && it->routine->is_active())
                    it->routine->set_should_terminate(true);
                m_tasks.erase(it);
            }
        }

        op.entry = { nullptr, "" };
        op.active.store(false, std::memory_order_release);
        m_pending_count.fetch_sub(1, std::memory_order_relaxed);
    }
}

}
