provisioning-platform/crates/catalog-registry/src/oci/client.rs

use std::time::Duration;

use async_trait::async_trait;
use bytes::Bytes;
use chrono::Utc;
use reqwest::Client;
use tracing::debug;

use crate::client::traits::{
    BackendType, DistributionClient, ExtensionClient, LayerInfo, ManifestInfo,
};
use crate::config::OciConfig;
use crate::error::{RegistryError, Result};
use crate::models::{Extension, ExtensionSource, ExtensionType, ExtensionVersion};
use crate::oci::models::{OciCatalog, OciManifest, OciTagsList};

/// OCI registry client
pub struct OciClient {
    backend_id: String,
    registry: String,
    namespace: String,
    auth_token: Option<String>,
    client: Client,
}

impl OciClient {
    /// Create new OCI client with a pre-resolved auth token.
    ///
    /// Token resolution (file read or vault:// fetch) is the caller's
    /// responsibility. Pass `None` for unauthenticated registry access.
    pub fn new(config: &OciConfig, auth_token: Option<String>) -> Result<Self> {
        let client = Client::builder()
            .timeout(Duration::from_secs(config.timeout_seconds))
            .danger_accept_invalid_certs(!config.verify_ssl)
            .build()
            .map_err(|e| RegistryError::Oci(format!("Failed to create HTTP client: {}", e)))?;

        let backend_id = config
            .id
            .clone()
            .unwrap_or_else(|| format!("oci-{}", config.registry));

        Ok(Self {
            backend_id,
            registry: config.registry.clone(),
            namespace: config.namespace.clone(),
            auth_token,
            client,
        })
    }

    /// List all extensions from OCI registry
    pub async fn list_extensions(
        &self,
        extension_type: Option<ExtensionType>,
    ) -> Result<Vec<Extension>> {
        debug!("Fetching artifacts from OCI registry: {}", self.registry);

        let catalog = self.list_catalog().await?;
        let mut extensions = Vec::new();

        for repo_name in catalog.repositories {
            // Skip if not in our namespace
            if !repo_name.starts_with(&format!("{}/", self.namespace)) {
                continue;
            }

            // Parse extension type and name from repository
            let Some((ext_type, _name)) = self.parse_repository_name(&repo_name) else {
                continue;
            };

            // Filter by type if specified
            if let Some(filter_type) = extension_type {
                if ext_type != filter_type {
                    continue;
                }
            }

            // Get tags for this repository
            let Ok(tags) = self.list_tags(&repo_name).await else {
                continue;
            };

            let Some(latest_tag) = tags.tags.first() else {
                continue;
            };

            if let Ok(manifest) = self.get_manifest(&repo_name, latest_tag).await {
                if let Some(extension) =
                    self.manifest_to_extension(&repo_name, latest_tag, &manifest, ext_type)
                {
                    extensions.push(extension);
                }
            }
        }

        Ok(extensions)
    }

    /// Get specific extension metadata
    pub async fn get_extension(
        &self,
        extension_type: ExtensionType,
        name: &str,
    ) -> Result<Extension> {
        let repo_name = self.format_repository_name(extension_type, name);
        debug!("Fetching extension: {}", repo_name);

        let tags = self.list_tags(&repo_name).await?;
        let latest_tag = tags
            .tags
            .first()
            .ok_or_else(|| RegistryError::NotFound(format!("No tags found for {}", repo_name)))?;

        let manifest = self.get_manifest(&repo_name, latest_tag).await?;

        self.manifest_to_extension(&repo_name, latest_tag, &manifest, extension_type)
            .ok_or_else(|| {
                RegistryError::NotFound(format!("Invalid extension metadata for {}", repo_name))
            })
    }

    /// List all versions for an extension
    pub async fn list_versions(
        &self,
        extension_type: ExtensionType,
        name: &str,
    ) -> Result<Vec<ExtensionVersion>> {
        let repo_name = self.format_repository_name(extension_type, name);
        debug!("Fetching versions for: {}", repo_name);

        let tags = self.list_tags(&repo_name).await?;

        let mut versions = Vec::new();
        for tag in tags.tags {
            if let Ok(manifest) = self.get_manifest(&repo_name, &tag).await {
                let size = manifest.layers.iter().map(|l| l.size).sum();

                versions.push(ExtensionVersion {
                    version: tag.clone(),
                    published_at: Utc::now(), /* OCI doesn't provide creation time without
                                               * additional metadata */
                    download_url: None,
                    checksum: Some(manifest.config.digest.clone()),
                    size: Some(size),
                });
            }
        }

        Ok(versions)
    }

