rialo-build-lib 0.2.0

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Rialo custom Rust toolchain management
//!
//! This module manages the Rialo custom Rust toolchain with the
//! riscv64emac-solana-solana target for building Rust programs.

use std::path::{Path, PathBuf};

use anyhow::{Context, Result};

use super::{
    extract_tar_gz, get_platform, get_s3_bucket, get_toolchain_root, S3StorageBackend, Toolchain,
    ToolchainConfig,
};

/// Rust nightly version to use for source builds
pub const RUST_NIGHTLY_VERSION: &str = "nightly-2025-05-10";

/// How the toolchain version was resolved
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolchainSource {
    /// From `RIALO_RUST_TOOLCHAIN_VERSION` environment variable
    Environment,
    /// From release manifest (rialoman-managed binary)
    Rialoman,
    /// First installed toolchain found on disk
    Installed,
}

impl std::fmt::Display for ToolchainSource {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Environment => write!(f, "environment"),
            Self::Rialoman => write!(f, "manifest"),
            Self::Installed => write!(f, "installed"),
        }
    }
}

/// Rust commit hash for source builds
pub const RUST_COMMIT_HASH: &str = "dcecb99176edf2eec51613730937d21cdd5c8f6e";

/// Embedded RISC-V backend patch
/// This patch adds the riscv64emac-solana-solana target to the Rust compiler
const RISCV_BACKEND_PATCH: &str = include_str!("../../patches/0001-risc-v-backend.patch");

/// Rialo Rust toolchain manager
///
/// Manages the Rialo custom Rust toolchain which includes:
/// - rustc with riscv64emac-solana-solana target
/// - cargo
/// - std library for the custom target
#[derive(Debug, Clone)]
pub struct RialoRustToolchain {
    pub(crate) config: ToolchainConfig,
    /// Whether this toolchain was built from source
    #[allow(dead_code)]
    built_from_source: bool,
}

impl RialoRustToolchain {
    /// Create a new Rialo Rust toolchain with auto-resolved version.
    ///
    /// Resolution order:
    /// 1. `RIALO_RUST_TOOLCHAIN_VERSION` environment variable
    /// 2. Release manifest `rust_toolchain` field (if running from rialoman-installed binary)
    /// 3. First installed toolchain
    pub fn new() -> Result<Self> {
        let (version, _) = Self::resolve_version()?;
        Self::with_version(&version)
    }

    /// Create a new Rialo Rust toolchain with an explicit version.
    pub fn with_version(version: &str) -> Result<Self> {
        // Backwards compatibility: redirect magic strings to new()
        if version == "auto" || version == "latest" {
            return Self::new();
        }

        let toolchain_root = get_toolchain_root()?;
        let install_path = toolchain_root.join(format!("rialo-rust-{}", version));
        let platform = get_platform()?;
        let download_url = Self::get_download_url(version, &platform)?;

        Ok(Self {
            config: ToolchainConfig {
                name: "rialo-rust".to_string(),
                version: version.to_string(),
                download_url,
                install_path,
                checksum: None,
            },
            built_from_source: false,
        })
    }

    /// Resolve the toolchain version automatically and its source.
    ///
    /// Resolution order:
    /// 1. `RIALO_RUST_TOOLCHAIN_VERSION` environment variable
    /// 2. Release manifest `rust_toolchain` field (rialoman-installed binary)
    /// 3. First installed toolchain
    pub fn resolve_version() -> Result<(String, ToolchainSource)> {
        if let Ok(version) = std::env::var("RIALO_RUST_TOOLCHAIN_VERSION") {
            log::info!("Using toolchain {} from environment variable", version);
            return Ok((version, ToolchainSource::Environment));
        }

        if let Some(version) = super::detect_rialoman_release_toolchain_version()? {
            log::info!("Using toolchain {} from Rialoman release manifest", version);
            return Ok((version, ToolchainSource::Rialoman));
        }

        if let Some(version) = Self::find_installed_version()? {
            log::info!("Using installed Rialo Rust toolchain: {}", version);
            return Ok((version, ToolchainSource::Installed));
        }

        Err(anyhow::anyhow!(
            "No Rialo Rust toolchains installed. Install one with:\n  \
             rialoman toolchain install rialo-rust --version <VERSION>\n\n\
             Or build from source with:\n  \
             rialoman toolchain build"
        ))
    }

