entrenar/sovereign/nix/

config.rs

//! Nix flake configuration

use super::crate_spec::CrateSpec;
use super::system::NixSystem;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Nix flake configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NixFlakeConfig {
    /// Crate specifications
    pub crates: Vec<CrateSpec>,
    /// Rust toolchain version
    pub rust_version: String,
    /// Features to enable per crate
    pub features: HashMap<String, Vec<String>>,
    /// Target systems
    pub systems: Vec<NixSystem>,
    /// Flake description
    pub description: String,
    /// Enable GPU support
    pub gpu_support: bool,
    /// Include dev shell
    pub include_dev_shell: bool,
    /// Include CI checks
    pub include_checks: bool,
}

impl NixFlakeConfig {
    /// Create a new flake config
    pub fn new(description: impl Into<String>) -> Self {
        Self {
            crates: Vec::new(),
            rust_version: "1.75.0".to_string(),
            features: HashMap::new(),
            systems: NixSystem::all(),
            description: description.into(),
            gpu_support: false,
            include_dev_shell: true,
            include_checks: true,
        }
    }

    /// Create the sovereign stack configuration (all PAIML crates)
    pub fn sovereign_stack() -> Self {
        let mut config = Self::new("PAIML Sovereign ML Stack - Air-gapped deployment ready");

        // Add all PAIML crates
        config.crates = vec![
            CrateSpec::crates_io("trueno", "0.2"),
            CrateSpec::crates_io("aprender", "0.1"),
            CrateSpec::crates_io("renacer", "0.1"),
            CrateSpec::crates_io("entrenar", "0.2"),
            CrateSpec::crates_io("realizar", "0.1"),
        ];

        // Set features
        config.features.insert("trueno".to_string(), vec!["simd".to_string()]);
        config.features.insert("entrenar".to_string(), vec!["full".to_string()]);

        config.rust_version = "1.75.0".to_string();
        config.include_dev_shell = true;
        config.include_checks = true;

        config
    }

    /// Add a crate to the configuration
    pub fn add_crate(mut self, spec: CrateSpec) -> Self {
        self.crates.push(spec);
        self
    }

    /// Set features for a crate
    pub fn with_features(
        mut self,
        crate_name: impl Into<String>,
        features: impl IntoIterator<Item = impl Into<String>>,
    ) -> Self {
        self.features.insert(crate_name.into(), features.into_iter().map(Into::into).collect());
        self
    }

    /// Set Rust version
    pub fn with_rust_version(mut self, version: impl Into<String>) -> Self {
        self.rust_version = version.into();
        self
    }

    /// Set target systems
    pub fn with_systems(mut self, systems: Vec<NixSystem>) -> Self {
        self.systems = systems;
        self
    }

    /// Enable GPU support
    pub fn with_gpu_support(mut self, enabled: bool) -> Self {
        self.gpu_support = enabled;
        self
    }

