#include <build/bazel/remote/asset/v1/remote_asset_mock.grpc.pb.h>

#include <gtest/gtest.h>

#include <buildboxcommon_assetclient.h>
#include <buildboxcommon_grpcclient.h>

#include <memory>

using namespace buildboxcommon;
using namespace testing;

class ClientTestFixture : public AssetClient, public Test {
  protected:
    std::shared_ptr<MockFetchStub> fetchClient =
        std::make_shared<MockFetchStub>();
    std::shared_ptr<MockPushStub> pushClient =
        std::make_shared<MockPushStub>();

    ClientTestFixture() : AssetClient(std::make_shared<GrpcClient>())
    {
        this->init(fetchClient, pushClient);
        this->getGrpcClient()->setInstanceName(clientInstanceName);
        this->getGrpcClient()->setRetryLimit(1);
    }

    /********************** MOCK TEST DATA ********************************/
  public:
    static const std::vector<std::string> uris;
    static const std::string selectedUri;
    static const std::vector<std::pair<std::string, std::string>> qualifiers;
    static const std::vector<Qualifier> returnedQualifiers;
    static const std::string clientInstanceName;
    static const Digest digest;
    static const std::chrono::time_point<std::chrono::steady_clock> time;
    static const google::protobuf::Timestamp protobuf_time;
};

const std::vector<std::string> ClientTestFixture::uris = {"www.bloomberg.com",
                                                          "www.google.com"};

const std::string ClientTestFixture::selectedUri = "www.google.com";

const std::vector<std::pair<std::string, std::string>>
    ClientTestFixture::qualifiers = {{"key1", "value1"}, {"key2", "value2"}};

const std::vector<Qualifier> ClientTestFixture::returnedQualifiers = [] {
    Qualifier q1;
    q1.set_name("key1");
    q1.set_value("value1");

    Qualifier q2;
    q2.set_name("key2");
    q2.set_value("value2");
    return std::vector<Qualifier>({std::move(q1), std::move(q2)});
}();

const std::string ClientTestFixture::clientInstanceName = "CasTestInstance123";

const Digest ClientTestFixture::digest = [] {
    Digest d;
    d.set_hash("123");
    d.set_size_bytes(100);
    return d;
}();

const std::chrono::time_point<std::chrono::steady_clock>
    ClientTestFixture::time = std::chrono::steady_clock::now();

const google::protobuf::Timestamp ClientTestFixture::protobuf_time = [] {
    google::protobuf::Timestamp ts;

    const auto seconds =
        std::chrono::time_point_cast<std::chrono::seconds>(time);

    ts.set_seconds(seconds.time_since_epoch().count());
    ts.set_nanos(static_cast<int32_t>(
        (std::chrono::time_point_cast<std::chrono::nanoseconds>(time) -
         std::chrono::time_point_cast<std::chrono::nanoseconds>(seconds))
            .count()));
    return ts;
}();

MATCHER(FetchRequestMatcher, "")
{
    if (std::vector<std::string>(arg.uris().begin(), arg.uris().end()) !=
        ClientTestFixture::uris) {
        return false;
    }

    if (arg.instance_name() != ClientTestFixture::clientInstanceName) {
        return false;
    }

    for (std::size_t i = 0; i < ClientTestFixture::qualifiers.size(); i++) {
        const Qualifier &qualifier = arg.qualifiers()[i];
        if (qualifier.name() != ClientTestFixture::qualifiers[i].first) {
            return false;
        }

        if (qualifier.value() != ClientTestFixture::qualifiers[i].second) {
            return false;
        }
    }

    return true;
}

TEST_F(ClientTestFixture, FetchBlobTest)
{
    EXPECT_CALL(*fetchClient, FetchBlob(_, FetchRequestMatcher(), _))
        .WillOnce([](::grpc::ClientContext *context,
                     const FetchBlobRequest &request,
                     FetchBlobResponse *response) {
            response->set_uri(ClientTestFixture::selectedUri);
            for (const Qualifier &qualifier :
                 ClientTestFixture::returnedQualifiers) {
                *response->add_qualifiers() = qualifier;
            }

            *response->mutable_blob_digest() = ClientTestFixture::digest;
            *response->mutable_expires_at() = ClientTestFixture::protobuf_time;
            return grpc::Status::OK;
        });

    const FetchResult res = fetchBlob(uris, qualifiers);

    ASSERT_EQ(res.uri, ClientTestFixture::selectedUri);
    ASSERT_EQ(res.qualifiers, ClientTestFixture::qualifiers);
    ASSERT_EQ(res.expiresAt, ClientTestFixture::time);
    ASSERT_EQ(res.digest, ClientTestFixture::digest);
}

