#include "download.h"

#include <saas/library/http_utils/http_utils.h>

#include <robot/library/oxygen/base/generic/retry.h>
#include <kernel/yt/logging/log.h>
#include <robot/library/oxygen/base/protos/recipe.pb.h>
#include <robot/library/oxygen/base/system/fs_utils.h>
#include <robot/library/oxygen/base/system/rsync.h>
#include <robot/library/oxygen/indexer/processor/protos/config.pb.h>
#include <robot/library/skynet/skynet.h>

#include <google/protobuf/text_format.h>

#include <util/datetime/base.h>
#include <util/folder/path.h>
#include <util/folder/tempdir.h>
#include <util/generic/hash.h>
#include <util/generic/hash_set.h>
#include <util/generic/yexception.h>
#include <util/stream/file.h>
#include <util/string/join.h>
#include <library/cpp/string_utils/url/url.h>
#include <util/system/fs.h>

using namespace NOxygen;
using namespace NSaas;

namespace {

    bool DownloadViaRsync(const TString& rsyncPath, const TFsPath& destination, bool throwOnError, bool quiet, std::function<void(const TString& path)> f) {
        f(rsyncPath);
        auto rsyncF = [&]() { DoRsync(rsyncPath, destination, TRsyncOptions().WithQuiet(quiet)); };
        return DoWithRetry(rsyncF, TRetryOptions(3, TDuration::Seconds(5)), throwOnError);
    }

    bool DownloadViaHttpClient(const TString& address, const TFsPath& destination, bool throwOnError, std::function<void(const TString& path)> f) {
        f(address);
        MkDirs(destination);
        auto getF = [&]() { NSaas::DoDownloadFileViaHttp(address, destination); };
        return DoWithRetry(getF, TRetryOptions(3, TDuration::Seconds(5)), throwOnError);
    }

    void CopyLocalFile(const TFsPath& source, const TFsPath& destination, bool throwOnError) {
        const TFsPath& resolved = source.IsAbsolute() ? source : TFsPath::Cwd() / source;
        bool created = NFs::SymLink(resolved, destination);
        if (!created) {
            L_ERROR << "Failed to make symlink from " << source << " to " << destination;
            if (throwOnError) {
                ythrow yexception() << "Failed to make symlink from " << source << " to " << destination;
            }
        } else {
            L_INFO << "symlink " << source << ' ' << destination;
        }
    }

    void MoveFiles(const TFsPath& sourceDir, const TFsPath& destinationDir) {
        TVector<TFsPath> toMove;
        sourceDir.List(toMove);
        L_INFO << "About to move " << toMove.size() << " files from " << sourceDir << " to " << destinationDir;
        for (const auto& source : toMove) {
            TFsPath dest = destinationDir / source.GetName();
            if (dest.IsDirectory()) { // Temporary workaround for 'current-stable' directory
                RmRf(dest);
            }
            L_INFO << "mv " << source << " " << dest;
            source.RenameTo(dest);
        }
        TVector<TString> leftFiles;
        sourceDir.ListNames(leftFiles);
        Y_ENSURE(leftFiles.empty(), "Expected all files moved from " << sourceDir << ": " << JoinSeq(", ", leftFiles));
    }

} // namespace


TSandboxDownloader::TSandboxDownloader(const TFsPath& dir, const TSandboxOptions& options, bool keepDownloadDir)
    : Dir(dir)
    , Options(options)
    , ResType2Info(CreateResType2Info(dir, options))
    , DownloadDir(Dir / ".sandbox-skynet-tmp")
    , KeepDownloadDir(keepDownloadDir)
{
    if (!KeepDownloadDir) {
        RmRf(DownloadDir);
    }
    MkDirs(DownloadDir);
}

TSandboxDownloader::~TSandboxDownloader() {
    if (!KeepDownloadDir) {
        RmRf(DownloadDir);
    }
}

void TSandboxDownloader::SetCustomJsonConfigPath(const TString& path) {
    Options.UseCustomJsonConfigPath = true;
    Options.CustomJsonConfigPath = path;
}

