rialo-build-lib 0.11.2

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

//! GNU RISC-V toolchain management
//!
//! This module manages the GNU RISC-V toolchain (gcc, binutils, etc.)
//! used for compiling C/C++ programs for RISC-V targets.

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,
};

/// Default GNU RISC-V toolchain version
pub const DEFAULT_GNU_RISCV_VERSION: &str = "13.2.0";

/// GNU RISC-V toolchain manager
///
/// Manages the GNU RISC-V cross-compilation toolchain which includes:
/// - riscv64-unknown-elf-gcc
/// - riscv64-unknown-elf-ld
/// - riscv64-unknown-elf-objcopy
/// - And other binutils
#[derive(Debug, Clone)]
pub struct GnuRiscvToolchain {
    config: ToolchainConfig,
}

impl GnuRiscvToolchain {
    /// Create a new GNU RISC-V toolchain with the default version
    pub fn new() -> Result<Self> {
        Self::with_version(DEFAULT_GNU_RISCV_VERSION)
    }

    /// Create a new GNU RISC-V toolchain with a specific version
    pub fn with_version(version: &str) -> Result<Self> {
        let toolchain_root = get_toolchain_root()?;
        let install_path = toolchain_root.join(format!("gnu-riscv-{version}"));

        let platform = get_platform()?;
        let download_url = Self::get_download_url(version, &platform)?;

        Ok(Self {
            config: ToolchainConfig {
                name: "gnu-riscv".to_string(),
                version: version.to_string(),
                download_url,
                install_path,
                checksum: None, // Checksums can be added per version if needed
            },
        })
    }

    /// Get the download URL for the GNU RISC-V toolchain based on platform
    fn get_download_url(version: &str, platform: &str) -> Result<String> {
        // Using RISC-V GNU Toolchain prebuilt releases
        // These are hosted on GitHub releases
        let base_url = "https://github.com/riscv-collab/riscv-gnu-toolchain/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-linux-ubuntu14",
            "aarch64-unknown-linux-gnu" => "aarch64-linux-ubuntu20",
            _ => {
                return Err(anyhow::anyhow!(
                    "No GNU RISC-V toolchain available for platform: {platform}"
                ))
            }
        };

        // Format: riscv64-elf-ubuntu-20.04-gcc-nightly-2024.09.03-nightly.tar.gz
        // We'll use a stable release format based on version
        let archive_name = format!("riscv64-elf-{toolchain_platform}-{version}.tar.gz");
        let url = format!("{base_url}/{version}/{archive_name}");

        Ok(url)
    }

    /// Get the path to a specific tool in the toolchain
    fn get_tool_path(&self, tool: &str) -> PathBuf {
        self.config
            .install_path
            .join("bin")
            .join(format!("riscv64-unknown-elf-{tool}"))
    }

    /// Check if a specific tool exists and is executable
    fn tool_exists(&self, tool: &str) -> bool {
        let tool_path = self.get_tool_path(tool);
        tool_path.exists() && tool_path.is_file()
    }
}

impl Default for GnuRiscvToolchain {
    fn default() -> Self {
        Self::new().expect("Failed to create default GNU RISC-V toolchain")
    }
}

impl Toolchain for GnuRiscvToolchain {
    fn is_installed(&self) -> Result<bool> {
        // Check if the install path exists and contains the expected tools
        if !self.config.install_path.exists() {
            return Ok(false);
        }

        // Check for essential tools
        let required_tools = ["gcc", "ld", "objcopy", "objdump", "ar", "as"];
        for tool in &required_tools {
            if !self.tool_exists(tool) {
                return Ok(false);
            }
        }

        Ok(true)
    }

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

