/*
 * Copyright 2025 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_fslocalassetstorage.h>
#include <buildboxcasd_localrainstance.h>
#include <buildboxcasd_metricnames.h>
#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_fslocalcas.h>
#include <buildboxcommon_temporarydirectory.h>
#include <buildboxcommonmetrics_countingmetricvalue.h>
#include <buildboxcommonmetrics_metricguard.h>
#include <buildboxcommonmetrics_testingutils.h>

#include <filesystem>
#include <grpcpp/server_context.h>
#include <gtest/gtest.h>

using namespace buildboxcasd;
using namespace buildboxcommon;

namespace {
const auto digestFunctionInitializer = []() {
    buildboxcommon::DigestGenerator::init();
    return 0;
}();
} // namespace

class LocalRaInstanceTest : public ::testing::Test {
  protected:
    LocalRaInstanceTest()
        : storage_root_directory(), cas_storage(std::make_shared<FsLocalCas>(
                                        storage_root_directory.name())),
          asset_storage(std::make_shared<FsLocalAssetStorage>(
              storage_root_directory.name())),
          ra_instance(asset_storage)
    {
    }

    void SetUp() override
    {
        // Clear metrics before each test
        buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();

        test_digest = DigestGenerator::hash("test_data");
        test_uri = "urn:fdc:example.com:2020:test";

        // Store test blob in CAS
        cas_storage->writeBlob(test_digest, "test_data");
    }

  protected:
    TemporaryDirectory storage_root_directory;
    std::shared_ptr<FsLocalCas> cas_storage;
    std::shared_ptr<FsLocalAssetStorage> asset_storage;
    LocalRaInstance ra_instance;
    grpc::ServerContext server_context;

    Digest test_digest;
    std::string test_uri;
};

TEST_F(LocalRaInstanceTest, PushAndFetchBlob)
{
    // Push blob to asset storage
    PushBlobRequest push_request;
    PushBlobResponse push_response;
    push_request.add_uris(test_uri);
    push_request.mutable_blob_digest()->CopyFrom(test_digest);

    grpc::Status push_status =
        ra_instance.PushBlob(push_request, &push_response);
    ASSERT_TRUE(push_status.ok());

    // Fetch blob from asset storage
    FetchBlobRequest fetch_request;
    FetchBlobResponse fetch_response;
    fetch_request.add_uris(test_uri);

    grpc::Status fetch_status =
        ra_instance.FetchBlob(fetch_request, &fetch_response);
    ASSERT_TRUE(fetch_status.ok());
    EXPECT_EQ(fetch_response.blob_digest().hash(), test_digest.hash());
    EXPECT_EQ(fetch_response.blob_digest().size_bytes(),
              test_digest.size_bytes());

    // Verify metrics were recorded
    EXPECT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue>(
            MetricNames::COUNTER_NAME_RA_FETCH_BLOB_HITS,
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)));
}

TEST_F(LocalRaInstanceTest, FetchNonExistentBlob)
{
    FetchBlobRequest fetch_request;
    FetchBlobResponse fetch_response;
    fetch_request.add_uris("urn:fdc:example.com:2020:nonexistent");

    grpc::Status fetch_status =
        ra_instance.FetchBlob(fetch_request, &fetch_response);
    EXPECT_EQ(fetch_status.error_code(), grpc::StatusCode::NOT_FOUND);

    // Verify miss metric was recorded
    EXPECT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue>(
            MetricNames::COUNTER_NAME_RA_FETCH_BLOB_MISSES,
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)));
}

TEST_F(LocalRaInstanceTest, PushAndFetchDirectory)
{
    // Create a simple directory structure
    Directory test_dir;
    FileNode *file_node = test_dir.add_files();
    file_node->set_name("test_file");
    file_node->mutable_digest()->CopyFrom(test_digest);

    std::string dir_data;
    test_dir.SerializeToString(&dir_data);
    Digest dir_digest = DigestGenerator::hash(dir_data);

    // Store directory in CAS
    cas_storage->writeBlob(dir_digest, dir_data);

    // Push directory to asset storage
    PushDirectoryRequest push_request;
    PushDirectoryResponse push_response;
    push_request.add_uris(test_uri + "_dir");
    push_request.mutable_root_directory_digest()->CopyFrom(dir_digest);

    grpc::Status push_status =
        ra_instance.PushDirectory(push_request, &push_response);
    ASSERT_TRUE(push_status.ok());

    // Fetch directory from asset storage
    FetchDirectoryRequest fetch_request;
    FetchDirectoryResponse fetch_response;
    fetch_request.add_uris(test_uri + "_dir");

    grpc::Status fetch_status =
        ra_instance.FetchDirectory(fetch_request, &fetch_response);
    ASSERT_TRUE(fetch_status.ok());
    EXPECT_EQ(fetch_response.root_directory_digest().hash(),
              dir_digest.hash());
    EXPECT_EQ(fetch_response.root_directory_digest().size_bytes(),
              dir_digest.size_bytes());

    // Verify metrics were recorded
    EXPECT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue>(
            MetricNames::COUNTER_NAME_RA_FETCH_DIRECTORY_HITS,
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)));
}

TEST_F(LocalRaInstanceTest, FetchNonExistentDirectory)
{
    FetchDirectoryRequest fetch_request;
    FetchDirectoryResponse fetch_response;
    fetch_request.add_uris("urn:fdc:example.com:2020:nonexistent_dir");

    grpc::Status fetch_status =
        ra_instance.FetchDirectory(fetch_request, &fetch_response);
    EXPECT_EQ(fetch_status.error_code(), grpc::StatusCode::NOT_FOUND);

    // Verify miss metric was recorded
    EXPECT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue>(
            MetricNames::COUNTER_NAME_RA_FETCH_DIRECTORY_MISSES,
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)));
}

TEST_F(LocalRaInstanceTest, Truncated)
{
    // Push blob to asset storage
    PushBlobRequest push_request;
    PushBlobResponse push_response;
    push_request.add_uris(test_uri);
    push_request.mutable_blob_digest()->CopyFrom(test_digest);

    grpc::Status push_status =
        ra_instance.PushBlob(push_request, &push_response);
    ASSERT_TRUE(push_status.ok());

    // Truncate asset storage entry to check error handling after system crash
    for (const auto &dirEntry : std::filesystem::recursive_directory_iterator{
             storage_root_directory.strname() + "/assets/objects"}) {
        if (dirEntry.is_regular_file()) {
            ASSERT_EQ(truncate(dirEntry.path().c_str(), 0), 0);
        }
    }

    // Fetch blob from asset storage
    FetchBlobRequest fetch_request;
    FetchBlobResponse fetch_response;
    fetch_request.add_uris(test_uri);

    grpc::Status fetch_status =
        ra_instance.FetchBlob(fetch_request, &fetch_response);
    EXPECT_EQ(fetch_status.error_code(), grpc::StatusCode::NOT_FOUND);

    // Verify metrics were recorded
    EXPECT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue>(
            MetricNames::COUNTER_NAME_RA_FETCH_BLOB_MISSES,
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)));
}