    /// Find the first installed Rialo Rust toolchain version.
    ///
    /// Returns `None` if no toolchains are installed.
    pub fn find_installed_version() -> Result<Option<String>> {
        let toolchains = crate::toolchain::list_installed_toolchains()?;
        for (name, version) in toolchains {
            if name == "rialo-rust" {
                log::debug!("Found installed Rialo Rust toolchain: {}", version);
                return Ok(Some(version));
            }
        }
        Ok(None)
    }

    /// Create a new Rialo Rust toolchain for building from source
    pub fn new_from_source() -> Result<Self> {
        let toolchain_root = get_toolchain_root()?;
        let install_path = toolchain_root.join("rialo-rust-source");

        Ok(Self {
            config: ToolchainConfig {
                name: "rialo-rust".to_string(),
                version: RUST_NIGHTLY_VERSION.to_string(),
                download_url: String::new(), // Not used for source builds
                install_path,
                checksum: None,
            },
            built_from_source: true,
        })
    }

    /// Get the download URL for prebuilt binaries
    fn get_download_url(version: &str, platform: &str) -> Result<String> {
        // Using Rialo prebuilt toolchain releases
        // These are hosted on GitHub releases at subzerolabs/rialo-toolchains
        let base_url = "https://github.com/subzerolabs/rialo-toolchains/releases/download";

        // Map our platform strings to the toolchain naming convention
        let toolchain_platform = match platform {
            "x86_64-apple-darwin" => "x86_64-apple-darwin",
            "aarch64-apple-darwin" => "aarch64-apple-darwin",
            "x86_64-unknown-linux-gnu" => "x86_64-unknown-linux-gnu",
            "aarch64-unknown-linux-gnu" => "aarch64-unknown-linux-gnu",
            _ => {
                return Err(anyhow::anyhow!(
                    "No Rialo Rust toolchain available for platform: {platform}"
                ))
            }
        };

        // Format: rialo-rust-{version}-{platform}.tar.gz
        let archive_name = format!("rialo-rust-{version}-{toolchain_platform}.tar.gz");
        let url = format!("{base_url}/{version}/{archive_name}");

        Ok(url)
    }

