r2x_python/

initialization.rs

//! Python environment initialization and setup
//!
//! This module handles all Python interpreter initialization, virtual environment
//! configuration, and environment setup required before the bridge can be used.

use super::utils::{resolve_python_path, resolve_site_package_path};
use crate::errors::BridgeError;
use once_cell::sync::OnceCell;
use pyo3::prelude::*;
use pyo3::types::PyModule;
use r2x_config::Config;
use r2x_logger as logger;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;
use which::which;

#[cfg(windows)]
const PYTHON_BIN_DIR_NAME: &str = "Scripts";
#[cfg(not(windows))]
const PYTHON_BIN_DIR_NAME: &str = "bin";
#[cfg(windows)]
const SYSTEM_PYTHON_CANDIDATES: &[&str] = &["python.exe", "py.exe"];
#[cfg(not(windows))]
const SYSTEM_PYTHON_CANDIDATES: &[&str] = &["python3", "python"];
const REQUIRED_PYTHON_MAJOR: i32 = 3;
const REQUIRED_PYTHON_MINOR: i32 = 12;

pub struct Bridge {}

#[derive(Debug, Clone)]
pub struct PythonEnvironment {
    pub interpreter: PathBuf,
    pub python_home: Option<PathBuf>,
}

static BRIDGE_INSTANCE: OnceCell<Result<Bridge, BridgeError>> = OnceCell::new();

impl Bridge {
    /// Get or initialize the bridge singleton
    pub fn get() -> Result<&'static Bridge, BridgeError> {
        match BRIDGE_INSTANCE.get_or_init(Bridge::initialize) {
            Ok(bridge) => Ok(bridge),
            Err(e) => Err(BridgeError::Initialization(format!("{}", e))),
        }
    }

    /// Initialize Python interpreter and configure environment
    ///
    /// This performs:
    /// - Configure venv/python environment
    /// - Add venv site-packages to sys.path
    /// - Ensure r2x-core is installed
    fn initialize() -> Result<Bridge, BridgeError> {
        let start_time = std::time::Instant::now();

        let python_env = configure_python_venv()?;
        let python_path = python_env.interpreter.clone();

        let mut config = Config::load()
            .map_err(|e| BridgeError::Initialization(format!("Failed to load config: {}", e)))?;
        let cache_path = config.ensure_cache_path().map_err(|e| {
            BridgeError::Initialization(format!("Failed to ensure cache path: {}", e))
        })?;
        let venv_path = PathBuf::from(config.get_venv_path());
        let site_packages = resolve_site_package_path(&venv_path)?;
        if let Some(ref home) = python_env.python_home {
            configure_embedded_python_env(home, &site_packages);
        } else {
            logger::debug("Using default embedded Python search paths");
        }

        logger::debug(&format!(
            "Initializing Python bridge with: {}",
            python_path.display()
        ));

        let pyo3_start = std::time::Instant::now();
        pyo3::Python::initialize();
        logger::debug(&format!(
            "pyo3::Python::initialize took: {:?}",
            pyo3_start.elapsed()
        ));

        // Enable Python bytecode generation for faster subsequent imports
        // This overrides PYTHONDONTWRITEBYTECODE if set in the environment
        pyo3::Python::attach(|py| {
            let sys = PyModule::import(py, "sys")
                .map_err(|e| BridgeError::Python(format!("Failed to import sys module: {}", e)))?;
            sys.setattr("dont_write_bytecode", false).map_err(|e| {
                BridgeError::Python(format!("Failed to enable bytecode generation: {}", e))
            })?;
            Ok::<(), BridgeError>(())
        })?;
        logger::debug("Enabled Python bytecode generation");

        // Add site-packages from venv to sys.path so imports work as expected
        logger::debug(&format!(
            "site_packages: {}, exists: {}",
            site_packages.display(),
            site_packages.exists()
        ));

        pyo3::Python::attach(|py| {
            let site = PyModule::import(py, "site")
                .map_err(|e| BridgeError::Python(format!("Failed to import site module: {}", e)))?;
            site.call_method1("addsitedir", (site_packages.to_str().unwrap(),))
                .map_err(|e| BridgeError::Python(format!("Failed to add site directory: {}", e)))?;
            Ok::<(), BridgeError>(())
        })?;

        let sitedir_start = std::time::Instant::now();
        logger::debug(&format!(
            "Site packages setup completed in: {:?}",
            sitedir_start.elapsed()
        ));

        // Detect and store the compiled Python version in config if not already set
        let version_start = std::time::Instant::now();
        detect_and_store_python_version()?;
        logger::debug(&format!(
            "Python version detection took: {:?}",
            version_start.elapsed()
        ));

        configure_python_cache(&cache_path)?;

        // r2x_core is now installed during venv creation, so no need to check here

        // Configure Python loguru to write to the same log file as Rust
        // Python logs always go to file, --log-python flag controls console output
        logger::debug("Starting Python logging configuration...");
        if let Err(e) = Self::configure_python_logging() {
            logger::warn(&format!("Python logging configuration failed: {}", e));
        }
        logger::debug("Python logging configuration completed");

        logger::debug(&format!(
            "Total bridge initialization took: {:?}",
            start_time.elapsed()
        ));
        Ok(Bridge {})
    }

    /// Configure Python loguru logging to integrate with Rust logger
    fn configure_python_logging() -> Result<(), BridgeError> {
        let log_file = logger::get_log_path_string();
        let verbosity = logger::get_verbosity();
        let log_level = match verbosity {
            0 => "WARNING",
            1 => "INFO",
            2 => "DEBUG",
            _ => "TRACE",
        };

        // Format to match Rust logger: [YYYY-MM-DD HH:MM:SS] [PYTHON] LEVEL message
        let fmt = "[{time:YYYY-MM-DD HH:mm:ss}] [PYTHON] {level: <8} {message}";

        // Check if Python logs should be shown on console
        let enable_console = logger::get_log_python();

        logger::debug(&format!(
            "Configuring Python logging with level={}, file={}, enable_console={}",
            log_level, log_file, enable_console
        ));

        pyo3::Python::attach(|py| {
            let logger_module = PyModule::import(py, "r2x_core.logger").map_err(|e| {
                logger::warn(&format!("Failed to import r2x_core.logger: {}", e));
                BridgeError::Import("r2x_core.logger".to_string(), format!("{}", e))
            })?;
            let setup_logging = logger_module.getattr("setup_logging").map_err(|e| {
                logger::warn(&format!("Failed to get setup_logging function: {}", e));
                BridgeError::Python(format!("setup_logging not found: {}", e))
            })?;
            let kwargs = pyo3::types::PyDict::new(py);
            kwargs.set_item("level", log_level)?;
            kwargs.set_item("log_file", &log_file)?;
            kwargs.set_item("fmt", fmt)?;
            kwargs.set_item("enable_console_log", enable_console)?;
            setup_logging.call((), Some(&kwargs))?;

            // Explicitly enable logging for r2x modules
            let loguru = PyModule::import(py, "loguru")?;
            let logger = loguru.getattr("logger")?;
            logger.call_method1("enable", ("r2x_core",))?;
            logger.call_method1("enable", ("r2x_reeds",))?;
            logger.call_method1("enable", ("r2x_plexos",))?;
            logger.call_method1("enable", ("r2x_sienna",))?;

            Ok::<(), BridgeError>(())
        })
    }
}

