#pragma once

#include "address.h"
#include "function.h"
#include "http.h"
#include "storage.h"
#include "unistat.h"

#include <contrib/libs/yaml-cpp/include/yaml-cpp/yaml.h>

#include <util/datetime/base.h>
#include <util/generic/hash.h>
#include <util/generic/scope.h>
#include <util/string/builder.h>
#include <util/string/cast.h>

namespace YAML {
    template <typename T>
    struct convert {
        static bool decode(const Node& n, T& r) {
            return n.IsScalar() && TryFromString(n.Scalar(), r);
        }
    };

    template <typename T, typename Period>
    struct convert<std::chrono::duration<T, Period>> {
        static bool decode(const Node& n, std::chrono::duration<T, Period>& r) {
            if (n.Scalar() == "inf") {
                r = std::chrono::duration<T, Period>::max();
                return true;
            }
            TDuration d;
            if (!convert<TDuration>::decode(n, d)) {
                return false;
            }
            // Despite the parser accepting "ns" as a suffix, TDuration can only represent microseconds.
            r = std::chrono::duration_cast<std::chrono::duration<T, Period>>(std::chrono::microseconds(d.MicroSeconds()));
            return true;
        }
    };
}

// If an expression is false, throw an exception that points to the location of the given
// node as the source of the error.
#define CHECK_NODE(node, expr, ...) do if (!(expr)) \
    throw YAML::Exception((node).Mark(), (TStringBuilder() << __VA_ARGS__)); while (0)

#define FAIL_NODE(node, ...) CHECK_NODE(node, false, __VA_ARGS__)

namespace NSv {
    // Try to convert the node's contents into the needed type, then evaluate a check;
    // if either step fails, throw an error that points to the node as the reason.
    template <typename T, typename F = bool(*)(const T&)>
    static T Required(const YAML::Node& arg, F&& check = [](const T&) { return true; }) {
        T val;
        CHECK_NODE(arg, YAML::convert<T>::decode(arg, val) && check(val), "invalid value");
        return val;
    }

    // Same as `Required`, but return a default value if the node does not exist at all.
    template <typename T, typename F = bool(*)(const T&)>
    static T Optional(const YAML::Node& arg, T fallback, F&& check = [](const T&) { return true; }) {
        return arg ? Required<T>(arg, std::forward<F>(check)) : fallback;
    }

    // Compare two YAML nodes for value equality, ignoring tags. (`YAML::Node::operator==`
    // checks for reference equality.)
    bool EqYAML(const YAML::Node&, const YAML::Node&);

    // Parse a YAML config, but with extensions.
    //
    //  * Within any line, {{X}} is replaced with the value of the env var X.
    //
    //  * Lines starting with `#include: ` are replaced with the contents of the
    //    specified file. (Anchors can be used across file boundaries, too.)
    //
    //  * `#default: X Y` defines a default value for env var `X`. Defaults are scoped
    //    to a single file and everything it `#include`s.
    //
    YAML::Node LoadYAML(const TString& path);

    // Remap a line number in a YAML node produced by the above call to its actual position,
    // which might be in an `#include`d file.
    std::pair<TString, int> MapYAMLLine(const TString& path, int lineno);

    // Apply a patch to a YAML node. **Modifies both the node and the patch in-place.**
    //
    // Lists:
    //    - !add x  # append element
    //    - !del x  # remove first element with same value
    //
    // Lists of maps:
    //    - !add x: y  # append element
    //    - !del x     # remove first element with first key = `x`
    //    - !set x: y  # replace element
    //    - !mod x: y  # apply patch `y` to an element
    //    - !set-arg x: y  # replace the value of key `x` in an element in which it is the first key
    //    - !mod-arg x: y  # apply patch `y` to the value of key `x`
    //
    // Maps:
    //    - !add x: y  # add a new key, fail if exists
    //    - !del x:    # remove a key
    //    - !set x: y  # replace a value
    //    - !mod x: y  # apply patch `y` to a value
    //
    // If the patch is a sequence, each element is applied in order.
    //
    YAML::Node PatchYAML(YAML::Node node, YAML::Node patch);