    /// Get the embedded RISC-V backend patch
    pub fn get_patch_content() -> &'static str {
        RISCV_BACKEND_PATCH
    }

    /// Write the patch file to a directory
    pub fn write_patch_file(dest_dir: &Path) -> Result<PathBuf> {
        std::fs::create_dir_all(dest_dir)
            .with_context(|| format!("Failed to create directory {}", dest_dir.display()))?;

        let patch_path = dest_dir.join("0001-risc-v-backend.patch");
        std::fs::write(&patch_path, RISCV_BACKEND_PATCH)
            .with_context(|| format!("Failed to write patch file to {}", patch_path.display()))?;

        Ok(patch_path)
    }

    /// Check if rustup is installed
    fn check_rustup() -> Result<()> {
        which::which("rustup")
            .context("rustup not found. Please install rustup from https://rustup.rs/")?;
        Ok(())
    }

    /// Check if the toolchain is currently registered with rustup
    fn check_rustup_registration() -> Result<bool> {
        let output = std::process::Command::new("rustup")
            .args(["toolchain", "list"])
            .output()
            .context("Failed to run rustup toolchain list")?;

        let output_str = String::from_utf8_lossy(&output.stdout);
        Ok(output_str.lines().any(|line| line.starts_with("rialo")))
    }

    /// Register the toolchain with rustup as 'rialo', possibly overwriting the
    /// old alias, if exists.
    ///
    /// # Errors
    ///
    /// Returns an error if rustup is not installed or the link operation fails.
    pub fn register_with_rustup(&self) -> Result<()> {
        log::info!("Registering toolchain with rustup as 'rialo'...");

        // `rustup toolchain link` is idempotent - it overwrites existing links
        let status = std::process::Command::new("rustup")
            .args([
                "toolchain",
                "link",
                "rialo",
                self.config.install_path.to_str().unwrap(),
            ])
            .status()
            .context("Failed to link toolchain with rustup")?;

        if !status.success() {
            return Err(anyhow::anyhow!("Failed to register toolchain with rustup"));
        }

        log::info!("Toolchain registered as 'rialo'");
        Ok(())
    }

    /// Unregister the toolchain from rustup
    fn unregister_from_rustup() -> Result<()> {
        if Self::check_rustup_registration()? {
            log::info!("Unregistering toolchain from rustup...");
            std::process::Command::new("rustup")
                .args(["toolchain", "remove", "rialo"])
                .status()
                .context("Failed to unregister toolchain")?;
            log::info!("Toolchain unregistered");
        }
        Ok(())
    }

    /// Validate that a toolchain exists and is complete at the given path
    ///
    /// Checks for the presence of rustc and cargo binaries in the bin/ directory
    fn validate_toolchain_at(path: &Path) -> Result<()> {
        if !path.exists() {
            return Err(anyhow::anyhow!(
                "Toolchain path does not exist: {}",
                path.display()
            ));
        }

        let bin_path = path.join("bin");
        if !bin_path.exists() {
            return Err(anyhow::anyhow!(
                "Toolchain bin directory not found at {}",
                bin_path.display()
            ));
        }

        let rustc = bin_path.join("rustc");
        let cargo = bin_path.join("cargo");

        if !rustc.exists() {
            return Err(anyhow::anyhow!(
                "rustc not found in toolchain at {}",
                bin_path.display()
            ));
        }

        if !cargo.exists() {
            return Err(anyhow::anyhow!(
                "cargo not found in toolchain at {}",
                bin_path.display()
            ));
        }

        Ok(())
    }
}

impl Default for RialoRustToolchain {
    fn default() -> Self {
        Self::new().expect("Failed to create default Rialo Rust toolchain")
    }
}

impl Toolchain for RialoRustToolchain {
    fn is_installed(&self) -> Result<bool> {
        // Check if toolchain files exist at the install path
        if Self::validate_toolchain_at(&self.config.install_path).is_ok() {
            log::debug!(
                "Found Rialo Rust toolchain at {}",
                self.config.install_path.display()
            );

            // Require rustup registration for the toolchain to be considered installed
            if !Self::check_rustup_registration()? {
                log::debug!("Toolchain files exist but not registered with rustup");
                return Ok(false);
            }

            return Ok(true);
        }

        Ok(false)
    }

    fn install(&self) -> Result<()> {
        if self.is_installed()? {
            log::info!(
                "Rialo Rust toolchain {} is already installed at {}",
                self.config.version,
                self.config.install_path.display()
            );
            return Ok(());
        }

        // Check rustup is available
        Self::check_rustup()?;

        log::info!("Installing Rialo Rust toolchain {}...", self.config.version);

        // Create toolchain root directory
        let toolchain_root = get_toolchain_root()?;
        std::fs::create_dir_all(&toolchain_root).with_context(|| {
            format!(
                "Failed to create toolchain directory {}",
                toolchain_root.display()
            )
        })?;

        // Create a temporary directory for download
        let temp_dir = tempfile::tempdir()
            .context("Failed to create temporary directory for toolchain download")?;

        let archive_path = temp_dir.path().join("rialo-rust-toolchain.tar.gz");

        self.download_with_fallback(&archive_path)
            .with_context(|| {
                format!(
                    "Failed to download Rialo Rust toolchain version {}",
                    self.config.version
                )
            })?;

        // Create install directory
        std::fs::create_dir_all(&self.config.install_path).with_context(|| {
            format!(
                "Failed to create install directory {}",
                self.config.install_path.display()
            )
        })?;

        // Extract the archive
        extract_tar_gz(&archive_path, &self.config.install_path)?;

        // Register with rustup
        self.register_with_rustup()?;

        log::info!(
            "Rialo Rust toolchain {} installed successfully at {}",
            self.config.version,
            self.config.install_path.display()
        );

        Ok(())
    }

