#include <lix/config.h> // IWYU pragma: keep

// doesn't exist on macOS
// IWYU pragma: no_include <bits/types/struct_rusage.h>

#include <lix/libexpr/attr-path.hh>
#include <lix/libstore/local-fs-store.hh>
#include <lix/libcmd/installable-flake.hh>
#include <sys/resource.h>
#include <stdio.h>
#include <stdlib.h>
#include <lix/libexpr/attr-set.hh>
#include <lix/libutil/canon-path.hh>
#include <lix/libcmd/common-eval-args.hh>
#include <lix/libutil/current-process.hh>
#include <lix/libutil/error.hh>
#include <lix/libexpr/eval-cache.hh>
#include <lix/libexpr/eval-inline.hh>
#include <lix/libexpr/eval.hh>
#include <lix/libexpr/flake/flakeref.hh>
#include <lix/libexpr/get-drvs.hh>
#include <lix/libutil/input-accessor.hh>
#include <lix/libutil/json.hh>
#include <lix/libutil/logging.hh>
#include <lix/libexpr/nixexpr.hh>
#include <lix/libutil/ref.hh>
#include <lix/libstore/store-api.hh>
#include <lix/libexpr/symbol-table.hh>
#include <lix/libutil/types.hh>
#include <lix/libexpr/value.hh>
#include <lix/libutil/terminal.hh>
#include <exception>
#include <numeric>
#include <optional>
#include <sstream>
#include <string>
#include <string_view>
#include <utility>

#include "worker.hh"
#include "drv.hh"
#include "buffered-io.hh"
#include "eval-args.hh"

static nix::Value releaseExprTopLevelValue(nix::EvalState &state,
                                           nix::Bindings &autoArgs,
                                           MyArgs &args) {
    nix::Value vTop;

    if (args.fromArgs) {
        nix::Expr &e = state.ctx.parseExprFromString(args.releaseExpr,
                                                     nix::CanonPath::fromCwd());
        vTop = state.eval(e);
    } else {
        vTop = state.evalFile(
            state.aio.blockOn(nix::lookupFileArg(state.ctx, args.releaseExpr))
                .unwrap());
    }

    return state.autoCallFunction(autoArgs, vTop, {});
}

static std::string attrPathJoin(nix::JSON input) {
    return std::accumulate(input.begin(), input.end(), std::string(),
                           [](std::string ss, std::string s) {
                               // Escape token if containing dots
                               if (s.find(".") != std::string::npos) {
                                   s = "\"" + s + "\"";
                               }
                               return ss.empty() ? s : ss + "." + s;
                           });
}

static std::optional<Constituents>
readConstituents(const nix::Value *v, nix::box_ptr<nix::EvalState> &state,
                 nix::ref<nix::eval_cache::CachingEvaluator> &evaluator) {
    auto a = v->attrs()->get(state->ctx.symbols.create("_hydraAggregate"));
    if (a && state->forceBool(a->value, a->pos,
                              "while evaluating the "
                              "`_hydraAggregate` attribute")) {
        std::vector<std::string> constituents;
        std::vector<std::string> namedConstituents;
        auto a = v->attrs()->get(state->ctx.symbols.create("constituents"));
        if (!a)
            state->ctx.errors
                .make<nix::EvalError>("derivation must have a ‘constituents’ "
                                      "attribute")
                .debugThrow(nix::always_progresses); // we can't have a debugger here

        nix::NixStringContext context;
        state->coerceToString(a->pos, a->value, context,
                              "while evaluating the `constituents` attribute",
                              nix::StringCoercionMode::ToString, false);
        for (auto &c : context)
            std::visit(nix::overloaded{
                           [&](const nix::NixStringContextElem::Built &b) {
                               constituents.push_back(
                                   b.drvPath.to_string(*evaluator->store));
                           },
                           [&](const nix::NixStringContextElem::Opaque &) {},
                           [&](const nix::NixStringContextElem::DrvDeep &) {},
                       },
                       c.raw);

        state->forceList(a->value, a->pos,
                         "while evaluating the "
                         "`constituents` attribute");
        for (unsigned int n = 0; n < a->value.listSize(); ++n) {
            auto v = a->value.listElems()[n];
            state->forceValue(v, nix::noPos);
            if (v.type() == nix::nString)
                namedConstituents.emplace_back(v.str());
        }

        return Constituents(constituents, namedConstituents);
    }

    return std::nullopt;
}