    // An action is a small module that does something with the request, such as:
    //  * modifying its headers by using the provided pointer;
    //  * wrapping the provided stream into a custom `IStream` implementation to hijack
    //    calls to `WriteHead` and add headers to the response;
    //  * calling `WriteHead`, `Write`, and `Close` itself to respond to the request.
    using TAction = NSv::TFunction<bool(NSv::IStreamPtr&)>;

    // TODO name this properly
    struct TAuxData {
    private:
        // Toggle collection and set the prefix of all signals allocated between the call
        // to this function and the return value's destruction.
        auto CollectSignals(bool enable, TString prefix) noexcept {
            auto restore = [c = CollectingSignals_, p = SignalPrefix_](TAuxData* aux) noexcept {
                aux->SignalPrefix_ = p;
                aux->CollectingSignals_ = c;
            };
            SignalPrefix_ = prefix;
            CollectingSignals_ = enable;
            return std::unique_ptr<TAuxData, decltype(restore)>{this, std::move(restore)};
        }

        template <typename F>
        TAction AttachCounters(F&& a) {
            if (!CollectingSignals_) {
                return std::forward<F>(a);
            }
            auto& requests    = Signal("requests_dmmm");
            auto& requestsNow = Signal("requests-active_ammv");
            auto& failures    = Signal("failures_dmmm");
            auto& cancelled   = Signal("cancelled_dmmm");
            auto& time        = Signal<NSv::THistogram>("time_dhhh");
            return [&, a = std::forward<F>(a)](NSv::IStreamPtr& req) mutable {
                requests++;
                requestsNow++;
                Y_DEFER {
                    requestsNow--;
                };
                auto timer = NSv::Timer(time);
                if (a(req)) {
                    return true;
                }
                if (mun_errno == ECANCELED) {
                    cancelled++;
                } else if (mun_errno != EREQDONE) {
                    failures++;
                }
                return false;
            };
        }

        TAction ActionList(const YAML::Node& items) {
            if (items.size() == 0) {
                return AttachCounters([](NSv::IStreamPtr&) {
                    return true;
                });
            }
            if (items.size() == 1) {
                return AttachCounters(Action(items[0]));
            }
            TVector<TAction> parsed;
            for (const auto& item : items) {
                parsed.emplace_back(Action(item));
            }
            return AttachCounters([parsed = std::move(parsed)](NSv::IStreamPtr& req) mutable {
                for (const auto& action : parsed) {
                    if (!action(req)) {
                        return false;
                    }
                }
                return true;
            });
        }

    public:
        struct TActionMap : THashMap<TString, NSv::TFunction<TAction(const YAML::Node&, TAuxData&)>> {};

        // Construct a signal with a given type in a place where it will be accessible
        // by the unistat admin handle. Depending on the current collection settings,
        // the signal might have a prefix added to the name, or placed into the "ignore"
        // list. The return value must not be moved.
        template <typename T = TNumber<ui64>, typename... Args>
        T& Signal(TStringBuf name, Args&&... args) {
            if (!CollectingSignals_) {
                // Inactive signals are fully independent.
                auto p = std::make_shared<T>(TString{}, std::forward<Args>(args)...);
                Inactive_.emplace_back(p);
                return *p;
            }
            // Active signals with the same name map to the same object.
            auto fullName = TString::Join(SignalPrefix_, "-", name);
            auto p = NSv::StaticData(fullName, [&]{
                return T{fullName, std::forward<Args>(args)...};
            });
            if (ActiveSet_.emplace(p.get()).second) {
                Active_.push_back(p);
            }
            return *p;
        }