/// Helper: get python3.X directory inside venv lib/
/// Detect the Python version from the embedded interpreter and store it in config
///
/// This function:
/// 1. Gets the Python version from sys.version_info (the compiled/embedded version)
/// 2. Compares it with what's stored in config
/// 3. If missing or mismatched, updates config to the actual version
/// 4. Logs warnings if there's a mismatch (indicates config was manually edited)
fn detect_and_store_python_version() -> Result<(), BridgeError> {
    let mut config = Config::load()
        .map_err(|e| BridgeError::Initialization(format!("Failed to load config: {}", e)))?;

    // Get Python version from sys.version_info (the actual compiled version)
    let version_str = pyo3::Python::attach(|py| {
        let sys = PyModule::import(py, "sys")
            .map_err(|e| BridgeError::Python(format!("Failed to import sys: {}", e)))?;
        let version_info = sys
            .getattr("version_info")
            .map_err(|e| BridgeError::Python(format!("Failed to get version_info: {}", e)))?;

        let major = version_info
            .getattr("major")
            .map_err(|e| BridgeError::Python(format!("Failed to get major: {}", e)))?
            .extract::<i32>()
            .map_err(|e| BridgeError::Python(format!("Failed to extract major: {}", e)))?;

        let minor = version_info
            .getattr("minor")
            .map_err(|e| BridgeError::Python(format!("Failed to get minor: {}", e)))?
            .extract::<i32>()
            .map_err(|e| BridgeError::Python(format!("Failed to extract minor: {}", e)))?;

        Ok::<String, BridgeError>(format!("{}.{}", major, minor))
    })?;

    logger::debug(&format!("Detected Python version: {}", version_str));

    // Check if config version matches detected version
    if let Some(ref config_version) = config.python_version {
        if config_version == &version_str {
            // Versions match, nothing to do
            return Ok(());
        } else {
            // Mismatch detected - config was likely manually edited
            logger::warn(&format!(
                "Python version mismatch: binary was compiled with {}, but config shows {}. Updating config to match compiled version.",
                version_str, config_version
            ));
        }
    } else {
        // First time detection
        logger::debug("First time detecting Python version for this binary");
    }

    // Store/update the actual compiled version in config
    config.python_version = Some(version_str.clone());
    config
        .save()
        .map_err(|e| BridgeError::Initialization(format!("Failed to save config: {}", e)))?;

    logger::info(&format!("Python version {} stored in config", version_str));

    Ok(())
}

