#include "config.h"

#include <util/folder/path.h>
#include <util/stream/file.h>
#include <util/string/split.h>
#include <util/string/strip.h>
#include <util/system/env.h>

#include <sstream>

template <>
void Out<YAML::Node>(IOutputStream& out, TTypeTraits<YAML::Node>::TFuncParam node) {
    out << (std::ostringstream{} << node).str();
}

bool NSv::EqYAML(const YAML::Node& a, const YAML::Node& b) {
    if (a.IsScalar() && b.IsScalar()) {
        return a.Scalar() == b.Scalar();
    }
    bool seq = a.IsSequence() && b.IsSequence();
    bool map = a.IsMap() && b.IsMap();
    if (seq || map) {
        if (a.size() != b.size()) {
            return false;
        }
        for (auto i1 = a.begin(), i2 = b.begin(); i1 != a.end(); i1++, i2++) {
            if (seq ? !EqYAML(*i1, *i2) : !EqYAML(i1->first, i2->first) || !EqYAML(i1->second, i2->second)) {
                return false;
            }
        }
        return true;
    }
    return false;
}

static TFsPath RelativeToDir(const TFsPath& a, const TFsPath& b) {
    return b.IsRelative() ? a.Parent() / b : b;
}

static TMaybe<std::pair<TFsPath, size_t>> IsIncludeLine(const TFsPath& prev, TStringBuf line) {
    size_t indent = 0;
    while (line.SkipPrefix(" ")) {
        indent++;
    }
    if (!line.SkipPrefix("#include: ")) {
        return {};
    }
    return std::make_pair(RelativeToDir(prev, line), indent);
}

static TMaybe<std::pair<TString, TString>> IsDefaultLine(TStringBuf line) {
    while (line.SkipPrefix(" ")) {}
    TStringBuf k, v;
    if (!line.SkipPrefix("#default: ") || !line.TrySplit(' ', k, v)) {
        return {};
    }
    return std::make_pair(TString(k), TString(v));
}

static void EnvReplace(TString& line, const THashMap<TString, TString>& env) {
    for (size_t i = line.size() + 1; (i = line.rfind("{{", i - 1)) != TString::npos; ) {
        size_t j = line.find("}}", i);
        if (j == TString::npos) {
            break;
        }
        auto name = StripString(line.substr(i + 2, j - i - 2));
        auto it = env.find(name);
        line.replace(i, j - i + 2, GetEnv(name, it != env.end() ? it->second : ""));
    }
}

static std::string LoadYAML(THashMap<TString, TString> env, const TFsPath& path, size_t indent) {
    std::string result;
    TFileInput in(path);
    for (TString line; in.ReadLine(line);) {
        EnvReplace(line, env);
        if (auto def = IsDefaultLine(line)) {
            env[def->first] = def->second;
        } else if (auto include = IsIncludeLine(path, line)) {
            result += LoadYAML(env, include->first, indent + include->second) + "\n";
        } else {
            result += TString(indent, ' ') + line + "\n";
        }
    }
    return result;
}

YAML::Node NSv::LoadYAML(const TString& path) {
    return YAML::Load(LoadYAML({}, path, 0));
}

static std::pair<TString, int> MapYAMLLine(THashMap<TString, TString> env, const TFsPath& path, int& lineno) {
    TFileInput in(path);
    TString line;
    for (int i = 0; in.ReadLine(line); i++) {
        EnvReplace(line, env);
        if (auto def = IsDefaultLine(line)) {
            env[def->first] = def->second;
            continue;
        }
        if (auto include = IsIncludeLine(path, line)) {
            auto ret = MapYAMLLine(env, include->first, lineno);
            if (ret.second >= 0) {
                return ret;
            }
        }
        if (!lineno--) {
            return {TString(path), i};
        }
    }
    return {path, -1};
}

std::pair<TString, int> NSv::MapYAMLLine(const TString& path, int lineno) {
    return MapYAMLLine({}, path, lineno);
}

static auto FindByKey(YAML::Node node, YAML::Node key) {
    for (auto it = node.begin(); it != node.end(); it++) {
        if (node.IsMap() ? NSv::EqYAML(it->first, key) : it->IsMap() && it->size() && NSv::EqYAML(it->begin()->first, key)) {
            return it;
        }
    }
    return node.end();
}

