/*
 * Copyright 2026 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 <buildboxcommon_credentialhelperauthenticator.h>

#include <buildboxcommon_logging.h>
#include <buildboxcommon_systemutils.h>

#include <build/buildbox/credential_helper.pb.h>
#include <google/protobuf/util/json_util.h>

#include <chrono>

namespace buildboxcommon {

CredentialHelperAuthenticator::CredentialHelperAuthenticator(
    const std::string &helperPath)
    : d_helperPath(helperPath)
{
}

bool CredentialHelperAuthenticator::IsBlocking() const { return false; }

grpc::Status CredentialHelperAuthenticator::GetMetadata(
    grpc::string_ref service_url, grpc::string_ref method_name,
    const grpc::AuthContext &,
    std::multimap<grpc::string, grpc::string> *metadata)
{
    // Construct the full URI
    std::string uri = std::string(service_url.data(), service_url.size()) +
                      std::string(method_name.data(), method_name.size());

    BUILDBOX_LOG_TRACE(
        "CredentialHelper: Getting credentials for URI: " << uri);

    // Check cache for valid, non-expired credentials
    // Note: We only cache credentials that have an expiration time.
    // Credentials without expiration are fetched on every request.
    std::optional<CachedCredentials> cachedCreds;
    {
        std::shared_lock<std::shared_mutex> lock(d_cacheMutex);
        auto it = d_cache.find(uri);
        if (it != d_cache.end()) {
            const auto &cached = it->second;
            const auto now = std::chrono::system_clock::now();
            if (cached.expiresAt.has_value() &&
                now < cached.expiresAt.value()) {
                // Make a copy of valid, non-expired credentials
                cachedCreds = cached;
            }
        }
    }

    if (cachedCreds.has_value()) {
        BUILDBOX_LOG_TRACE("CredentialHelper: Using cached credentials for "
                           << uri);
        for (const auto &[key, value] : cachedCreds->headers) {
            metadata->emplace(key, value);
        }
        return grpc::Status::OK;
    }

    // Cache miss or expired: invoke credential helper
    if (d_helperPath.empty()) {
        BUILDBOX_LOG_WARNING(
            "CredentialHelper: No credential helper configured for " << uri);
        return grpc::Status(grpc::StatusCode::UNAUTHENTICATED,
                            "No credential helper configured");
    }

    BUILDBOX_LOG_DEBUG("CredentialHelper: Invoking helper: "
                       << d_helperPath << " for " << uri);

    CachedCredentials credentials;
    if (!invokeHelper(d_helperPath, uri, &credentials)) {
        return grpc::Status(grpc::StatusCode::UNAUTHENTICATED,
                            "Credential helper failed");
    }

    // Cache the credentials ONLY if they have an expiration time.
    if (credentials.expiresAt.has_value()) {
        {
            std::unique_lock<std::shared_mutex> lock(d_cacheMutex);
            d_cache[uri] = credentials;
        }
        BUILDBOX_LOG_DEBUG("CredentialHelper: Cached credentials for " << uri);
    }

    // Add headers to metadata
    for (const auto &[key, value] : credentials.headers) {
        metadata->emplace(key, value);
    }

    return grpc::Status::OK;
}

bool CredentialHelperAuthenticator::invokeHelper(
    const std::string &helperPath, const std::string &uri,
    CachedCredentials *credentials)
{
    // Prepare the request message
    build::buildbox::GetCredentialsRequest request;
    request.set_uri(uri);

    // Serialize request to JSON
    std::string requestJson;
    google::protobuf::util::JsonPrintOptions options;
    options.add_whitespace = false;
    options.preserve_proto_field_names = true;

    const auto jsonStatus = google::protobuf::util::MessageToJsonString(
        request, &requestJson, options);
    if (!jsonStatus.ok()) {
        BUILDBOX_LOG_ERROR("CredentialHelper: Failed to serialize request: "
                           << jsonStatus.message());
        return false;
    }

    // Invoke the helper
    const std::vector<std::string> command = {helperPath, "get"};
    const auto result = SystemUtils::executeCommandWithResult(
        command, true, true, std::nullopt, requestJson,
        std::chrono::milliseconds(5000));

    if (result.d_exitCode != 0) {
        BUILDBOX_LOG_WARNING("CredentialHelper: Helper exited with code "
                             << result.d_exitCode << ": " << helperPath);
        if (!result.d_stdErr.empty()) {
            BUILDBOX_LOG_WARNING(
                "CredentialHelper: Helper stderr: " << result.d_stdErr);
        }
        return false;
    }

    // Parse the response
    build::buildbox::GetCredentialsResponse response;
    const auto parseStatus = google::protobuf::util::JsonStringToMessage(
        result.d_stdOut, &response);
    if (!parseStatus.ok()) {
        BUILDBOX_LOG_ERROR(
            "CredentialHelper: Failed to parse helper response");
        return false;
    }

    // Extract headers from Struct
    // Expected format: {"Authorization": ["Bearer token"]}
    const auto &headersStruct = response.headers();
    for (const auto &[headerName, headerValue] : headersStruct.fields()) {
        // Headers must be lowercase to be added to gRPC metadata
        std::string headerNameLower = headerName;
        std::transform(headerNameLower.begin(), headerNameLower.end(),
                       headerNameLower.begin(), ::tolower);
        if (headerValue.kind_case() == google::protobuf::Value::kListValue) {
            const auto &listValue = headerValue.list_value();
            for (const auto &item : listValue.values()) {
                if (item.kind_case() ==
                    google::protobuf::Value::kStringValue) {
                    credentials->headers.emplace(headerNameLower,
                                                 item.string_value());
                }
            }
        }
        else {
            BUILDBOX_LOG_WARNING(
                "CredentialHelper: Invalid header value format for '"
                << headerName << "', expected array of strings");
        }
    }

    // Parse expiration if present
    if (response.has_expires()) {
        const auto &timestamp = response.expires();
        // Convert google::protobuf::Timestamp to
        // std::chrono::system_clock::time_point
        credentials->expiresAt =
            std::chrono::time_point_cast<std::chrono::system_clock::duration>(
                std::chrono::system_clock::from_time_t(timestamp.seconds()) +
                std::chrono::nanoseconds(timestamp.nanos()));
    }

    return true;
}

} // namespace buildboxcommon