fn configure_python_cache(cache_path: &str) -> Result<(), BridgeError> {
    std::fs::create_dir_all(cache_path).map_err(|e| {
        BridgeError::Initialization(format!("Failed to create cache directory: {}", e))
    })?;
    std::env::set_var("R2X_CACHE_PATH", cache_path);

    let cache_path_escaped = cache_path.replace('\\', "\\\\");
    pyo3::Python::attach(|py| {
        let patch_code = format!(
            r#"from pathlib import Path
_R2X_CACHE_PATH = Path(r"{cache}")

def _r2x_cache_path_override():
    return _R2X_CACHE_PATH
"#,
            cache = cache_path_escaped
        );

        let code_cstr = std::ffi::CString::new(patch_code).map_err(|e| {
            BridgeError::Python(format!("Failed to prepare cache override script: {}", e))
        })?;
        let filename = std::ffi::CString::new("r2x_cache_patch.py").unwrap();
        let module_name = std::ffi::CString::new("r2x_cache_patch").unwrap();
        let patch_module = PyModule::from_code(
            py,
            code_cstr.as_c_str(),
            filename.as_c_str(),
            module_name.as_c_str(),
        )
        .map_err(|e| BridgeError::Python(format!("Failed to build cache override: {}", e)))?;

        let override_fn = patch_module
            .getattr("_r2x_cache_path_override")
            .map_err(|e| {
                BridgeError::Python(format!("Failed to obtain cache override function: {}", e))
            })?;

        let file_ops = PyModule::import(py, "r2x_core.utils.file_operations").map_err(|e| {
            BridgeError::Python(format!(
                "Failed to import r2x_core.utils.file_operations: {}",
                e
            ))
        })?;

        file_ops
            .setattr("get_r2x_cache_path", override_fn)
            .map_err(|e| BridgeError::Python(format!("Failed to override cache path: {}", e)))?;

        Ok::<(), BridgeError>(())
    })?;

    Ok(())
}

/// Configure the Python virtual environment before PyO3 initialization
pub fn configure_python_venv() -> Result<PythonEnvironment, BridgeError> {
    let mut config = Config::load()
        .map_err(|e| BridgeError::Initialization(format!("Failed to load config: {}", e)))?;

    let venv_path = PathBuf::from(config.get_venv_path());

    let python_path_result = resolve_python_path(&venv_path);

    if python_path_result.is_err() {
        logger::debug("Could not resolve Python path");
    }

    let mut python_path = python_path_result.unwrap_or_else(|_| PathBuf::new());

    // Create venv only when it doesn't exist
    if !venv_path.exists() {
        logger::step(&format!(
            "Creating Python virtual environment at: {}",
            venv_path.display()
        ));

        let uv_path = config
            .ensure_uv_path()
            .map_err(|e| BridgeError::Initialization(format!("Failed to ensure uv: {}", e)))?;

        // Use the Python version from config, or default to 3.12
        let python_version = config.python_version.as_deref().unwrap_or("3.12");

        let output = Command::new(&uv_path)
            .arg("venv")
            .arg(&venv_path)
            .arg("--python")
            .arg(python_version)
            .output()?;

        logger::capture_output(&format!("uv venv --python {}", python_version), &output);

        if !output.status.success() {
            return Err(BridgeError::Initialization(
                "Failed to create Python virtual environment".to_string(),
            ));
        }

        python_path = resolve_python_path(&venv_path).unwrap_or_else(|_| PathBuf::new());

        if python_path.as_os_str().is_empty() || !python_path.exists() {
            if let Ok(entries) = std::fs::read_dir(venv_path.join(PYTHON_BIN_DIR_NAME)) {
                let names: Vec<String> = entries
                    .filter_map(|e| e.ok())
                    .filter_map(|e| e.file_name().into_string().ok())
                    .collect();
                logger::debug(&format!("Venv bin contents after creation: {:?}", names));
            }
            return Err(BridgeError::Initialization(
                "Failed to locate Python executable after creating venv".to_string(),
            ));
        }
    }

    if python_path.as_os_str().is_empty() || !python_path.exists() {
        logger::warn("Python binary not found in configured venv; attempting system fallback");
        if let Some((fallback, home)) = find_system_python() {
            return Ok(PythonEnvironment {
                interpreter: fallback,
                python_home: Some(home),
            });
        }

        return Err(BridgeError::Initialization(
            "Failed to locate a usable Python interpreter".to_string(),
        ));
    }

    let python_home = resolve_python_home(&venv_path);

    Ok(PythonEnvironment {
        interpreter: python_path,
        python_home,
    })
}

