#include "Render_Graph_Runtime.hpp"

#include <condition_variable>
#include <exception>
#include <limits>
#include <mutex>
#include <stdexcept>
#include <unordered_map>
#include <utility>
#include <vector>

namespace renderive::render_graph::detail {

struct Render_Graph_Runtime::State
    : std::enable_shared_from_this<Render_Graph_Runtime::State> {
    struct Node_State {
        std::size_t remaining_predecessors{};
        Node_Execution_Status status{Node_Execution_Status::pending};
    };

    State(const Render_Plan& plan,
          std::span<Node_Execution* const> execution_slots,
          Execute_Node execute,
          Schedule scheduler,
          Current_Worker_Id worker_id)
        : nodes(plan.graph.nodes.size()),
          successors(plan.graph.nodes.size()),
          executions(execution_slots.begin(), execution_slots.end()),
          execute_node(std::move(execute)),
          schedule(std::move(scheduler)),
          current_worker_id(std::move(worker_id)),
          unfinished_nodes(plan.graph.nodes.size()) {
        if (!execute_node)
            throw std::invalid_argument("render graph node executor is empty");
        if (!schedule)
            throw std::invalid_argument("render graph scheduler is empty");
        if (!current_worker_id)
            throw std::invalid_argument("render graph worker id source is empty");
        if (!executions.empty() && executions.size() != nodes.size())
            throw std::invalid_argument(
                "render graph execution slot count differs from plan");
        if (executions.empty())
            executions.resize(nodes.size());

        std::vector<bool> occupied(nodes.size());
        std::unordered_map<Render_Node_Id, std::size_t> indices;
        indices.reserve(nodes.size());
        for (const auto& node : plan.graph.nodes) {
            if (node.execution_index >= nodes.size() ||
                occupied[node.execution_index])
                throw std::invalid_argument(
                    "render plan execution indices are not dense and unique");
            occupied[node.execution_index] = true;
            if (!indices.emplace(node.node_id, node.execution_index).second)
                throw std::invalid_argument(
                    "render plan contains duplicate node ids");
        }
        for (const auto& edge : plan.graph.edges) {
            const auto from = indices.find(edge.from);
            const auto to = indices.find(edge.to);
            if (from == indices.end() || to == indices.end())
                throw std::invalid_argument(
                    "render plan edge references an unknown node");
            successors[from->second].push_back(to->second);
            ++nodes[to->second].remaining_predecessors;
        }

        std::vector<std::size_t> remaining;
        remaining.reserve(nodes.size());
        std::vector<std::size_t> topological;
        topological.reserve(nodes.size());
        for (const auto& node : nodes)
            remaining.push_back(node.remaining_predecessors);
        for (std::size_t index = 0; index < remaining.size(); ++index) {
            if (remaining[index] == 0)
                topological.push_back(index);
        }
        for (std::size_t cursor = 0; cursor < topological.size(); ++cursor) {
            for (const std::size_t successor : successors[topological[cursor]]) {
                if (--remaining[successor] == 0)
                    topological.push_back(successor);
            }
        }
        if (topological.size() != nodes.size())
            throw std::invalid_argument("render plan contains a cycle");
    }

    void execute() {
        std::vector<std::size_t> ready;
        {
            std::lock_guard lock(mutex);
            if (started)
                throw std::logic_error("render graph runtime already executed");
            started = true;
            if (nodes.empty()) {
                terminal = true;
            } else {
                for (std::size_t index = 0; index < nodes.size(); ++index) {
                    if (nodes[index].remaining_predecessors == 0)
                        make_ready_locked(index, ready);
                }
                if (ready.empty()) {
                    fail_graph_locked(std::make_exception_ptr(
                        std::logic_error("render plan contains a cycle")));
                }
            }
        }
        submit(ready);

        std::exception_ptr error;
        {
            std::unique_lock lock(mutex);
            completion.wait(lock, [this] { return terminal; });
            error = first_exception;
        }
        if (error)
            std::rethrow_exception(error);
    }

    void make_ready_locked(std::size_t index,
                           std::vector<std::size_t>& ready) {
        auto& node = nodes.at(index);
        if (failed || node.status != Node_Execution_Status::pending)
            return;
        node.status = Node_Execution_Status::ready;
        if (auto* execution = executions[index]) {
            execution->ready_time_ns = render_clock_now_ns();
            execution->status = Node_Execution_Status::ready;
        }
        ready.push_back(index);
    }

    void submit(const std::vector<std::size_t>& ready) noexcept {
        for (const std::size_t index : ready)
            submit(index);
    }

    void submit(std::size_t index) noexcept {
        {
            std::lock_guard lock(mutex);
            if (failed || nodes[index].status != Node_Execution_Status::ready) {
                finish_if_terminal_locked();
                return;
            }
            ++active_tasks;
        }
        try {
            auto self = shared_from_this();
            schedule([self = std::move(self), index] {
                self->run_node(index);
            });
        } catch (...) {
            std::lock_guard lock(mutex);
            --active_tasks;
            fail_node_locked(index, std::current_exception(),
                             render_clock_now_ns());
            finish_if_terminal_locked();
        }
    }

    void run_node(std::size_t index) noexcept {
        Node_Execution_Metrics* metrics{};
        {
            std::lock_guard lock(mutex);
            if (failed || nodes[index].status != Node_Execution_Status::ready) {
                --active_tasks;
                finish_if_terminal_locked();
                return;
            }
            nodes[index].status = Node_Execution_Status::running;
            if (auto* execution = executions[index]) {
                execution->start_time_ns = render_clock_now_ns();
                try {
                    execution->worker_id = current_worker_id();
                } catch (...) {
                    execution->worker_id =
                        std::numeric_limits<std::uint32_t>::max();
                }
                execution->status = Node_Execution_Status::running;
                metrics = &execution->metrics;
            }
        }

        Node_Execution_Result result = Node_Execution_Result::completed();
        std::exception_ptr error;
        try {
            result = execute_node(index, metrics);
        } catch (...) {
            error = std::current_exception();
        }
        const std::uint64_t cpu_end = render_clock_now_ns();

        if (error) {
            std::lock_guard lock(mutex);
            --active_tasks;
            fail_node_locked(index, std::move(error), cpu_end);
            finish_if_terminal_locked();
            return;
        }
        if (!result.is_external()) {
            complete_synchronous(index, cpu_end);
            return;
        }

        {
            std::lock_guard lock(mutex);
            --active_tasks;
            auto& node = nodes[index];
            node.status = Node_Execution_Status::waiting_external;
            ++waiting_external;
            if (auto* execution = executions[index]) {
                execution->cpu_end_time_ns = cpu_end;
                execution->external_start_time_ns = cpu_end;
                execution->status = Node_Execution_Status::waiting_external;
            }
        }
        try {
            auto self = shared_from_this();
            result.operation().on_complete(
                [self = std::move(self), index](std::exception_ptr completion_error) {
                    self->complete_external(index, std::move(completion_error));
                });
        } catch (...) {
            complete_external(index, std::current_exception());
        }
    }

    void complete_synchronous(std::size_t index,
                              std::uint64_t cpu_end) noexcept {
        std::vector<std::size_t> ready;
        {
            std::lock_guard lock(mutex);
            --active_tasks;
            auto& node = nodes[index];
            node.status = Node_Execution_Status::complete;
            if (auto* execution = executions[index]) {
                execution->cpu_end_time_ns = cpu_end;
                execution->end_time_ns = cpu_end;
                execution->status = Node_Execution_Status::complete;
            }
            --unfinished_nodes;
            unlock_successors_locked(index, ready);
            finish_if_terminal_locked();
        }
        submit(ready);
    }

    void complete_external(std::size_t index,
                           std::exception_ptr error) noexcept {
        std::vector<std::size_t> ready;
        {
            std::lock_guard lock(mutex);
            auto& node = nodes[index];
            if (node.status != Node_Execution_Status::waiting_external)
                return;
            --waiting_external;
            const std::uint64_t end = render_clock_now_ns();
            if (auto* execution = executions[index]) {
                execution->external_end_time_ns = end;
                execution->end_time_ns = end;
            }
            if (error) {
                fail_node_locked(index, std::move(error), end);
            } else {
                node.status = Node_Execution_Status::complete;
                if (auto* execution = executions[index])
                    execution->status = Node_Execution_Status::complete;
                --unfinished_nodes;
                unlock_successors_locked(index, ready);
            }
            finish_if_terminal_locked();
        }
        submit(ready);
    }

    void unlock_successors_locked(std::size_t index,
                                  std::vector<std::size_t>& ready) {
        if (failed)
            return;
        for (const std::size_t successor : successors[index]) {
            auto& state = nodes[successor];
            if (state.remaining_predecessors == 0)
                continue;
            --state.remaining_predecessors;
            if (state.remaining_predecessors == 0)
                make_ready_locked(successor, ready);
        }
    }

    void fail_node_locked(std::size_t index, std::exception_ptr error,
                          std::uint64_t end) {
        auto& node = nodes[index];
        node.status = Node_Execution_Status::failed;
        if (auto* execution = executions[index]) {
            if (execution->start_time_ns == 0)
                execution->start_time_ns = end;
            if (execution->cpu_end_time_ns == 0)
                execution->cpu_end_time_ns = end;
            execution->end_time_ns = end;
            execution->status = Node_Execution_Status::failed;
        }
        fail_graph_locked(std::move(error));
    }

    void fail_graph_locked(std::exception_ptr error) {
        failed = true;
        if (!first_exception)
            first_exception = std::move(error);
    }

    void finish_if_terminal_locked() {
        if (terminal)
            return;
        if ((!failed && unfinished_nodes == 0) ||
            (failed && active_tasks == 0 && waiting_external == 0)) {
            terminal = true;
            completion.notify_all();
        }
    }

    std::vector<Node_State> nodes;
    std::vector<std::vector<std::size_t>> successors;
    std::vector<Node_Execution*> executions;
    Execute_Node execute_node;
    Schedule schedule;
    Current_Worker_Id current_worker_id;
    std::mutex mutex;
    std::condition_variable completion;
    std::exception_ptr first_exception;
    std::size_t unfinished_nodes{};
    std::size_t active_tasks{};
    std::size_t waiting_external{};
    bool started{};
    bool failed{};
    bool terminal{};
};

Render_Graph_Runtime::Render_Graph_Runtime(
    const Render_Plan& plan,
    std::span<Node_Execution* const> executions,
    Execute_Node execute_node,
    Schedule schedule,
    Current_Worker_Id current_worker_id)
    : state_(std::make_shared<State>(
          plan, executions, std::move(execute_node), std::move(schedule),
          std::move(current_worker_id))) {}

Render_Graph_Runtime::~Render_Graph_Runtime() = default;

void Render_Graph_Runtime::execute() {
    state_->execute();
}

} // namespace renderive::render_graph::detail