void worker(nix::AutoCloseFD &to, nix::AutoCloseFD &from, MyArgs &args)
try {
    // Increase the default stack size for the evaluator and for
    // libstdc++'s std::regex.
    nix::ensureStackSizeAtLeast(64 * 1024 * 1024);

#if HAVE_BOEHMGC
    // We are doing the garbage collection by killing forks.
    GC_disable();

    // There is some memory usage overhead on top of the GC heap size.
    // A simple model using fixed and proportional overhead already gives reasonable results.
    // The parameters have been determined experimentally using evaluation of nixpkgs.
    size_t maxGcHeapSize = std::ldexp(std::max(0.0, 0.9 * (static_cast<double>(args.maxMemorySize) - 150.0)), 20);
    if (GC_get_heap_size() > maxGcHeapSize) {
        throw nix::Error("The heap is too large. Increase %s or reduce %s.",
                         "--max-memory-size", "GC_INITIAL_HEAP_SIZE");
    }
#endif

    auto & aio = args.aio();

    auto evalStore = aio.blockOn(args.evalStoreUrl
                         ? nix::openStore(*args.evalStoreUrl)
                         : nix::openStore());
    auto evaluator =
        nix::make_ref<nix::eval_cache::CachingEvaluator>(
            aio, args.searchPath, evalStore);
    nix::Bindings &autoArgs = *args.getAutoArgs(*evaluator);

    nix::Value vRoot = [&]() {
        auto state = evaluator->begin(aio);
        if (args.flake) {
            auto [flakeRef, fragment, outputSpec] =
                nix::parseFlakeRefWithFragmentAndExtendedOutputsSpec(
                    args.releaseExpr, nix::absPath("."));
            nix::InstallableFlake flake{
                {}, evaluator, std::move(flakeRef), fragment, outputSpec,
                {}, {},        args.lockFlags};

            return flake.toValue(*state).first;
        } else {
            return releaseExprTopLevelValue(*state, autoArgs, args);
        }
    }();

    LineReader fromReader(from.release());
    auto state = evaluator->begin(aio);

    while (true) {
        /* Wait for the collector to send us a job name. */
        if (tryWriteLine(to.get(), "next") < 0) {
            return; // main process died
        }

        auto s = fromReader.readLine();
        if (s == "exit") {
            break;
        }
        if (!s.starts_with("do ")) {
            std::cerr << "worker error: received invalid command '" << s << "'\n";
            abort();
        }
        auto path = nix::json::parse(s.substr(3));
        auto attrPathS = attrPathJoin(path);

        /* Evaluate it and send info back to the collector. */
        nix::JSON reply =
            nix::JSON{{"attr", attrPathS}, {"attrPath", path}};
        try {
            auto vTmp =
                nix::findAlongAttrPath(*state, attrPathS, autoArgs, vRoot)
                    .first;

            nix::Value v = state->autoCallFunction(autoArgs, vTmp, {});

            if (v.type() == nix::nAttrs) {
                if (auto drvInfo = nix::getDerivation(*state, v, false)) {
                    std::optional<Constituents> maybeConstituents;
                    if (args.constituents) {
                        maybeConstituents =
                            readConstituents(&v, state, evaluator);
                    }
                    auto drv = Drv(attrPathS, *state, *drvInfo, args,
                                   maybeConstituents);
                    reply.update(drv);

                    /* Register the derivation as a GC root.  !!! This
                       registers roots for jobs that we may have already
                       done. */
                    register_gc_root(args.gcRootsDir, drv.drvPath, evaluator->store, aio);
                } else {
                    auto attrs = nix::JSON::array();
                    bool recurse =
                        args.forceRecurse ||
                        path.size() == 0; // Dont require `recurseForDerivations
                                          // = true;` for top-level attrset

                    for (auto &i :
                         v.attrs()->lexicographicOrder(evaluator->symbols)) {
                        const std::string_view name = evaluator->symbols[i->name];
                        attrs.emplace_back(name);

                        if (name == "recurseForDerivations" &&
                            !args.forceRecurse) {
                            auto attrv = v.attrs()->get(
                                evaluator->symbols.sym_recurseForDerivations);
                            recurse = state->forceBool(
                                attrv->value, attrv->pos,
                                "while evaluating recurseForDerivations");
                        }
                    }
                    if (recurse)
                        reply["attrs"] = std::move(attrs);
                    else
                        reply["attrs"] = nix::JSON::array();
                }
            } else {
                // We ignore everything that cannot be build
                reply["attrs"] = nix::JSON::array();
            }
        } catch (nix::EvalError &e) {
            auto err = e.info();
            std::ostringstream oss;
            nix::showErrorInfo(oss, err, nix::loggerSettings.showTrace.get());
            auto msg = oss.str();

            // Transmits the error we got from the previous evaluation
            // in the JSON output.
            reply["error"] = nix::filterANSIEscapes(msg, true);
            // Don't forget to print it into the STDERR log, this is
            // what's shown in the Hydra UI.
            std::cerr << msg << "\n";
        } catch (const nix::Interrupted &) {
            throw;
        } catch ( // NOLINT(lix-foreign-exceptions)
            const std::exception &e) { // FIXME: for some reason the catch block
                                       // above, doesn't trigger on macOS (?)
            auto msg = e.what();
            reply["error"] = nix::filterANSIEscapes(msg, true);
            std::cerr << msg << "\n";
        }

        if (tryWriteLine(to.get(), reply.dump()) < 0) {
            return; // main process died
        }

        // If the memory limit is exceeded, exit.
        // The collector will start a new process.
#if HAVE_BOEHMGC
        if (GC_get_heap_size() > maxGcHeapSize) {
            break;
        }
#else
        struct rusage r;
        getrusage(RUSAGE_SELF, &r);
        if ((size_t)r.ru_maxrss > args.maxMemorySize * 1024) {
            break;
        }
#endif
    }

    if (tryWriteLine(to.get(), "restart") < 0) {
        return; // main process died
    };
} catch (const nix::Interrupted &) {
    // The coordinator should get the interrupt too, so it doesn't need to be logged.
} catch (nix::Error &e) {
    nix::JSON err;
    auto msg = e.msg();
    err["error"] = nix::filterANSIEscapes(msg, true);
    printError("%1%", nix::Uncolored(msg));
    if (tryWriteLine(to.get(), err.dump()) < 0) {
        return; // main process died
    };
    // Don't forget to print it into the STDERR log, this is
    // what's shown in the Hydra UI.
    if (tryWriteLine(to.get(), "restart") < 0) {
        return; // main process died
    }
}