    /// Pull extension artifact
    pub async fn download_extension(
        &self,
        extension_type: ExtensionType,
        name: &str,
        tag: &str,
    ) -> Result<Bytes> {
        let repo_name = self.format_repository_name(extension_type, name);
        debug!("Pulling extension: {} tag {}", repo_name, tag);

        let manifest = self.get_manifest(&repo_name, tag).await?;

        // Download first layer (assuming single-layer artifacts)
        let layer = manifest
            .layers
            .first()
            .ok_or_else(|| RegistryError::Oci("No layers found in manifest".to_string()))?;

        self.download_blob(&repo_name, &layer.digest).await
    }

    /// List catalog (all repositories)
    async fn list_catalog(&self) -> Result<OciCatalog> {
        let url = format!("https://{}/v2/_catalog", self.registry);

        let mut request = self.client.get(&url);

        if let Some(ref token) = self.auth_token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request
            .send()
            .await
            .map_err(|e| RegistryError::Oci(format!("Request failed: {}", e)))?;

        if !response.status().is_success() {
            return Err(RegistryError::Oci(format!(
                "Failed to list catalog: {}",
                response.status()
            )));
        }

        response
            .json()
            .await
            .map_err(|e| RegistryError::Oci(format!("Failed to parse response: {}", e)))
    }

    /// List tags for repository
    async fn list_tags(&self, repository: &str) -> Result<OciTagsList> {
        let url = format!("https://{}/v2/{}/tags/list", self.registry, repository);

        let mut request = self.client.get(&url);

        if let Some(ref token) = self.auth_token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request
            .send()
            .await
            .map_err(|e| RegistryError::Oci(format!("Request failed: {}", e)))?;

        if response.status().as_u16() == 404 {
            return Err(RegistryError::NotFound(format!(
                "Repository not found: {}",
                repository
            )));
        }

        if !response.status().is_success() {
            return Err(RegistryError::Oci(format!(
                "Failed to list tags: {}",
                response.status()
            )));
        }

        response
            .json()
            .await
            .map_err(|e| RegistryError::Oci(format!("Failed to parse response: {}", e)))
    }

    /// Get manifest for tag
    async fn get_manifest(&self, repository: &str, tag: &str) -> Result<OciManifest> {
        let url = format!(
            "https://{}/v2/{}/manifests/{}",
            self.registry, repository, tag
        );

        let mut request = self
            .client
            .get(&url)
            .header("Accept", "application/vnd.oci.image.manifest.v1+json");

        if let Some(ref token) = self.auth_token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request
            .send()
            .await
            .map_err(|e| RegistryError::Oci(format!("Request failed: {}", e)))?;

        if response.status().as_u16() == 404 {
            return Err(RegistryError::NotFound(format!(
                "Manifest not found: {}:{}",
                repository, tag
            )));
        }

        if !response.status().is_success() {
            return Err(RegistryError::Oci(format!(
                "Failed to get manifest: {}",
                response.status()
            )));
        }

        response
            .json()
            .await
            .map_err(|e| RegistryError::Oci(format!("Failed to parse manifest: {}", e)))
    }

    /// Download blob by digest
    async fn download_blob(&self, repository: &str, digest: &str) -> Result<Bytes> {
        let url = format!(
            "https://{}/v2/{}/blobs/{}",
            self.registry, repository, digest
        );

        let mut request = self.client.get(&url);

        if let Some(ref token) = self.auth_token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request
            .send()
            .await
            .map_err(|e| RegistryError::Oci(format!("Download failed: {}", e)))?;

        if !response.status().is_success() {
            return Err(RegistryError::Oci(format!(
                "Download failed: {}",
                response.status()
            )));
        }

        response
            .bytes()
            .await
            .map_err(|e| RegistryError::Oci(format!("Failed to read response: {}", e)))
    }

    /// Parse repository name into extension type and name
    fn parse_repository_name(&self, repo_name: &str) -> Option<(ExtensionType, String)> {
        let name = repo_name.strip_prefix(&format!("{}/", self.namespace))?;

        name.strip_suffix("-provider")
            .map(|base_name| (ExtensionType::Provider, base_name.to_string()))
            .or_else(|| {
                name.strip_suffix("-taskserv")
                    .map(|base_name| (ExtensionType::Taskserv, base_name.to_string()))
            })
            .or_else(|| {
                name.strip_suffix("-cluster")
                    .map(|base_name| (ExtensionType::Cluster, base_name.to_string()))
            })
    }