        log::info!("Installing GNU RISC-V 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("riscv-toolchain.tar.gz");

        self.download_with_fallback(&archive_path)
            .with_context(|| {
                format!(
                    "Failed to download GNU RISC-V toolchain version {}\n\
                    \n\
                    Pre-built binaries may not be available for your platform.\n\
                    \n\
                    If you're building Rust programs:\n\
                      You don't need this toolchain! The Rialo Rust toolchain is already available.\n\
                      Use auto-detection: rialo-build --program-path /path/to/program\n\
                    \n\
                    If you're building C programs:\n\
                      The GNU RISC-V toolchain must be obtained through alternative channels\n\
                      or built from source (not currently supported via rialo-build).",
                    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)?;

        // The archive may contain a nested directory, so we need to move contents up
        Self::flatten_installation(&self.config.install_path)?;

        log::info!(
            "GNU RISC-V 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!(
                "GNU RISC-V toolchain {} is not installed.\n\
                \n\
                For Rust programs:\n\
                  You don't need the GNU toolchain! Use auto-detection instead:\n\
                  rialo-build --program-path /path/to/program\n\
                \n\
                For C programs:\n\
                  Install the GNU toolchain (note: pre-built binaries may not be available):\n\
                  rialo-build toolchain install gnu-riscv",
                self.config.version
            ));
        }

        // Validate that we can run gcc --version
        let gcc_path = self.get_tool_path("gcc");
        let output = std::process::Command::new(&gcc_path)
            .arg("--version")
            .output()
            .with_context(|| {
                format!(
                    "Failed to run gcc at {}. The toolchain may be corrupted.",
                    gcc_path.display()
                )
            })?;

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

        let version_output = String::from_utf8_lossy(&output.stdout);
        log::info!("GNU RISC-V toolchain validation:");
        log::info!("{}", version_output.lines().next().unwrap_or("Unknown"));
        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 GnuRiscvToolchain {
    /// Flatten the installation directory if the archive contains a nested directory
    fn flatten_installation(install_path: &Path) -> Result<()> {
        // Check if there's a single subdirectory
        let entries: Vec<_> = std::fs::read_dir(install_path)
            .with_context(|| format!("Failed to read directory {}", install_path.display()))?
            .filter_map(|e| e.ok())
            .collect();

        // If there's exactly one entry and it's a directory, move its contents up
        if entries.len() == 1 && entries[0].path().is_dir() {
            let nested_dir = entries[0].path();

            // Move all contents from nested_dir to install_path
            for entry in std::fs::read_dir(&nested_dir)
                .with_context(|| format!("Failed to read directory {}", nested_dir.display()))?
            {
                let entry = entry?;
                let src = entry.path();
                let dest = install_path.join(entry.file_name());

                std::fs::rename(&src, &dest).with_context(|| {
                    format!("Failed to move {} to {}", src.display(), dest.display())
                })?;
            }

            // Remove the now-empty nested directory
            std::fs::remove_dir(&nested_dir).with_context(|| {
                format!("Failed to remove empty directory {}", nested_dir.display())
            })?;
        }

        Ok(())
    }

    /// Get the prefix for RISC-V tools (e.g., "riscv64-unknown-elf-")
    pub fn get_tool_prefix(&self) -> &str {
        "riscv64-unknown-elf"
    }

    /// Get environment variables needed for building with this toolchain
    pub fn get_env_vars(&self) -> Result<Vec<(String, String)>> {
        let bin_path = self.get_bin_path()?;

        Ok(vec![
            (
                "RISCV_TOOLCHAIN_PATH".to_string(),
                self.config.install_path.display().to_string(),
            ),
            ("RISCV_BIN_PATH".to_string(), bin_path.display().to_string()),
            (
                "RISCV_PREFIX".to_string(),
                self.get_tool_prefix().to_string(),
            ),
        ])
    }

    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!("riscv64-elf-{}-{}", platform, self.config.version);

        client
            .download_toolchain("gnu-riscv", &self.config.version, &archive_name, dest)
            .with_context(|| {
                format!(
                    "Failed to download GNU RISC-V 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::info!("Packaging GNU RISC-V 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("gnu-riscv-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("gnu-riscv", &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_gnu_riscv_toolchain() {
        let toolchain = GnuRiscvToolchain::new();
        assert!(toolchain.is_ok());
    }

    #[test]
    fn test_toolchain_with_version() {
        let toolchain = GnuRiscvToolchain::with_version("13.2.0");
        assert!(toolchain.is_ok());
        let tc = toolchain.unwrap();
        assert_eq!(tc.config.version, "13.2.0");
        assert_eq!(tc.config.name, "gnu-riscv");
    }

    #[test]
    fn test_get_tool_prefix() {
        let toolchain = GnuRiscvToolchain::new().unwrap();
        assert_eq!(toolchain.get_tool_prefix(), "riscv64-unknown-elf");
    }
}