    /// Generate the flake.nix content
    pub fn generate_flake_nix(&self) -> String {
        let systems_list: Vec<&str> = self.systems.iter().map(NixSystem::as_str).collect();
        let systems_str =
            systems_list.iter().map(|s| format!("\"{s}\"")).collect::<Vec<_>>().join(" ");

        let crate_names: Vec<&str> = self.crates.iter().map(|c| c.name.as_str()).collect();

        let mut flake = String::new();

        // Header
        flake.push_str(&format!(
            r#"# Nix Flake for PAIML Sovereign Stack
# {}
# Generated by entrenar sovereign deployment tooling
{{
  description = "{}";

  inputs = {{
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    rust-overlay = {{
      url = "github:oxalica/rust-overlay";
      inputs.nixpkgs.follows = "nixpkgs";
    }};
    crane = {{
      url = "github:ipetkov/crane";
      inputs.nixpkgs.follows = "nixpkgs";
    }};
    flake-utils.url = "github:numtide/flake-utils";
  }};

  outputs = {{ self, nixpkgs, rust-overlay, crane, flake-utils, ... }}:
    flake-utils.lib.eachSystem [ {} ] (system:
      let
        overlays = [ (import rust-overlay) ];
        pkgs = import nixpkgs {{
          inherit system overlays;
        }};

        rustToolchain = pkgs.rust-bin.stable."{}" .default.override {{
          extensions = [ "rust-src" "rust-analyzer" ];
        }};

        craneLib = (crane.mkLib pkgs).overrideToolchain rustToolchain;

"#,
            chrono::Utc::now().format("%Y-%m-%d"),
            self.description,
            systems_str,
            self.rust_version,
        ));

        // Build inputs
        flake.push_str("        buildInputs = with pkgs; [\n");
        flake.push_str("          openssl\n");
        flake.push_str("          pkg-config\n");
        if self.gpu_support {
            flake.push_str("          # GPU support\n");
            flake.push_str("          cudatoolkit\n");
            flake.push_str("          cudnn\n");
        }
        flake.push_str("        ];\n\n");

        // Common args
        flake.push_str(&format!(
            r"        commonArgs = {{
          src = craneLib.cleanCargoSource ./.;
          inherit buildInputs;
          nativeBuildInputs = with pkgs; [ pkg-config ];
        }};

        # Build dependencies first (for caching)
        cargoArtifacts = craneLib.buildDepsOnly commonArgs;

        # Main packages
{}",
            self.generate_package_definitions(&crate_names),
        ));

        // Outputs
        flake.push_str(&format!(
            r"
      in {{
        packages = {{
{}          default = {};
        }};
",
            self.generate_packages_attr(&crate_names),
            crate_names.first().unwrap_or(&"entrenar"),
        ));

        // Dev shell
        if self.include_dev_shell {
            flake.push_str(&format!(
                r"
        devShells.default = pkgs.mkShell {{
          inputsFrom = [ {} ];
          buildInputs = with pkgs; [
            rustToolchain
            rust-analyzer
            cargo-watch
            cargo-edit
            cargo-expand
          ];
        }};
",
                crate_names.first().unwrap_or(&"entrenar"),
            ));
        }

        // Checks
        if self.include_checks {
            flake.push_str(
                r#"
        checks = {
          clippy = craneLib.cargoClippy (commonArgs // {
            inherit cargoArtifacts;
            cargoClippyExtraArgs = "--all-targets -- -D warnings";
          });

          test = craneLib.cargoNextest (commonArgs // {
            inherit cargoArtifacts;
            partitions = 1;
            partitionType = "count";
          });

          fmt = craneLib.cargoFmt {
            src = craneLib.cleanCargoSource ./.;
          };
        };
"#,
            );
        }

        // Close outputs
        flake.push_str("      });\n}\n");

        flake
    }

    /// Generate package definitions
    fn generate_package_definitions(&self, crate_names: &[&str]) -> String {
        use std::fmt::Write;
        let mut result = String::new();
        for name in crate_names {
            let features = self
                .features
                .get(*name)
                .map(|f| {
                    let feature_list = f.join(",");
                    format!(r#"cargoExtraArgs = "--features {feature_list}";"#)
                })
                .unwrap_or_default();

            let _ = writeln!(
                &mut result,
                "        {name} = craneLib.buildPackage (commonArgs // {{\n\
          inherit cargoArtifacts;\n\
          pname = \"{name}\";\n\
          {features}\n\
        }});\n"
            );
        }
        result
    }

    /// Generate packages attribute
    fn generate_packages_attr(&self, crate_names: &[&str]) -> String {
        use std::fmt::Write;
        let mut result = String::new();
        for name in crate_names {
            let _ = writeln!(&mut result, "          {name} = {name};");
        }
        result
    }

    /// Generate Cachix configuration
    pub fn generate_cachix_config(&self) -> String {
        format!(
            r#"# Cachix configuration for PAIML Sovereign Stack
# Push: cachix push paiml $(nix-build)
# Use:  cachix use paiml

{{
  "name": "paiml",
  "signing_key_path": "$HOME/.config/cachix/cachix.dhall",
  "binary_caches": [
    {{
      "url": "https://paiml.cachix.org",
      "public_signing_keys": ["paiml.cachix.org-1:..."]
    }}
  ],
  "crates": {:?},
  "rust_version": "{}",
  "generated": "{}"
}}
"#,
            self.crates.iter().map(|c| &c.name).collect::<Vec<_>>(),
            self.rust_version,
            chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ"),
        )
    }

    /// Generate a minimal flake for a single crate
    pub fn minimal_flake(crate_name: &str, crate_path: &str) -> String {
        format!(
            r#"{{
  description = "{crate_name} - PAIML component";

  inputs = {{
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    rust-overlay.url = "github:oxalica/rust-overlay";
    crane.url = "github:ipetkov/crane";
    flake-utils.url = "github:numtide/flake-utils";
  }};

  outputs = {{ self, nixpkgs, rust-overlay, crane, flake-utils, ... }}:
    flake-utils.lib.eachDefaultSystem (system:
      let
        pkgs = import nixpkgs {{
          inherit system;
          overlays = [ (import rust-overlay) ];
        }};
        craneLib = crane.mkLib pkgs;
      in {{
        packages.default = craneLib.buildPackage {{
          src = ./{crate_path};
          pname = "{crate_name}";
        }};

        devShells.default = pkgs.mkShell {{
          inputsFrom = [ self.packages.${{system}}.default ];
          buildInputs = with pkgs; [ rust-analyzer ];
        }};
      }});
}}
"#
        )
    }
}

impl Default for NixFlakeConfig {
    fn default() -> Self {
        Self::sovereign_stack()
    }
}