mecha10_cli/services/

remote_image.rs

//! Remote image detection service
//!
//! Handles detection of whether to use pre-built images from the registry
//! or build locally based on project configuration.

use crate::types::project::TargetsConfig;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;

/// URL for the remote image manifest hosted in user-tools
const MANIFEST_URL: &str = "https://raw.githubusercontent.com/laboratory-one/user-tools/main/mecha10-remote/manifest.json";

/// Manifest describing available pre-built remote images
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RemoteImageManifest {
    /// Version of the manifest schema
    pub version: String,
    /// Available pre-built configurations
    pub configurations: Vec<RemoteImageConfig>,
    /// Hash of the template Dockerfile.remote
    pub dockerfile_template_hash: String,
}

/// A pre-built image configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RemoteImageConfig {
    /// Hash of the sorted node array (for matching)
    pub nodes_hash: String,
    /// List of nodes included in this image
    pub nodes: Vec<String>,
    /// Docker image tag to pull
    pub image_tag: String,
    /// Image digest for verification
    pub image_digest: Option<String>,
}

/// Result of pre-built image detection
#[derive(Debug)]
pub struct PrebuiltDetectionResult {
    /// Whether a pre-built image can be used
    pub can_use_prebuilt: bool,
    /// The image tag to use (if pre-built available)
    pub image_tag: Option<String>,
    /// Reason why pre-built cannot be used
    pub reason: Option<String>,
}

/// Service for detecting and managing remote images
pub struct RemoteImageService;

impl RemoteImageService {
    /// Create a new remote image service
    pub fn new() -> Self {
        Self
    }

    /// Detect whether a pre-built image can be used
    ///
    /// Returns information about whether to use pre-built or build locally.
    pub async fn detect_prebuilt(
        &self,
        targets: &TargetsConfig,
        project_dir: &Path,
    ) -> Result<PrebuiltDetectionResult> {
        // Check for custom nodes first (quick check, no network needed)
        if targets.has_custom_remote_nodes() {
            return Ok(PrebuiltDetectionResult {
                can_use_prebuilt: false,
                image_tag: None,
                reason: Some("Project contains custom @local/* nodes".to_string()),
            });
        }

        // Check if there are any remote nodes
        if !targets.has_remote_nodes() {
            return Ok(PrebuiltDetectionResult {
                can_use_prebuilt: false,
                image_tag: None,
                reason: Some("No remote nodes configured".to_string()),
            });
        }

        // Check if Dockerfile.remote has been modified
        let dockerfile_path = project_dir.join("docker/Dockerfile.remote");
        if dockerfile_path.exists() {
            // Try to fetch manifest and compare dockerfile hash
            match self.fetch_manifest().await {
                Ok(manifest) => {
                    let local_hash = self.hash_file(&dockerfile_path).await?;
                    if local_hash != manifest.dockerfile_template_hash {
                        return Ok(PrebuiltDetectionResult {
                            can_use_prebuilt: false,
                            image_tag: None,
                            reason: Some("Dockerfile.remote has been modified".to_string()),
                        });
                    }

                    // Check if we have a matching configuration
                    let nodes_hash = self.hash_nodes(targets);
                    for config in &manifest.configurations {
                        if config.nodes_hash == nodes_hash {
                            return Ok(PrebuiltDetectionResult {
                                can_use_prebuilt: true,
                                image_tag: Some(config.image_tag.clone()),
                                reason: None,
                            });
                        }
                    }

                    // No matching configuration found
                    Ok(PrebuiltDetectionResult {
                        can_use_prebuilt: false,
                        image_tag: None,
                        reason: Some("No pre-built image for this node combination".to_string()),
                    })
                }
                Err(e) => {
                    // Network error or manifest not available
                    Ok(PrebuiltDetectionResult {
                        can_use_prebuilt: false,
                        image_tag: None,
                        reason: Some(format!("Could not fetch manifest: {}", e)),
                    })
                }
            }
        } else {
            // No Dockerfile.remote, can't use prebuilt
            Ok(PrebuiltDetectionResult {
                can_use_prebuilt: false,
                image_tag: None,
                reason: Some("Dockerfile.remote not found".to_string()),
            })
        }
    }

    /// Fetch the remote image manifest from user-tools
    pub async fn fetch_manifest(&self) -> Result<RemoteImageManifest> {
        let response = reqwest::get(MANIFEST_URL)
            .await
            .context("Failed to fetch remote image manifest")?;

        if !response.status().is_success() {
            anyhow::bail!(
                "Failed to fetch manifest: HTTP {}",
                response.status()
            );
        }

        let manifest: RemoteImageManifest = response
            .json()
            .await
            .context("Failed to parse remote image manifest")?;

        Ok(manifest)
    }

    /// Compute a hash of the remote nodes for matching against manifest
    pub fn hash_nodes(&self, targets: &TargetsConfig) -> String {
        // Get sorted nodes for deterministic hashing
        let sorted_nodes = targets.sorted_remote_nodes();
        let nodes_str = sorted_nodes.join(",");

        let hash = blake3::hash(nodes_str.as_bytes());
        // Return first 16 hex chars for brevity
        hash.to_hex()[..16].to_string()
    }

    /// Hash a file's contents
    async fn hash_file(&self, path: &Path) -> Result<String> {
        let content = tokio::fs::read(path)
            .await
            .context("Failed to read file for hashing")?;

        let hash = blake3::hash(&content);
        Ok(hash.to_hex().to_string())
    }

    /// Get the image environment variable name
    pub fn image_env_var() -> &'static str {
        "MECHA10_REMOTE_IMAGE"
    }
}

impl Default for RemoteImageService {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_hash_nodes_deterministic() {
        let service = RemoteImageService::new();

        let mut targets1 = TargetsConfig::default();
        targets1.remote = vec![
            "@mecha10/object-detector".to_string(),
            "@mecha10/image-classifier".to_string(),
        ];

        let mut targets2 = TargetsConfig::default();
        targets2.remote = vec![
            "@mecha10/image-classifier".to_string(),
            "@mecha10/object-detector".to_string(),
        ];

        // Should produce same hash regardless of order
        let hash1 = service.hash_nodes(&targets1);
        let hash2 = service.hash_nodes(&targets2);
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn test_hash_nodes_different_for_different_configs() {
        let service = RemoteImageService::new();

        let mut targets1 = TargetsConfig::default();
        targets1.remote = vec!["@mecha10/object-detector".to_string()];

        let mut targets2 = TargetsConfig::default();
        targets2.remote = vec!["@mecha10/image-classifier".to_string()];

        let hash1 = service.hash_nodes(&targets1);
        let hash2 = service.hash_nodes(&targets2);
        assert_ne!(hash1, hash2);
    }
}