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

#ifndef BUILDBOXCOMMON_OCICLIENT_H
#define BUILDBOXCOMMON_OCICLIENT_H

#include <buildboxcommon_assetclient.h>
#include <buildboxcommon_casclient.h>
#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_futuregroup.h>
#include <buildboxcommon_httpclient.h>
#include <buildboxcommon_ocimanifest.h>
#include <buildboxcommon_temporarydirectory.h>
#include <buildboxcommon_temporaryfile.h>

#include <ThreadPool.h>

#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace buildboxcommon {

constexpr const char *const OCI_MANIFEST_MEDIA_TYPE =
    "application/vnd.oci.image.manifest.v1+json";

constexpr const char *const DOCKER_MANIFEST_MEDIA_TYPE =
    "application/vnd.docker.distribution.manifest.v2+json";

constexpr const char *const OCI_IMAGE_INDEX_MEDIA_TYPE =
    "application/vnd.oci.image.index.v1+json";

constexpr const char *const DOCKER_MANIFEST_LIST_MEDIA_TYPE =
    "application/vnd.docker.distribution.manifest.list.v2+json";

constexpr const char *const MANIFEST_MEDIA_TYPE =
    "application/vnd.oci.image.manifest.v1+json, "
    "application/vnd.docker.distribution.manifest.v2+json, "
    "application/vnd.oci.image.index.v1+json, "
    "application/vnd.docker.distribution.manifest.list.v2+json";

constexpr const char *const DOCKER_LAYER_MEDIA_TYPE =
    "application/vnd.docker.image.rootfs.diff.tar.gzip";

constexpr const char *const OCI_LAYER_MEDIA_TYPE =
    "application/vnd.oci.image.layer.v1.tar+gzip";

constexpr int SHA256_PREFIX_LEN = 7; // Length of "sha256:" prefix

// Asset service qualifier constants
constexpr const char *const BUILDBOX_OCI_VERSION_QUALIFIER =
    "buildbox_oci_version";
constexpr const char *const RESOURCE_TYPE_QUALIFIER = "resource_type";
constexpr const char *const PLATFORM_OS_QUALIFIER = "platform_os";
constexpr const char *const PLATFORM_ARCH_QUALIFIER = "platform_arch";

// Asset service version for layer caching
static constexpr const char *BUILDBOX_OCI_VERSION = "v1";

// Default platform values
constexpr const char *const DEFAULT_PLATFORM_OS = "linux";
constexpr const char *const DEFAULT_PLATFORM_ARCH = "amd64";

// Parsed result for a Docker/OCI image URI
struct OciUriComponents {
    std::string registry;
    std::string repository;
    std::optional<std::string> tag;          // tag, if present
    std::optional<std::string> sha256Digest; // sha256 digest, if present

    // An image of the form "registry/image-name@sha256:hash" is a fully
    // stable/resolved identifier
    constexpr bool isFullyResolved() const { return sha256Digest.has_value(); }
};

// OCI Registry error codes
class OciRegistryException : public std::runtime_error {
  public:
    explicit OciRegistryException(const std::string &message)
        : std::runtime_error(message)
    {
    }
};

class OciInvalidUriException : public OciRegistryException {
  public:
    explicit OciInvalidUriException(const std::string &message)
        : OciRegistryException(message)
    {
    }
};

class OciRegistryApiCallException : public OciRegistryException {
  public:
    explicit OciRegistryApiCallException(const std::string &message)
        : OciRegistryException(message)
    {
    }
};

class OciUnsupportedPlatformException : public OciRegistryException {
  public:
    explicit OciUnsupportedPlatformException(const std::string &message)
        : OciRegistryException(message)
    {
    }
};

class OciTarExtractionException : public OciRegistryException {
  public:
    explicit OciTarExtractionException(const std::string &message)
        : OciRegistryException(message)
    {
    }
};

// OCI client for parsing image URIs and registry interaction
class OciClient {
  private:
    using Qualifiers = std::vector<std::pair<std::string, std::string>>;

  public:
    // Constructor that takes HTTPClient, CASClient, and AssetClient pointers
    explicit OciClient(const std::shared_ptr<HTTPClient> &httpClient,
                       const std::shared_ptr<CASClient> &casClient,
                       const std::shared_ptr<AssetClient> &assetClient);

    // Constructor that takes HTTPClient, CASClient, AssetClient pointers,
    // optional tar binary path, and optional auth token path
    explicit OciClient(const std::shared_ptr<HTTPClient> &httpClient,
                       const std::shared_ptr<CASClient> &casClient,
                       const std::shared_ptr<AssetClient> &assetClient,
                       const std::optional<std::string> &tarBinaryPath,
                       const std::optional<std::string> &authTokenPath);

    /**
     * Parse a Docker/OCI image URI into registry, repository, and reference.
     * Handles tags, digests, and implicit defaults (docker.io, library,
     * latest). Throws std::invalid_argument on malformed URIs.
     */
    static OciUriComponents parseOCIuri(const std::string &originalUri);

    // Get OCI manifest for the given OCI URI
    // Validates that the URI references a valid OCI image and returns the
    // manifest. Only sha256 digests are supported to ensure deterministic
    // actions. The sha256 digest is verified using content verification.
    // For manifest lists/indexes, selects the manifest matching the specified
    // platform (defaults to linux/amd64).
    OciManifest
    getOCIManifest(const OciUriComponents &components,
                   const std::string &platformOs = "linux",
                   const std::string &platformArchitecture = "amd64");