    fn validate(&self) -> Result<()> {
        if !self.is_installed()? {
            return Err(anyhow::anyhow!(
                "Rialo Rust toolchain {} is not installed. Run 'rialo-build toolchain install rialo-rust' to install it.",
                self.config.version
            ));
        }

        // Validate that we can run cargo +rialo --version
        let output = std::process::Command::new("cargo")
            .args(["+rialo", "--version"])
            .output()
            .context("Failed to run cargo +rialo --version. Is rustup configured correctly?")?;

        if !output.status.success() {
            return Err(anyhow::anyhow!(
                "cargo +rialo returned non-zero exit code. The toolchain may be corrupted."
            ));
        }

        log::info!("Rialo Rust toolchain validation:");
        let version_output = String::from_utf8_lossy(&output.stdout);
        log::info!("  {}", version_output.trim());

        // Check that the custom target is available
        let output = std::process::Command::new("rustc")
            .args(["+rialo", "--print", "target-list"])
            .output()
            .context("Failed to run rustc +rialo --print target-list")?;

        let targets = String::from_utf8_lossy(&output.stdout);
        if !targets.contains("riscv64emac-solana-solana") {
            return Err(anyhow::anyhow!(
                "Custom target riscv64emac-solana-solana not found in toolchain. The toolchain may be incorrectly built."
            ));
        }

        log::info!("  Custom target: riscv64emac-solana-solana");
        log::info!("Toolchain is functional");

        Ok(())
    }

    fn get_bin_path(&self) -> Result<PathBuf> {
        let bin_path = self.config.install_path.join("bin");
        if !bin_path.exists() {
            return Err(anyhow::anyhow!(
                "Toolchain bin directory not found at {}",
                bin_path.display()
            ));
        }
        Ok(bin_path)
    }

    fn get_config(&self) -> &ToolchainConfig {
        &self.config
    }
}

impl RialoRustToolchain {
    /// Uninstall the toolchain, including rustup registration
    pub fn uninstall(&self) -> Result<()> {
        // Unregister from rustup first
        Self::unregister_from_rustup()?;

        // Remove installation directory
        if self.config.install_path.exists() {
            std::fs::remove_dir_all(&self.config.install_path).with_context(|| {
                format!(
                    "Failed to remove installation directory {}",
                    self.config.install_path.display()
                )
            })?;
            log::info!(
                "Removed Rialo Rust toolchain from {}",
                self.config.install_path.display()
            );
        }

        Ok(())
    }

    fn download_with_fallback(&self, dest: &Path) -> Result<()> {
        self.download_from_http(dest)
    }

    /// Download toolchain from HTTP (public S3 bucket via HTTPS)
    ///
    /// Downloads from rialo-artifacts S3 bucket via HTTP - no AWS credentials needed.
    fn download_from_http(&self, dest: &Path) -> Result<()> {
        use crate::toolchain::HttpToolchainClient;

        let client =
            HttpToolchainClient::new().context("Failed to create HTTP toolchain client")?;

        let platform = get_platform()?;
        let archive_name = format!("rialo-rust-{}-{}", self.config.version, platform);

        client
            .download_toolchain("rialo-rust", &self.config.version, &archive_name, dest)
            .with_context(|| {
                format!(
                    "Failed to download Rialo Rust toolchain version {}",
                    self.config.version
                )
            })
    }