YAML::Node NSv::PatchYAML(YAML::Node n, YAML::Node p) {
    if (p.IsSequence()) {
        // P(X, []) = X
        // P(X, [p, .1]) = P(P(X, p), [.1])
        return std::accumulate(p.begin(), p.end(), n, PatchYAML);
    } else if (n.IsMap()) {
        if (p.Tag() == "!del") {
            // P({k: v, .1}, !del k) = {.1}
            auto nk = FindByKey(n, p);
            CHECK_NODE(p, nk != n.end(), "this key is not set");
            n.remove(nk->first);
        } else if (p.IsMap() && p.size() == 1) {
            auto k = p.begin()->first;
            auto v = p.begin()->second;
            auto nk = FindByKey(n, k);
            if (k.Tag() == "!add") {
                // P({k: v, .1}, {!add K: V}) = {k: v, K: V, .1}
                CHECK_NODE(k, nk == n.end(), "this key is already set; use !set");
                k.SetTag("");
                n[k] = v;
            } else if (k.Tag() == "!set") {
                // P({k: v, .1}, {!set k: V}) = {k: V, .1}
                CHECK_NODE(k, nk != n.end(), "this key is not set; use !add");
                nk->second = v;
            } else if (k.Tag() == "!mod") {
                // P({k: v, .1}, {!mod k: p}) = {k: P(v, p), .1}
                CHECK_NODE(k, nk != n.end(), "this key is not set");
                nk->second = PatchYAML(nk->second, v);
            } else {
                FAIL_NODE(p, "invalid patch for a map");
            }
        } else {
            FAIL_NODE(p, "invalid patch for a map");
        }
    } else if (n.IsSequence()) {
        if (p.Tag() == "!add") {
            // P([.1], !add x) = [.1, x]
            p.SetTag("");
            n.push_back(p);
        } else if (p.Tag() == "!del") {
            // P([.1, x, .2], !del x) = [.1, .2]
            // P([.1, {k: v, .2}, .3], !del k) = [.1, .3]
            FAIL_NODE(p, "!del on list items requires updating yaml-cpp");
        } else if (p.IsMap() && p.size() > 0) {
            auto k = p.begin()->first;
            auto v = p.begin()->second;
            auto nk = FindByKey(n, k);
            if (k.Tag() == "!add") {
                // P([.1], {!add K: V, .2}) = [.1, {K: V, .2}]
                k.SetTag("");
                n.push_back(p);
            } else if (k.Tag() == "!set") {
                // P([.1, {k: v, .2}, .3], {!set k: V, .4}) = [.1, {k: V, .4}, .3]
                CHECK_NODE(k, nk != n.end(), "this key is not in list; use !add");
                k.SetTag("");
                static_cast<YAML::Node&&>(*nk) = p;
            } else if (k.Tag() == "!set-arg") {
                // P([.1, {k: v, .2}, .3], {!set-arg k: V}) = [.1, {k: V, .2}, .3]
                CHECK_NODE(k, nk != n.end(), "this key is not in list");
                nk->begin()->second = v;
            } else if (k.Tag() == "!mod") {
                // P([.1, {k: v, .2}, .3], {!mod k: p}) = [.1, P({k: v, .2}, p), .3]
                CHECK_NODE(k, nk != n.end(), "this key is not in list");
                static_cast<YAML::Node&&>(*nk) = PatchYAML(*nk, v);
            } else if (k.Tag() == "!mod-arg") {
                // P([.1, {k: v, .2}, .3], {!mod-arg k: p}) = [.1, {k: P(v, p), .2}, .3]
                CHECK_NODE(k, nk != n.end(), "this key is not in list");
                nk->begin()->second = PatchYAML(nk->begin()->second, v);
            } else {
                FAIL_NODE(p, "invalid patch for a sequence");
            }
        } else {
            FAIL_NODE(p, "invalid patch for a sequence");
        }
    } else {
        FAIL_NODE(p, "cannot apply a patch to a scalar value");
    }
    return n;
}