TEST_F(ClientTestFixture, FetchDirectoryTest)
{
    EXPECT_CALL(*fetchClient, FetchDirectory(_, FetchRequestMatcher(), _))
        .WillOnce([](::grpc::ClientContext *context,
                     const FetchDirectoryRequest &request,
                     FetchDirectoryResponse *response) {
            response->set_uri(ClientTestFixture::selectedUri);
            for (const Qualifier &qualifier :
                 ClientTestFixture::returnedQualifiers) {
                *response->add_qualifiers() = qualifier;
            }

            *response->mutable_root_directory_digest() =
                ClientTestFixture::digest;
            *response->mutable_expires_at() = ClientTestFixture::protobuf_time;
            return grpc::Status::OK;
        });

    const FetchResult res = fetchDirectory(uris, qualifiers);

    ASSERT_EQ(res.uri, ClientTestFixture::selectedUri);
    ASSERT_EQ(res.qualifiers, ClientTestFixture::qualifiers);
    ASSERT_EQ(res.expiresAt, ClientTestFixture::time);
    ASSERT_EQ(res.digest, ClientTestFixture::digest);
}

MATCHER(PushBlobRequestMatcher, "")
{
    if (std::vector<std::string>(arg.uris().begin(), arg.uris().end()) !=
        ClientTestFixture::uris) {
        return false;
    }

    if (arg.instance_name() != ClientTestFixture::clientInstanceName) {
        return false;
    }

    if (arg.blob_digest() != ClientTestFixture::digest) {
        return false;
    }

    for (std::size_t i = 0; i < ClientTestFixture::qualifiers.size(); i++) {
        const Qualifier &qualifier = arg.qualifiers()[i];
        if (qualifier.name() != ClientTestFixture::qualifiers[i].first) {
            return false;
        }

        if (qualifier.value() != ClientTestFixture::qualifiers[i].second) {
            return false;
        }
    }

    return true;
}

MATCHER(PushDirectoryRequestMatcher, "")
{
    if (std::vector<std::string>(arg.uris().begin(), arg.uris().end()) !=
        ClientTestFixture::uris) {
        return false;
    }

    if (arg.instance_name() != ClientTestFixture::clientInstanceName) {
        return false;
    }

    if (arg.root_directory_digest() != ClientTestFixture::digest) {
        return false;
    }

    for (std::size_t i = 0; i < ClientTestFixture::qualifiers.size(); i++) {
        const Qualifier &qualifier = arg.qualifiers()[i];
        if (qualifier.name() != ClientTestFixture::qualifiers[i].first) {
            return false;
        }

        if (qualifier.value() != ClientTestFixture::qualifiers[i].second) {
            return false;
        }
    }

    return true;
}

TEST_F(ClientTestFixture, PushBlobTest)
{
    EXPECT_CALL(*pushClient, PushBlob(_, PushBlobRequestMatcher(), _))
        .WillOnce(Return(grpc::Status::OK));

    pushBlob(uris, qualifiers, digest);
}

TEST_F(ClientTestFixture, PushDirectoryTest)
{
    EXPECT_CALL(*pushClient,
                PushDirectory(_, PushDirectoryRequestMatcher(), _))
        .WillOnce(Return(grpc::Status::OK));

    pushDirectory(uris, qualifiers, digest);
}

TEST_F(ClientTestFixture, TryFetchBlobSuccessTest)
{
    EXPECT_CALL(*fetchClient, FetchBlob(_, FetchRequestMatcher(), _))
        .WillOnce([](::grpc::ClientContext *context,
                     const FetchBlobRequest &request,
                     FetchBlobResponse *response) {
            response->set_uri(ClientTestFixture::selectedUri);
            for (const Qualifier &qualifier :
                 ClientTestFixture::returnedQualifiers) {
                *response->add_qualifiers() = qualifier;
            }

            *response->mutable_blob_digest() = ClientTestFixture::digest;
            *response->mutable_expires_at() = ClientTestFixture::protobuf_time;
            return grpc::Status::OK;
        });

    const auto res = tryFetchBlob(uris, qualifiers);

    ASSERT_TRUE(res.has_value());
    if (res.has_value()) {
        ASSERT_EQ(res->uri, ClientTestFixture::selectedUri);
        ASSERT_EQ(res->qualifiers, ClientTestFixture::qualifiers);
        ASSERT_EQ(res->expiresAt, ClientTestFixture::time);
        ASSERT_EQ(res->digest, ClientTestFixture::digest);
    }
}

TEST_F(ClientTestFixture, TryFetchBlobNotFoundTest)
{
    EXPECT_CALL(*fetchClient, FetchBlob(_, FetchRequestMatcher(), _))
        .WillOnce(Return(grpc::Status(grpc::StatusCode::NOT_FOUND,
                                      "Asset not found in local cache")));

    // tryFetchBlob should not throw on NOT_FOUND
    const auto res = tryFetchBlob(uris, qualifiers);

    // Should return nullopt for NOT_FOUND
    ASSERT_FALSE(res.has_value());
}

