#pragma once

#include <coroutine>
#include <iostream>
#include <stdexcept>
#include <utility>
#include <optional>
#include <vector>
#include <memory>

enum class task_result {
    failure,
    success
};

class node_task
{
public:
    class promise_type;

    using handle_type = std::coroutine_handle<promise_type>;

    class promise_type {
    public:
        std::optional<task_result> result;
        handle_type get_return_object() noexcept {
            return handle_type::from_promise(*this);
        }

        std::suspend_always initial_suspend() noexcept { return {}; }

        std::suspend_always final_suspend() noexcept { return {}; }

        void return_value(task_result r) noexcept {
            result = r;
        }

        void unhandled_exception() noexcept {
            std::terminate(); // TODO: handle exception
        }
    };

    class awaiter {
    public:
        awaiter(handle_type handle) noexcept : _handle(handle) {}

        bool await_ready() const noexcept {
            return false;
        }

        auto await_suspend(std::coroutine_handle<>) noexcept {
            return _handle;
        }

        task_result await_resume() noexcept {
            return *_handle.promise().result;
        }

    private:
        handle_type _handle;
    };

    node_task(handle_type handle) noexcept : _handle(handle) {}

    node_task(node_task&& other) noexcept : _handle(other._handle) {
        other._handle = nullptr;
    }

    ~node_task() {
        if (_handle) {
            _handle.destroy();
        }
    }

    node_task& operator=(node_task&& other) noexcept {
        if (&other != this) {
            if (_handle) {
                _handle.destroy();
            }
            _handle = other._handle;
            other._handle = nullptr;
        }
        return *this;
    }

    awaiter operator co_await() noexcept {
        return awaiter{_handle};
    }

    void resume() noexcept {
        _handle.resume();
    }

    bool done() const noexcept {
        return _handle.done();
    }

    task_result result() noexcept {
        return *_handle.promise().result;
    }

    std::optional<task_result> try_result() noexcept {
        return _handle.promise().result;
    }

private:
    handle_type _handle;
};

using yield = std::suspend_always;

template<typename Context>
class behavior_node {
public:
    virtual node_task tick(Context& ctx) noexcept = 0;

    virtual ~behavior_node() = default;
};

template<typename Context>
class sequence_node : public behavior_node<Context> {
public:
    void add_child(std::unique_ptr<behavior_node<Context>>&& child) {
        _children.push_back(std::move(child));
    }

    node_task tick(Context& ctx) noexcept override {
        std::size_t idx = 0;
        while (idx < _children.size()) {
            auto task = _children[idx]->tick(ctx);
            while (!task.done()) {
                task.resume();

                if (idx < _children.size() - 1)
                    co_await yield{};
            }

            if (task.result() == task_result::failure) {
                co_return task_result::failure;
            }

            ++idx;
        }

        co_return task_result::success;
    }
    
private:
    std::vector<std::unique_ptr<behavior_node<Context>>> _children;
};

template<typename Context>
auto sequence() {
    return std::make_unique<sequence_node<Context>>();
}

template<typename Context>
class selector_node : public behavior_node<Context> {
public:
    void add_child(std::unique_ptr<behavior_node<Context>>&& child) {
        _children.push_back(std::move(child));
    }

    node_task tick(Context& ctx) noexcept override {
        for (std::size_t idx = 0; idx < _children.size(); ++idx) {
            auto task = _children[idx]->tick(ctx);
            while (!task.done()) {
                task.resume();
                co_await yield{};
            }

            if (task.result() == task_result::success) {
                co_return task_result::success;
            }
        }
        co_return task_result::failure;
    }

private:
    std::vector<std::unique_ptr<behavior_node<Context>>> _children;
};

template<typename Context>
auto selector() {
    return std::make_unique<selector_node<Context>>();
}

template<typename Context, typename TickFn>
class action_node : public behavior_node<Context> {
public:
    action_node(TickFn tick) noexcept : _tick(tick) {}

    node_task tick(Context& ctx) noexcept override {
        return _tick(ctx);
    }

private:
    TickFn _tick;
};

template<typename Context, typename TickFn>
auto action(TickFn tick) {
    return std::make_unique<action_node<Context, TickFn>>(tick);
}

template<typename Context, typename ConditionFn>
class conditional_node : public behavior_node<Context> {
public:
    conditional_node(ConditionFn condition) noexcept : _condition(condition) {}

    node_task tick(Context& ctx) noexcept override {
        if (_condition(ctx)) {
            co_return task_result::success;
        }
        co_return task_result::failure;
    }

private:
    ConditionFn _condition;
};

template<typename Context, typename ConditionFn>
auto conditional(ConditionFn condition) {
    return std::make_unique<conditional_node<Context, ConditionFn>>(condition);
}