        // Construct a signal, but don't take in acccount signal prefix and tags
        // Custom signals are always active
        template <typename T = TNumber<ui64>, typename... Args>
        T& CustomSignal(TStringBuf name, Args&&... args) {
            auto fullName = TString(name);
            auto p = NSv::StaticData(fullName, [&]{
                return T{fullName, std::forward<Args>(args)...};
            });
            if (ActiveSet_.emplace(p.get()).second) {
                Active_.push_back(p);
            }
            return *p;
        }

        // Whether the signal created by the above method will appear in the unistat handle.
        bool CollectingSignals() const noexcept {
            return CollectingSignals_;
        }

        // Return all signals allocated while signal collection was enabled.
        const TVector<TSignal>& Signals() const noexcept {
            return Active_;
        }

        bool AddAction(const TString& name, const YAML::Node& body) {
            // FIXME should accept single actions too (that's, uh, in the README...)
            CHECK_NODE(body, body.IsSequence(), "action definition must be a sequence of actions");
            auto res = Actions_.emplace(name, [=, ref = TAction()](auto& call, auto& aux) mutable {
                if (!ref) {
                    auto prev = std::find(Constructing_.begin(), Constructing_.end(), name);
                    if (prev != Constructing_.end()) {
                        TStringBuilder msg;
                        while (prev != Constructing_.end()) {
                            msg << *prev++ << " -> ";
                        }
                        FAIL_NODE(call, "action cycle: " << msg << name);
                    }
                    Constructing_.emplace_back(name);
                    Y_DEFER {
                        Constructing_.pop_back();
                    };
                    // Signals named <name>-<requests|errors|...> count totals:
                    ref = aux.CollectSignals(true, name)->ActionList(body);
                }
                // <another section>-<stat id>-<requests|errors|...> count uses of one reference:
                return aux.AttachCounters(ref);
            });
            if (!res.second) {
                return false;
            }
            ConfigDefined_.insert(res.first->first);
            return true;
        }

        // Construct an action given its config description.
        TAction Action(const YAML::Node& node) {
            CHECK_NODE(node, !node.IsMap() || (node.size() && node.begin()->first.IsScalar()),
                       "an action reference must be `name` or `{name: primary argument, ...}`");
            auto name = node.IsMap() ? node.begin()->first : node;
            auto statId = TStringBuf(name.Tag()).Skip(1);
            auto prefix = TString::Join(SignalPrefix_, statId ? "-" : "", statId);
            // Cannot create new signals after freezing the action set.
            auto collect = CollectSignals(TStringBuf(name.Tag()).StartsWith("!") && !Frozen_, prefix);
            return node.IsSequence() ? ActionList(node) : Action(name.Scalar(), node);
        }

        TAction Action(TStringBuf name, const YAML::Node& node = {}) {
            // After freezing the action set, no new objects can be created. However,
            // config-defined actions have been pre-constructed, and constructing them
            // again simply yields the same object, which is safe. That could be used
            // e.g. for sending a request to an action specified as a parameter.
            if (Frozen_ && !ConfigDefined_.contains(name)) {
                return {};
            }
            auto it = Actions_.find(name);
            CHECK_NODE(node, it != Actions_.end(), "action " << name << " is not defined");
            return it->second(node, *this);
        }

        void Freeze() {
            for (auto name : ConfigDefined_) {
                Action(name, YAML::Node{});
            }
            Frozen_ = true;
        }

    private:
        TActionMap Actions_{*Singleton<TActionMap>()};
        THashSet<TStringBuf> ConfigDefined_;
        TVector<TStringBuf> Constructing_;
        TVector<TSignal> Inactive_;
        TVector<TSignal> Active_;
        THashSet<const void*> ActiveSet_;
        TString SignalPrefix_;
        bool CollectingSignals_ = false;
        bool Frozen_ = false;
    };
}

#define SV_DEFINE_ACTION(name, fn)                                                 \
    static auto Y_GENERATE_UNIQUE_ID(Action) = []{                                 \
        Y_ENSURE(Singleton<NSv::TAuxData::TActionMap>()->emplace(name, fn).second, \
                 "action " << name << " defined in multiple source files");        \
        return 0;                                                                  \
    }()