TEST_F(ClientTestFixture, TryFetchBlobOtherErrorTest)
{
    EXPECT_CALL(*fetchClient, FetchBlob(_, FetchRequestMatcher(), _))
        .WillOnce(Return(grpc::Status(grpc::StatusCode::INTERNAL,
                                      "Internal server error")));

    // tryFetchBlob should still throw on other errors (not NOT_FOUND)
    ASSERT_THROW(tryFetchBlob(uris, qualifiers), std::exception);
}

// Helper function to create a digest for testing
Digest makeTestDigest(const std::string &hash, int64_t size = 100)
{
    Digest d;
    d.set_hash(hash);
    d.set_size_bytes(size);
    return d;
}

// Matcher for checking PushBlobRequest with referenced blobs
MATCHER_P3(PushBlobRequestMatcher, expectedDigest, expectedReferencedBlobsSize,
           expectedReferencedDirsSize, "")
{
    if (arg.blob_digest().hash() != expectedDigest.hash() ||
        arg.blob_digest().size_bytes() != expectedDigest.size_bytes()) {
        return false;
    }
    if (arg.references_blobs_size() != expectedReferencedBlobsSize) {
        return false;
    }
    if (arg.references_directories_size() != expectedReferencedDirsSize) {
        return false;
    }
    return true;
}

// Matcher for checking PushDirectoryRequest with referenced blobs
MATCHER_P3(PushDirectoryRequestMatcher, expectedDigest,
           expectedReferencedBlobsSize, expectedReferencedDirsSize, "")
{
    if (arg.root_directory_digest().hash() != expectedDigest.hash() ||
        arg.root_directory_digest().size_bytes() !=
            expectedDigest.size_bytes()) {
        return false;
    }
    if (arg.references_blobs_size() != expectedReferencedBlobsSize) {
        return false;
    }
    if (arg.references_directories_size() != expectedReferencedDirsSize) {
        return false;
    }
    return true;
}

// Test pushBlob with referenced blobs
TEST_F(ClientTestFixture, PushBlobWithReferencedBlobs)
{
    const Digest blobDigest = makeTestDigest("testblob123");
    const Digest refBlob1 = makeTestDigest("refblob1");
    const Digest refBlob2 = makeTestDigest("refblob2");
    const Digest refDir1 = makeTestDigest("refdir1");

    std::vector<Digest> referencedBlobs = {refBlob1, refBlob2};
    std::vector<Digest> referencedDirectories = {refDir1};

    // Expect pushBlob call with references_blobs and references_directories
    // populated
    EXPECT_CALL(*pushClient,
                PushBlob(_, PushBlobRequestMatcher(blobDigest, 2, 1), _))
        .WillOnce(Return(grpc::Status::OK));

    // Test the enhanced pushBlob with referenced blobs/directories
    ASSERT_NO_THROW(pushBlob(uris, qualifiers, blobDigest, referencedBlobs,
                             referencedDirectories));
}

// Test pushDirectory with referenced blobs
TEST_F(ClientTestFixture, PushDirectoryWithReferencedBlobs)
{
    const Digest rootDigest = makeTestDigest("rootdir123");
    const Digest refBlob1 = makeTestDigest("refblob1");
    const Digest refDir1 = makeTestDigest("refdir1");
    const Digest refDir2 = makeTestDigest("refdir2");

    std::vector<Digest> referencedBlobs = {refBlob1};
    std::vector<Digest> referencedDirectories = {refDir1, refDir2};

    // Expect pushDirectory call with references populated
    EXPECT_CALL(
        *pushClient,
        PushDirectory(_, PushDirectoryRequestMatcher(rootDigest, 1, 2), _))
        .WillOnce(Return(grpc::Status::OK));

    // Test the enhanced pushDirectory with referenced blobs/directories
    ASSERT_NO_THROW(pushDirectory(uris, qualifiers, rootDigest,
                                  referencedBlobs, referencedDirectories));
}

// Test pushBlob with empty references (backward compatibility)
TEST_F(ClientTestFixture, PushBlobWithEmptyReferences)
{
    const Digest blobDigest = makeTestDigest("testblob456");

    // Expect pushBlob call with empty references
    EXPECT_CALL(*pushClient,
                PushBlob(_, PushBlobRequestMatcher(blobDigest, 0, 0), _))
        .WillOnce(Return(grpc::Status::OK));

    // Test pushBlob with default empty references (backward compatibility)
    ASSERT_NO_THROW(pushBlob(uris, qualifiers, blobDigest));
}

// Test pushDirectory with empty references (backward compatibility)
TEST_F(ClientTestFixture, PushDirectoryWithEmptyReferences)
{
    const Digest rootDigest = makeTestDigest("rootdir456");

    // Expect pushDirectory call with empty references
    EXPECT_CALL(
        *pushClient,
        PushDirectory(_, PushDirectoryRequestMatcher(rootDigest, 0, 0), _))
        .WillOnce(Return(grpc::Status::OK));

    // Test pushDirectory with default empty references (backward
    // compatibility)
    ASSERT_NO_THROW(pushDirectory(uris, qualifiers, rootDigest));
}