    // Verify content blob matches expected SHA256 digest
    // TODO: Remove the restriction that the digest function must be SHA256
    static bool verifyContent(const std::string &blob,
                              const std::string &expectedSha256);

    // Stream and extract layer content directly to a directory
    // Downloads the layer from the registry and pipes it directly to
    // tar for extraction, with digest verification
    void streamAndExtractLayer(TemporaryDirectory *outputDir,
                               const std::string &registryUri,
                               const std::string &repository,
                               const OciManifestLayer &layer);

    // Capture extracted layer directory to CAS - returns both tree and root
    // digests or throws
    std::pair<Digest, Digest>
    captureLayerToDigest(const std::string &layerDir,
                         const OciManifestLayer &layer);

    // Combined stream + extract + capture - returns CAS tree digest or throws
    Digest processLayer(const OciManifest &manifest,
                        const OciManifestLayer &layer);

    // Fetch an OCI/Docker image, process all its layers sequentially, merge
    // them into a single tree digest, and upload the merged blobs to CAS
    Digest getImageTreeDigest(const std::string &ociUri);

    // Fetch an OCI/Docker image, process all its layers sequentially, merge
    // them into a single root digest, and upload the merged blobs to CAS
    Digest getImageRootDigest(const std::string &ociUri);

    // Fetch an OCI/Docker image, process all its layers sequentially, merge
    // them into both tree and root digests, and upload the merged blobs to CAS
    // Returns a pair (tree digest, root digest)
    std::pair<Digest, Digest> getImageDigests(const std::string &ociUri);

    // Overloaded version that accepts an optional ThreadPool for parallel
    // processing
    std::pair<Digest, Digest> getImageDigests(const std::string &ociUri,
                                              ThreadPool *threadPool);

    // Helper function to merge multiple layer tree digests. Takes tree
    // digests, converts each to DirectoryTree, merges them with layer
    // semantics, uploads merged blobs to CAS, and returns both tree and root
    // digests as a pair (tree digest, root digest)
    std::pair<Digest, Digest>
    mergeLayerTrees(const std::vector<Digest> &layerTreeDigests);

  private:
    // HTTP client for making registry API calls
    std::shared_ptr<HTTPClient> d_httpClient;

    // CAS client for interacting with local CAS server
    std::shared_ptr<CASClient> d_casClient;

    // Asset client for caching layer trees
    std::shared_ptr<AssetClient> d_assetClient;

    // Optional tar binary path
    std::optional<std::string> d_tarBinaryPath;

    // Optional auth token file path for Bearer authentication
    std::optional<std::string> d_authTokenPath;

    // Helper method to cleanup pipe file descriptors and child process
    void cleanupTarExtraction(std::array<int, 2> &pipeFd, pid_t pid) const;

    // Helper method to make registry API calls
    // OCI spec recommends not depending on docker-content-digest header but it
    // is still used by all relevant registries
    HTTPResponse
    makeRegistryApiCall(std::string *verifiedDigest,
                        const std::string &mediaType,
                        const std::vector<std::string> &pathParts) const;

    // Helper method to stream registry API calls
    HTTPResponse
    streamRegistryApiCall(const std::vector<std::string> &pathParts,
                          const StreamCallback &callback) const;

    // Helper to get tar command path
    std::string getTarCommand() const;

    // Helper to read auth token from file path
    std::string readAuthToken() const;

    // Helper to upload blob map to CAS after checking for missing blobs
    void uploadBlobMap(const digest_string_map &blobMap) const;

    // Generate asset service URI for a layer
    std::string generateLayerUri(const OciManifest &manifest,
                                 const OciManifestLayer &layer) const;

    // Generate asset service qualifiers for a layer
    Qualifiers generateLayerQualifiers(const OciManifestLayer &layer) const;

    // Cache layer tree in asset service with root digest as referenced
    // directory
    void cacheLayerTree(const OciManifest &manifest,
                        const OciManifestLayer &layer,
                        const Digest &treeDigest, const Digest &rootDigest);

    // Generate canonical image URI using OciUriComponents and SHA256 manifest
    // digest
    std::string
    generateCanonicalImageUri(const OciUriComponents &components,
                              const std::string &manifestDigest) const;

    // Generate image qualifiers for full image caching
    Qualifiers generateImageQualifiers(const OciManifest &manifest) const;

    // Generate image qualifiers without resource type (for fully resolved
    // images)
    Qualifiers generateImageQualifiers() const;

    // Cache entire image tree digest using SHA256 manifest as key
    void cacheImageTreeDigest(const std::string &canonicalImageUri,
                              const Digest &treeDigest,
                              const Digest &rootDigest,
                              const Qualifiers &qualifiers);

    // Helper to fetch blob from asset cache and verify it exists in CAS,
    // returns an optional pair (tree digest, root digest)
    std::optional<std::pair<Digest, Digest>>
    fetchCachedImageDigest(const std::string &imageUri,
                           const Qualifiers &qualifiers) const;

    // Helper to assemble a DirectoryTree using directoryTrees and
    // newDirectoryBlobs in-memory, avoiding the getTree gRPC
    DirectoryTree
    getMergedTree(const Digest &mergedRootDigest,
                  const std::vector<DirectoryTree> &directoryTrees,
                  const digest_string_map &newDirectoryBlobs);
};

} // namespace buildboxcommon

#endif // BUILDBOXCOMMON_OCICLIENT_H
