/*
 * Copyright 2019 Bloomberg Finance LP
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef INCLUDED_BUILDBOXRUN_USERCHROOT
#define INCLUDED_BUILDBOXRUN_USERCHROOT

#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_mergeutil.h>
#include <buildboxcommon_runner.h>
#include <buildboxcommon_runnerutils.h>
#include <buildboxcommon_systemutils.h>

#include <filesystem>

namespace buildboxcommon {
namespace buildboxrun {
namespace userchroot {

class UserChrootRunner : public Runner {
  public:
    typedef std::vector<
        std::pair<std::filesystem::path, std::filesystem::path>>
        MappedPaths;

    /**
     * Run a command in a userchroot sandbox.
     */
    ActionResult execute(const Command &command, const Digest &inputRootDigest,
                         const Platform &platform) override;

    // Parse custom arguments for this runner
    bool parseArg(std::string_view arg) override;

    // Print custom usage
    void printSpecialUsage() override;

    const std::string &user() { return d_user; }
    const std::string &group() { return d_group; }
    const std::string &userchrootBin() { return d_userchroot_bin; }
    const std::unordered_map<std::string, MappedPaths> &
    mappedPathsByProperties()
    {
        return d_mappedPathsByProperties;
    }

  private:
    // location of userchroot binary
    std::string d_userchroot_bin = "userchroot";
    // location of sudo binary
    std::string d_sudoBin = "sudo";
    // location of find binary
    std::string d_findBin = "find";
    // location of chmod binary
    std::string d_chmodBin = "chmod";
    // Run `userchroot` as this user, i.e., sudo -u
    std::string d_user;
    // Run `userchroot` as this group, i.e., sudo -g
    std::string d_group;
    // Vector of host paths to copy/hardlink into the chroot at construction
    // time.
    MappedPaths d_mappedPaths;
    // Conditionally inject file(s) into chroot by platform properties
    std::unordered_map<std::string, MappedPaths> d_mappedPathsByProperties;

    void mergeDigests(const Digest &inputDigest, const Digest &chrootDigest,
                      Digest *mergedRootDigest,
                      const Command_OutputDirectoryFormat
                          &chrootDirectoryFormat = Command::DIRECTORY_ONLY);

    void uploadMissingBlobs(const digest_string_map &data) const;

    ProcessCredentials getProcessCredentials() const;

    // Optional. If set, blobs in the merged chroot will be uploaded using this
    // instance name.
    // (This is useful in order to avoid those blobs from being forwarded to a
    // remote in case of being connected to a casd proxy instance.)
    std::string d_localOnlyCasInstanceName;
    ConnectionOptions d_casRemote;

    // This mode will be set in the root directory of the merged trees.
    // That field will signal casd to chmod() the root of the staged
    // directory with a value that allows userchroot to work.
    static const google::protobuf::uint32 s_rootDirectoryRequiredUnixMode =
        0755;

    // Call `LocalCAS.FetchTree()` for the given Digest.
    void fetchChrootIntoLocalCas(const buildboxcommon::Digest &digest) const;

    // If --user USER is specified, the userchroot command might be executed
    // as a different user and directories with limited permissions might be
    // created. Call `chown` on those directories so they can be correctly
    // unstaged.
    std::vector<std::string>
    createPreUnstageCommand(const std::string &stagedDirPath) const;

    // Helper function to inject file(s) at `src` to `chrootPath`
    void injectPath(const std::filesystem::path &srcPath,
                    const std::filesystem::path &chrootPath);

  protected:
    // (Methods marked `protected` for unit testing).

    // Certain directories must be present in all chroots.
    // Check whether that is the case by looking in the staged directory
    // and otherwise throw.
    void assertRequiredDirectoriesExist() const;

    // Clear environment to ensure it is not leaked into userchroot
    inline std::optional<std::unordered_map<std::string, std::string>>
    commandEnvironment() override
    {
        return std::optional(std::unordered_map<std::string, std::string>());
    }

    // Attempt to merge two trees. The root of the result will have `UnixMode`
    // set in its node properties.
    // On errors returns an empty Digest.
    static Digest mergeTreesAndSetRootDirectoryProperty(
        const DirectoryTree &inputTree, const DirectoryTree &chrootTree,
        digest_string_map *mergedDirectoryBlobs);
};

} // namespace userchroot
} // namespace buildboxrun
} // namespace buildboxcommon
#endif