    /// Format repository name from extension type and name
    fn format_repository_name(&self, extension_type: ExtensionType, name: &str) -> String {
        let suffix = match extension_type {
            ExtensionType::Provider => "-provider",
            ExtensionType::Taskserv => "-taskserv",
            ExtensionType::Cluster => "-cluster",
        };
        format!("{}/{}{}", self.namespace, name, suffix)
    }

    /// Convert OCI manifest to Extension
    fn manifest_to_extension(
        &self,
        repo_name: &str,
        tag: &str,
        manifest: &OciManifest,
        extension_type: ExtensionType,
    ) -> Option<Extension> {
        let (_, name) = self.parse_repository_name(repo_name)?;

        let description = manifest
            .annotations
            .as_ref()
            .and_then(|a| a.get("org.opencontainers.image.description"))
            .cloned()
            .unwrap_or_default();

        let author = manifest
            .annotations
            .as_ref()
            .and_then(|a| a.get("org.opencontainers.image.authors"))
            .cloned();

        let repository = manifest
            .annotations
            .as_ref()
            .and_then(|a| a.get("org.opencontainers.image.url"))
            .cloned();

        let size = manifest.layers.iter().map(|l| l.size).sum();

        Some(Extension {
            name,
            extension_type,
            version: tag.to_string(),
            description,
            author,
            repository,
            source: ExtensionSource::Oci,
            published_at: Utc::now(),
            download_url: None,
            checksum: Some(manifest.config.digest.clone()),
            size: Some(size),
            tags: None,
        })
    }

    /// Check health of OCI connection
    pub async fn health_check(&self) -> Result<()> {
        let url = format!("https://{}/v2/", self.registry);

        let mut request = self.client.get(&url).timeout(Duration::from_secs(5));

        if let Some(ref token) = self.auth_token {
            request = request.header("Authorization", format!("Bearer {}", token));
        }

        let response = request
            .send()
            .await
            .map_err(|e| RegistryError::Oci(format!("Health check failed: {}", e)))?;

        if response.status().is_success() || response.status().as_u16() == 401 {
            // 401 means registry is up but auth is required
            Ok(())
        } else {
            Err(RegistryError::Oci(format!(
                "Health check returned: {}",
                response.status()
            )))
        }
    }
}

/// ExtensionClient trait implementation for OCI
#[async_trait]
impl ExtensionClient for OciClient {
    fn backend_id(&self) -> &str {
        &self.backend_id
    }

    fn backend_type(&self) -> BackendType {
        BackendType::Oci
    }

    async fn list_extensions(
        &self,
        extension_type: Option<ExtensionType>,
    ) -> Result<Vec<Extension>> {
        OciClient::list_extensions(self, extension_type).await
    }

    async fn get_extension(&self, extension_type: ExtensionType, name: &str) -> Result<Extension> {
        OciClient::get_extension(self, extension_type, name).await
    }

    async fn list_versions(
        &self,
        extension_type: ExtensionType,
        name: &str,
    ) -> Result<Vec<ExtensionVersion>> {
        OciClient::list_versions(self, extension_type, name).await
    }

    async fn download_extension(
        &self,
        extension_type: ExtensionType,
        name: &str,
        version: &str,
    ) -> Result<Bytes> {
        OciClient::download_extension(self, extension_type, name, version).await
    }

    async fn health_check(&self) -> Result<()> {
        OciClient::health_check(self).await
    }
}

/// DistributionClient trait implementation for OCI
#[async_trait]
impl DistributionClient for OciClient {
    async fn get_manifest(&self, repo_name: &str, tag: &str) -> Result<ManifestInfo> {
        let manifest = self.get_manifest(repo_name, tag).await?;
        let total_size: u64 =
            manifest.layers.iter().map(|l| l.size).sum::<u64>() + manifest.config.size;

        Ok(ManifestInfo {
            config_digest: manifest.config.digest,
            layers: manifest
                .layers
                .into_iter()
                .map(|l| LayerInfo {
                    digest: l.digest,
                    size: l.size,
                    media_type: l.media_type,
                })
                .collect(),
            total_size,
        })
    }

    async fn list_catalog(&self) -> Result<Vec<String>> {
        let catalog = self.list_catalog().await?;
        Ok(catalog.repositories)
    }

    async fn get_digest(&self, repo_name: &str, tag: &str) -> Result<String> {
        let manifest = self.get_manifest(repo_name, tag).await?;
        Ok(manifest.config.digest)
    }

    async fn verify_artifact(&self, repo_name: &str, tag: &str, digest: &str) -> Result<bool> {
        let manifest = self.get_manifest(repo_name, tag).await?;
        Ok(manifest.config.digest == digest)
    }
}