/*
 * Copyright 2021 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.
 */

#include <buildboxcasd_acinstance.h>
#include <buildboxcommon_logging.h>

using namespace buildboxcasd;
using namespace buildboxcommon;

namespace {

void add_directory_to_request(const Directory &d, FindMissingBlobsRequest *req)
{
    for (const FileNode &f : d.files()) {
        req->add_blob_digests()->CopyFrom(f.digest());
    }
    for (const DirectoryNode &d_node : d.directories()) {
        req->add_blob_digests()->CopyFrom(d_node.digest());
    }
}

void collectFileDigests(const std::shared_ptr<CasInstance> &cas,
                        const Digest &digest, FindMissingBlobsRequest *req)
{
    const auto directory = cas->getDirectory(digest);
    for (const FileNode &f : directory.files()) {
        req->add_blob_digests()->CopyFrom(f.digest());
    }
    for (const DirectoryNode &d_node : directory.directories()) {
        collectFileDigests(cas, d_node.digest(), req);
    }
}

} // namespace

bool AcInstance::hasAllDigests(ActionResult *result)
{
    if (!d_cas) {
        // This action cache instance is not associated with a CAS instance,
        // skip CAS check.
        return true;
    }

    FindMissingBlobsRequest req;
    req.set_instance_name(d_cas->instanceName());
    bool hasMissing = false;

    if (result->has_stdout_digest()) {
        req.add_blob_digests()->CopyFrom(result->stdout_digest());
    }
    if (result->has_stderr_digest()) {
        req.add_blob_digests()->CopyFrom(result->stderr_digest());
    }

    for (const OutputFile &file : result->output_files()) {
        req.add_blob_digests()->CopyFrom(file.digest());
    }

    BatchReadBlobsRequest tree_req;
    tree_req.set_instance_name(d_cas->instanceName());
    BatchReadBlobsResponse tree_res;
    for (const OutputDirectory &out_dir : result->output_directories()) {
        if (out_dir.has_tree_digest()) {
            tree_req.add_digests()->CopyFrom(out_dir.tree_digest());
        }
        else if (out_dir.has_root_directory_digest()) {
            FetchTreeRequest fetchtree_req;
            FetchTreeResponse fetchtree_res;
            fetchtree_req.set_instance_name(d_cas->instanceName());
            fetchtree_req.mutable_root_digest()->CopyFrom(
                out_dir.root_directory_digest());
            const auto status =
                d_cas->FetchTree(fetchtree_req, &fetchtree_res);
            if (status.ok()) {
                collectFileDigests(d_cas, out_dir.root_directory_digest(),
                                   &req);
            }
            else if (status.error_code() == grpc::StatusCode::NOT_FOUND) {
                hasMissing = true;
                BUILDBOX_LOG_DEBUG("Tree with root directory digest="
                                   << out_dir.root_directory_digest()
                                   << " is missing or incomplete");
            }
            else {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "FetchTree for root directory digest="
                        << out_dir.root_directory_digest()
                        << " failed in AcInstance::hasAllDigests: "
                        << status.error_code() << ": "
                        << status.error_message());
            }
        }
    }

    if (!tree_req.digests().empty()) {
        if (!d_cas->BatchReadBlobs(tree_req, &tree_res).ok()) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "BatchReadBlobs failed in AcInstance::hasAllDigests");
        }

        for (const auto &tree : tree_res.responses()) {
            if (tree.status().code() == grpc::StatusCode::NOT_FOUND) {
                hasMissing = true;
                BUILDBOX_LOG_DEBUG("Tree with digest=" + tree.digest().hash() +
                                   " is missing");
                continue;
            }
            else if (tree.status().code() != grpc::StatusCode::OK) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "readBlob for digest="
                        << tree.digest()
                        << " failed in AcInstance::hasAllDigests: "
                        << tree.status().code() << ": "
                        << tree.status().message());
            }
            Tree t;
            t.ParseFromString(tree.data());

            add_directory_to_request(t.root(), &req);
            for (const Directory &d :
                 t.children()) { // children already recursive,
                                 // don't need to go deeper
                add_directory_to_request(d, &req);
            }
        }
    }

    if (!req.blob_digests().empty()) {
        FindMissingBlobsResponse res;

        grpc::Status code = d_cas->FindMissingBlobs(req, &res);

        if (!code.ok()) {
            if (code.error_code() != grpc::StatusCode::CANCELLED) {
                BUILDBOX_LOG_ERROR(
                    "ActionCache unable to check CAS, returning not in CAS");
            }
            else {
                BUILDBOX_LOG_INFO(
                    "ActionCache unable to check CAS, request cancelled");
            }
            return true;
        }
        if (!res.missing_blob_digests().empty()) {
            hasMissing = true;
            BUILDBOX_LOG_INFO("ActionResult missing "
                              << res.missing_blob_digests().size()
                              << " digests");
        }
    }

    return !hasMissing;
}