    /// Upload this toolchain to S3
    ///
    /// Packages the installed toolchain as a tarball and uploads it to S3.
    /// Requires AWS credentials to be configured.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The toolchain is not installed
    /// - AWS credentials are not available
    /// - The upload fails
    pub fn upload_to_s3(&self) -> Result<()> {
        if !self.is_installed()? {
            return Err(anyhow::anyhow!(
                "Toolchain is not installed. Cannot upload to S3."
            ));
        }

        log::debug!("Packaging Rialo Rust toolchain for upload...");

        // Create tarball
        let temp_dir =
            tempfile::tempdir().context("Failed to create temporary directory for tarball")?;
        let archive_path = temp_dir.path().join("rialo-rust-toolchain.tar.gz");

        self.create_tarball(&archive_path)?;

        log::info!("Uploading to S3...");

        // Upload to S3
        let runtime = tokio::runtime::Runtime::new()
            .context("Failed to create tokio runtime for S3 upload")?;

        runtime.block_on(async {
            let bucket = get_s3_bucket();
            let backend = S3StorageBackend::new(bucket).await?.context(
                "S3 backend not available. Ensure AWS credentials are configured \
                     (AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY).",
            )?;

            let platform = get_platform()?;
            backend
                .upload_toolchain("rialo-rust", &self.config.version, &platform, &archive_path)
                .await
        })
    }

    /// Create a tarball of the installed toolchain
    fn create_tarball(&self, dest: &Path) -> Result<()> {
        use std::fs::File;

        use flate2::{write::GzEncoder, Compression};
        use tar::Builder;

        log::debug!("Creating tarball at {}...", dest.display());

        let tar_gz = File::create(dest)
            .with_context(|| format!("Failed to create tarball file {}", dest.display()))?;

        let enc = GzEncoder::new(tar_gz, Compression::default());
        let mut tar = Builder::new(enc);

        // Add all files from the installation directory
        tar.append_dir_all(".", &self.config.install_path)
            .with_context(|| {
                format!(
                    "Failed to add files from {} to tarball",
                    self.config.install_path.display()
                )
            })?;

        tar.finish().context("Failed to finalize tarball")?;

        log::debug!("Tarball created");
        Ok(())
    }
}

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

    #[test]
    fn test_create_rialo_rust_toolchain() {
        // Test creating a toolchain struct with a specific version
        // This should work regardless of whether toolchains are actually installed
        let toolchain = RialoRustToolchain::with_version("v1.0.0");
        assert!(
            toolchain.is_ok(),
            "Failed to create toolchain: {:?}",
            toolchain.err()
        );

        let tc = toolchain.unwrap();
        assert_eq!(tc.config.version, "v1.0.0");
        assert_eq!(tc.config.name, "rialo-rust");
    }

    #[test]
    fn test_toolchain_with_different_version() {
        // Test with a different version to ensure version parameter works correctly
        let toolchain = RialoRustToolchain::with_version("v2.5.3");
        assert!(toolchain.is_ok());
        let tc = toolchain.unwrap();
        assert_eq!(tc.config.version, "v2.5.3");
        assert_eq!(tc.config.name, "rialo-rust");
    }

    #[test]
    fn test_patch_content_exists() {
        let patch = RialoRustToolchain::get_patch_content();
        assert!(!patch.is_empty());
        assert!(patch.contains("riscv64emac-solana-solana"));
        assert!(patch.contains("risc-v backend"));
    }

    #[test]
    fn test_source_build_constants() {
        assert_eq!(RUST_NIGHTLY_VERSION, "nightly-2025-05-10");
        assert_eq!(RUST_COMMIT_HASH, "dcecb99176edf2eec51613730937d21cdd5c8f6e");
    }
}