fn configure_embedded_python_env(python_home: &Path, site_packages: &Path) {
    let home = python_home.to_string_lossy().to_string();
    env::set_var("PYTHONHOME", &home);
    logger::debug(&format!("Set PYTHONHOME={}", home));

    let mut paths = vec![site_packages.to_path_buf()];
    if let Some(existing) = env::var_os("PYTHONPATH") {
        if !existing.is_empty() {
            paths.extend(env::split_paths(&existing));
        }
    }
    if let Ok(joined) = env::join_paths(paths) {
        env::set_var("PYTHONPATH", &joined);
        logger::debug(&format!(
            "Updated PYTHONPATH to include {}",
            site_packages.display()
        ));
    }
}

fn resolve_python_home(venv_path: &Path) -> Option<PathBuf> {
    let cfg_path = venv_path.join("pyvenv.cfg");
    let contents = fs::read_to_string(&cfg_path).ok()?;
    for line in contents.lines() {
        let trimmed = line.trim();
        if trimmed.starts_with("home") {
            let parts: Vec<_> = trimmed.splitn(2, '=').collect();
            if parts.len() == 2 {
                let mut path = PathBuf::from(parts[1].trim());
                if path.ends_with("bin") || path.ends_with("Scripts") {
                    path = path.parent().map(PathBuf::from).unwrap_or(path);
                }
                logger::debug(&format!(
                    "Resolved base Python home {} from {}",
                    path.display(),
                    cfg_path.display()
                ));
                return Some(path);
            }
        }
    }
    logger::debug(&format!(
        "Failed to resolve base Python home from {}",
        cfg_path.display()
    ));
    None
}

fn detect_python_runtime(
    python_bin: &Path,
) -> Option<(PathBuf, i32, i32)> {
    let output = Command::new(python_bin)
        .arg("-c")
        .arg(
            "import sys\nprint(sys.base_prefix)\nprint(sys.version_info.major)\nprint(sys.version_info.minor)",
        )
        .output()
        .ok()?;
    if !output.status.success() {
        logger::debug(&format!(
            "Failed to probe python runtime (status {:?})",
            output.status.code()
        ));
        return None;
    }
    let stdout = String::from_utf8_lossy(&output.stdout).to_string();
    let mut lines = stdout
        .lines()
        .map(|s| s.trim().to_string())
        .collect::<Vec<_>>();
    if lines.len() < 3 {
        return None;
    }
    let minor = lines
        .pop()
        .and_then(|v| v.parse::<i32>().ok())
        .unwrap_or(0);
    let major = lines
        .pop()
        .and_then(|v| v.parse::<i32>().ok())
        .unwrap_or(0);
    let prefix = PathBuf::from(lines.pop().unwrap_or_default());
    Some((prefix, major, minor))
}

fn find_system_python() -> Option<(PathBuf, PathBuf)> {
    for candidate in SYSTEM_PYTHON_CANDIDATES {
        if let Ok(path) = which(candidate) {
            if let Some((home, major, minor)) = detect_python_runtime(&path) {
                if major == REQUIRED_PYTHON_MAJOR && minor == REQUIRED_PYTHON_MINOR {
                    logger::warn(&format!(
                        "Falling back to system Python at {}",
                        path.display()
                    ));
                    return Some((path, home));
                } else {
                    logger::debug(&format!(
                        "Skipping system python {} (version {}.{})",
                        path.display(),
                        major,
                        minor
                    ));
                }
            }
        }
    }
    None
}