bool TSandboxDownloader::Download(const TSandboxFileInfoProto& sandboxInfo, const TFsPath& destination,
                                  bool throwOnError, bool quiet, bool keepSkyNetDownloadDir, std::function<void(const TString& path)> f)
{
    if (!Options.UseSandbox) {
        return false;
    }
    const auto ptr = ResType2Info.FindPtr(sandboxInfo.GetResourceType());
    if (!ptr) {
        if (!quiet) {
            L_WARN << "No sandbox info found for resource `" << sandboxInfo.GetResourceType() << '`';
        }
        return false;
    }
    const auto& resInfo = *ptr;

    if (Options.UseSkyNet) {
        if (!ForceRsyncResourceTypes.contains(sandboxInfo.GetResourceType()) &&
                DownloadResourceViaSkyNetIfNeeded(resInfo))
        {
            MoveSkyNetDownloadedFileToCatalog(sandboxInfo, destination, keepSkyNetDownloadDir);
            return true;
        }
        L_WARN << "Failed to download " << resInfo.ResourceType << " via skynet, falling back to rsync";
        ForceRsyncResourceTypes.insert(sandboxInfo.GetResourceType());
    }

    const TString& taskRsyncPath = resInfo.RsyncPath;
    TString rsyncPath;
    if (sandboxInfo.GetRelativePath().size()) {
        rsyncPath = taskRsyncPath + "/" + sandboxInfo.GetRelativePath();
    } else {
        rsyncPath = taskRsyncPath;
    }
    return DownloadViaRsync(rsyncPath, destination, throwOnError, quiet, f);
}

bool TSandboxDownloader::DownloadResourceViaSkyNetIfNeeded(const NSaas::TSandboxResourceInfo& resInfo) {
    if (DownloadedResourceTypes.contains(resInfo.ResourceType)) {
        L_INFO << "Already downloaded resource " << resInfo.ResourceType;
        return true;
    }

    L_INFO << "Downloading resource " << resInfo.ResourceType << " via skynet";
    try {
        SkyGet(DownloadDir, resInfo.RbTorrent);
    } catch (const yexception& e) {
        L_ERROR << "Failed to download " << resInfo.ResourceType << " via skynet: " << e.what();
        return false;
    }

    DownloadedResourceTypes.insert(resInfo.ResourceType);
    return true;
}

void TSandboxDownloader::MoveSkyNetDownloadedFileToCatalog(const TSandboxFileInfoProto& sandboxInfo, const TFsPath& destination, bool keepSource) {
    const auto& resInfo = ResType2Info.at(sandboxInfo.GetResourceType());

    const TFsPath tmpResPath = DownloadDir / resInfo.FileName;

    const TFsPath sourceFile = sandboxInfo.GetRelativePath().size()
                    ? tmpResPath / sandboxInfo.GetRelativePath()
                            : tmpResPath;

    if (destination.IsDirectory()) { // Workaround for 'current-stable' directory
        RmRf(destination);
    }

    L_INFO << "mv " << sourceFile << " " << destination;
    if (keepSource) {
        sourceFile.CopyTo(destination, /*force=*/true);
    } else {
        sourceFile.RenameTo(destination);
    }
}

TSandboxDownloader::TResType2Info TSandboxDownloader::CreateResType2Info(const TFsPath& dir, const TSandboxOptions& options) {
    TResType2Info res;
    if (!options.UseSandbox) {
        return res;
    }

    NSaas::TSandboxManager manager(dir);
    if (options.UseCustomJsonConfigPath) {
        manager.EnableCustomJsonPath(true);
        manager.SetCustomJsonPath(options.CustomJsonConfigPath);
    }
    const TVector<NSaas::TSandboxResourceInfo> resources = manager.GetAvailableSandboxResources();
    for (const auto& resInfo : resources) {
        res[resInfo.ResourceType] = resInfo;
    }
    return res;
}


TMaybe<TSkyFileInfo> NSaas::TSandboxDownloader::GetFileInfo(const TSandboxFileInfoProto& sandboxInfo) {
    const auto ptr = ResType2Info.FindPtr(sandboxInfo.GetResourceType());
    if (!ptr) {
        return Nothing();
    }
    const auto& resInfo = *ptr;
    auto infos = SkyFiles(ptr->RbTorrent);
    for (const auto& info : infos) {
        if (info.Name == resInfo.FileName + "/" + sandboxInfo.GetRelativePath()) {
            return info;
        }
    }
    return Nothing();
}

void NSaas::DownloadRecipeFiles(
        const TRecipeConfigProto& recipe,
        const TFsPath& destDir,
        bool throwOnError,
        bool toShard,
        bool quiet,
        bool keepSkyNetDownloadDir,
        bool useSandbox)
{
    auto f = [](const TString& path) {
        L_INFO << "Retrieving recipe file: " << path;
    };
    NSaas::TSandboxOptions sandboxOptions;
    sandboxOptions.UseSandbox = useSandbox;
    DownloadRecipeFilesImpl(recipe, destDir, throwOnError, toShard, quiet, keepSkyNetDownloadDir, sandboxOptions, f);
}

void NSaas::DownloadFiles(const TString& rbTorrent, const TFsPath& destDir, const TFsPath& tmpDir) {
    auto downloadF = [&]() {
        RmRfAndMkDirs(tmpDir);
        SkyGet(tmpDir, rbTorrent);
    };
    DoWithRetry(TStringBuilder() << "Download " << rbTorrent, downloadF, TRetryOptions::Default(), /*throwLast*/ true);

    MoveFiles(tmpDir, destDir);
}

void NSaas::DownloadRecipeFilesImpl(
        const TRecipeConfigProto& recipe,
        const TFsPath& destDir,
        bool throwOnError,
        bool toShard,
        bool quiet,
        bool keepSkyNetDownloadDir,
        const TSandboxOptions& options,
        std::function<void(const TString& path)> callback,
        const TMaybe<std::function<void(const TString& path)>>& afterCallback,
        std::function<bool(const TString& dstPath)> filterCallback)
{
    Y_ENSURE(destDir.Exists(), "Destination dir not found: " << destDir);

    TSandboxDownloader sandboxDownloader(destDir, options, keepSkyNetDownloadDir);

    for (size_t i = 0; i < recipe.FilesSize(); ++i) {
        const TRecipeFileConfigProto& file = recipe.GetFiles(i);

        if (toShard && !file.GetCopyToShard()) {
            continue;
        }

        TFsPath destination = destDir / file.GetName();

        if (filterCallback && !filterCallback(destination)) {
            continue;
        }

        bool downloaded = false;
        if (file.HasSandboxInfo() && sandboxDownloader.Download(file.GetSandboxInfo(), destination, throwOnError, quiet, keepSkyNetDownloadDir, callback)) {
            downloaded = true;
        }
        if (!downloaded && file.HasRsyncPath() &&
                DownloadViaRsync(file.GetRsyncPath(), destination, throwOnError && !file.HasHttpAddress(), quiet, callback))
        {
            downloaded = true;
        }

        if (!downloaded && file.HasHttpAddress() &&
                DownloadViaHttpClient(file.GetHttpAddress(), destDir, throwOnError && !file.HasLocalPath(), callback))
        {
            downloaded = true;
        }

        if (!downloaded && file.HasLocalPath()) {
            callback(file.GetLocalPath());
            CopyLocalFile(file.GetLocalPath(), destination, throwOnError);
        } else if (!downloaded) {
            ythrow yexception() << "Don't know how to download recipe file: " << file.ShortDebugString();
        }

        if (afterCallback) {
            (*afterCallback)(destination);
        }
    }
}

NOxygen::TRecipeConfigProto NSaas::BuildRecipeFromStaticDataOptions(const NOxygen::TStaticDataOptions& staticDataOptions) {
    NOxygen::TRecipeConfigProto recipe;
    Y_ENSURE(staticDataOptions.HasRecipeConfigPath() ^ staticDataOptions.HasRecipeConfig(), "One and only one of RecipeConfigPath and RecipeConfig directives must be present");

    if (staticDataOptions.HasRecipeConfig()) {
        recipe = staticDataOptions.GetRecipeConfig();
    } else {
        const TString& recipeFilePath = staticDataOptions.GetRecipeConfigPath();
        Y_ENSURE(NFs::Exists(recipeFilePath), "can't find file " + recipeFilePath);
        TUnbufferedFileInput fi(recipeFilePath);
        Y_ENSURE(::google::protobuf::TextFormat::ParseFromString(fi.ReadAll(), &recipe), "Incorrect " + recipeFilePath + " file data");
    }
    return recipe;
}
