lib.rs

use std::{
    fs,
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
};

use log::trace;
use pet_core::{
    env::PythonEnv,
    python_environment::{PythonEnvironment, PythonEnvironmentBuilder, PythonEnvironmentKind},
    pyvenv_cfg::PyVenvCfg,
    reporter::Reporter,
    Configuration, Locator, LocatorKind,
};
use pet_python_utils::executable::find_executables;
use serde::Deserialize;
pub struct Uv {
    pub workspace_directories: Arc<Mutex<Vec<PathBuf>>>,
}

/// Represents information stored in a `pyvenv.cfg` generated by uv
struct UvVenv {
    uv_version: String,
    python_version: String,
    prompt: String,
}

impl UvVenv {
    fn maybe_from_file(file: &Path) -> Option<Self> {
        let contents = fs::read_to_string(file).ok()?;
        let mut uv_version = None;
        let mut python_version = None;
        let mut prompt = None;
        for line in contents.lines() {
            if let Some(uv_version_value) = line.trim_start().strip_prefix("uv = ") {
                uv_version = Some(uv_version_value.trim_end().to_string())
            }
            if let Some(version_info) = line.trim_start().strip_prefix("version_info = ") {
                python_version = Some(version_info.to_string());
            }
            if let Some(prompt_value) = line.trim_start().strip_prefix("prompt = ") {
                prompt = Some(prompt_value.trim_end().to_string());
            }
            if uv_version.is_some() && python_version.is_some() && prompt.is_some() {
                // we've found all the values we need, stop parsing
                break;
            }
        }
        Some(Self {
            uv_version: uv_version?,
            python_version: python_version?,
            prompt: prompt?,
        })
    }
}

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

impl Uv {
    pub fn new() -> Self {
        Self {
            workspace_directories: Arc::new(Mutex::new(Vec::new())),
        }
    }
}

impl Locator for Uv {
    fn get_kind(&self) -> LocatorKind {
        LocatorKind::Uv
    }

    fn supported_categories(&self) -> Vec<PythonEnvironmentKind> {
        vec![
            PythonEnvironmentKind::Uv,
            PythonEnvironmentKind::UvWorkspace,
        ]
    }

    fn configure(&self, config: &Configuration) {
        if let Some(workspace_directories) = config.workspace_directories.as_ref() {
            let mut ws = self
                .workspace_directories
                .lock()
                .expect("workspace_directories mutex poisoned");
            ws.clear();
            ws.extend(workspace_directories.iter().cloned());
        }
    }

    fn try_from(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
        let cfg = env
            .executable
            .parent()
            .and_then(PyVenvCfg::find)
            .or_else(|| {
                env.prefix
                    .as_ref()
                    .and_then(|prefix| PyVenvCfg::find(prefix))
            })?;
        let uv_venv = UvVenv::maybe_from_file(&cfg.file_path)?;
        trace!(
            "uv-managed venv found in {}, made by uv {}",
            env.executable.display(),
            uv_venv.uv_version
        );
        let prefix = env.prefix.clone().or_else(|| {
            env.executable
                .parent()
                .and_then(|p| p.parent().map(|pp| pp.to_path_buf()))
        });
        let pyproject = prefix
            .as_ref()
            .and_then(|prefix| prefix.parent())
            .and_then(parse_pyproject_toml_in);
        let kind = if pyproject
            .and_then(|pyproject| pyproject.tool)
            .and_then(|t| t.uv)
            .and_then(|uv| uv.workspace)
            .is_some()
        {
            PythonEnvironmentKind::UvWorkspace
        } else {
            PythonEnvironmentKind::Uv
        };

        Some(
            PythonEnvironmentBuilder::new(Some(kind))
                .name(Some(uv_venv.prompt))
                .executable(Some(env.executable.clone()))
                .version(Some(uv_venv.python_version))
                .symlinks(prefix.as_ref().map(find_executables))
                .prefix(prefix)
                .build(),
        )
    }

    fn find(&self, reporter: &dyn Reporter) {
        // look through workspace directories for uv-managed projects and any of their workspaces
        let workspaces = self
            .workspace_directories
            .lock()
            .expect("workspace_directories mutex poisoned")
            .clone();
        for workspace in workspaces {
            // TODO: maybe check for workspace in parent folders?
            for env in list_envs_in_directory(&workspace) {
                reporter.report_environment(&env);
            }
        }
    }
}

fn find_workspace(path: &Path) -> Option<PythonEnvironment> {
    for candidate in path.ancestors() {
        let pyproject = parse_pyproject_toml_in(candidate);
        if pyproject
            .as_ref()
            .and_then(|pp| pp.tool.as_ref())
            .and_then(|tool| tool.uv.as_ref())
            .and_then(|uv| uv.workspace.as_ref())
            .is_none()
        {
            continue;
        }
        // TODO: check for workspace members/excludes
        trace!("Found workspace at {:?}", candidate);
        let prefix = candidate.join(".venv");
        let pyvenv_cfg = prefix.join("pyvenv.cfg");
        if !pyvenv_cfg.exists() {
            trace!(
                "Workspace at {} does not have a virtual environment",
                candidate.display()
            );
            return None;
        }
        let unix_executable = prefix.join("bin/python");
        let windows_executable = prefix.join("Scripts/python.exe");
        let executable = if unix_executable.exists() {
            Some(unix_executable)
        } else if windows_executable.exists() {
            Some(windows_executable)
        } else {
            None
        };
        if let Some(uv_venv) = UvVenv::maybe_from_file(&pyvenv_cfg) {
            return Some(
                PythonEnvironmentBuilder::new(Some(PythonEnvironmentKind::UvWorkspace))
                    .name(Some(uv_venv.prompt))
                    .executable(executable)
                    .version(Some(uv_venv.python_version))
                    .symlinks(Some(find_executables(&prefix)))
                    .prefix(Some(prefix))
                    .build(),
            );
        } else {
            trace!(
                "Workspace at {} does not have a uv-managed virtual environment",
                candidate.display()
            );
        }
        return None;
    }
    None
}

fn list_envs_in_directory(path: &Path) -> Vec<PythonEnvironment> {
    let mut envs = Vec::new();
    let pyproject = parse_pyproject_toml_in(path);
    let Some(pyproject) = pyproject else {
        return envs;
    };
    let pyvenv_cfg = path.join(".venv/pyvenv.cfg");
    let prefix = path.join(".venv");
    let unix_executable = prefix.join("bin/python");
    let windows_executable = prefix.join("Scripts/python.exe");
    let executable = if unix_executable.exists() {
        Some(unix_executable)
    } else if windows_executable.exists() {
        Some(windows_executable)
    } else {
        None
    };
    if pyproject
        .tool
        .and_then(|t| t.uv)
        .and_then(|uv| uv.workspace)
        .is_some()
    {
        trace!("Workspace found in {}", path.display());
        if let Some(uv_venv) = UvVenv::maybe_from_file(&pyvenv_cfg) {
            trace!("uv-managed venv found for workspace in {}", path.display());
            let env = PythonEnvironmentBuilder::new(Some(PythonEnvironmentKind::UvWorkspace))
                .name(Some(uv_venv.prompt))
                .symlinks(Some(find_executables(&prefix)))
                .prefix(Some(prefix))
                .executable(executable)
                .version(Some(uv_venv.python_version))
                .build();
            envs.push(env);
        } else {
            trace!(
                "No uv-managed venv found for workspace in {}",
                path.display()
            );
        }
    // prioritize the workspace over the project if it's the same venv
    } else if let Some(project) = pyproject.project {
        if let Some(uv_venv) = UvVenv::maybe_from_file(&pyvenv_cfg) {
            trace!("uv-managed venv found for project in {}", path.display());
            let env = PythonEnvironmentBuilder::new(Some(PythonEnvironmentKind::Uv))
                .name(Some(uv_venv.prompt))
                .symlinks(Some(find_executables(&prefix)))
                .prefix(Some(prefix))
                .version(Some(uv_venv.python_version))
                .display_name(project.name)
                .executable(executable)
                .build();
            envs.push(env);
        } else {
            trace!("No uv-managed venv found in {}", path.display());
        }
        if let Some(workspace) = path.parent().and_then(find_workspace) {
            envs.push(workspace);
        }
    }

    envs
}

fn parse_pyproject_toml_in(path: &Path) -> Option<PyProjectToml> {
    let contents = fs::read_to_string(path.join("pyproject.toml")).ok()?;
    toml::from_str(&contents).ok()
}

#[derive(Deserialize, Debug)]
struct PyProjectToml {
    project: Option<Project>,
    tool: Option<Tool>,
}

#[derive(Deserialize, Debug)]
struct Project {
    name: Option<String>,
}

#[derive(Deserialize, Debug)]
struct Tool {
    uv: Option<ToolUv>,
}

#[derive(Deserialize, Debug)]
struct ToolUv {
    workspace: Option<serde::de::IgnoredAny>,
}

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

    #[test]
    fn test_uv_venv_parse_valid_pyvenv_cfg() {
        let temp_dir = TempDir::new().unwrap();
        let cfg_path = temp_dir.path().join("pyvenv.cfg");

        let contents = r#"home = /usr/bin
include-system-site-packages = false
version = 3.11.0
executable = /usr/bin/python3.11
uv = 0.1.0
version_info = 3.11.0
prompt = test-env"#;

        std::fs::write(&cfg_path, contents).unwrap();

        let uv_venv = UvVenv::maybe_from_file(&cfg_path);
        assert!(uv_venv.is_some());

        let uv_venv = uv_venv.unwrap();
        assert_eq!(uv_venv.uv_version, "0.1.0");
        assert_eq!(uv_venv.python_version, "3.11.0");
        assert_eq!(uv_venv.prompt, "test-env");
    }

    #[test]
    fn test_uv_venv_parse_missing_uv_field() {
        let temp_dir = TempDir::new().unwrap();
        let cfg_path = temp_dir.path().join("pyvenv.cfg");

        let contents = r#"home = /usr/bin
version_info = 3.11.0
prompt = test-env"#;

        std::fs::write(&cfg_path, contents).unwrap();

        let uv_venv = UvVenv::maybe_from_file(&cfg_path);
        assert!(
            uv_venv.is_none(),
            "Should return None when 'uv' field is missing"
        );
    }

    #[test]
    fn test_uv_venv_parse_missing_version_info() {
        let temp_dir = TempDir::new().unwrap();
        let cfg_path = temp_dir.path().join("pyvenv.cfg");

        let contents = r#"home = /usr/bin
uv = 0.1.0
prompt = test-env"#;

        std::fs::write(&cfg_path, contents).unwrap();

        let uv_venv = UvVenv::maybe_from_file(&cfg_path);
        assert!(
            uv_venv.is_none(),
            "Should return None when 'version_info' field is missing"
        );
    }

    #[test]
    fn test_uv_venv_parse_missing_prompt() {
        let temp_dir = TempDir::new().unwrap();
        let cfg_path = temp_dir.path().join("pyvenv.cfg");

        let contents = r#"home = /usr/bin
uv = 0.1.0
version_info = 3.11.0"#;

        std::fs::write(&cfg_path, contents).unwrap();

        let uv_venv = UvVenv::maybe_from_file(&cfg_path);
        assert!(
            uv_venv.is_none(),
            "Should return None when 'prompt' field is missing"
        );
    }

    #[test]
    fn test_uv_venv_parse_with_whitespace() {
        let temp_dir = TempDir::new().unwrap();
        let cfg_path = temp_dir.path().join("pyvenv.cfg");

        let contents = r#"  uv = 0.2.5
  version_info = 3.12.1
  prompt = my-project  "#;

        std::fs::write(&cfg_path, contents).unwrap();

        let uv_venv = UvVenv::maybe_from_file(&cfg_path);
        assert!(uv_venv.is_some());

        let uv_venv = uv_venv.unwrap();
        assert_eq!(uv_venv.uv_version, "0.2.5");
        assert_eq!(uv_venv.python_version, "3.12.1");
        assert_eq!(uv_venv.prompt, "my-project");
    }

    #[test]
    fn test_uv_venv_parse_nonexistent_file() {
        let uv_venv = UvVenv::maybe_from_file(Path::new("/nonexistent/path/pyvenv.cfg"));
        assert!(uv_venv.is_none());
    }

    #[test]
    fn test_parse_pyproject_toml_with_workspace() {
        let temp_dir = TempDir::new().unwrap();
        let pyproject_path = temp_dir.path().join("pyproject.toml");

        let contents = r#"[project]
name = "my-workspace"

[tool.uv.workspace]
members = ["packages/*"]"#;

        std::fs::write(&pyproject_path, contents).unwrap();

        let pyproject = parse_pyproject_toml_in(temp_dir.path());
        assert!(pyproject.is_some());

        let pyproject = pyproject.unwrap();
        assert!(pyproject.project.is_some());
        assert_eq!(
            pyproject.project.unwrap().name,
            Some("my-workspace".to_string())
        );
        assert!(pyproject.tool.is_some());
        assert!(pyproject.tool.unwrap().uv.is_some());
    }

    #[test]
    fn test_parse_pyproject_toml_without_workspace() {
        let temp_dir = TempDir::new().unwrap();
        let pyproject_path = temp_dir.path().join("pyproject.toml");

        let contents = r#"[project]
name = "my-project"

[tool.uv]
dev-dependencies = ["pytest"]"#;

        std::fs::write(&pyproject_path, contents).unwrap();

        let pyproject = parse_pyproject_toml_in(temp_dir.path());
        assert!(pyproject.is_some());

        let pyproject = pyproject.unwrap();
        assert!(pyproject.project.is_some());
        assert_eq!(
            pyproject.project.unwrap().name,
            Some("my-project".to_string())
        );
    }

    #[test]
    fn test_parse_pyproject_toml_missing_file() {
        let temp_dir = TempDir::new().unwrap();
        let pyproject = parse_pyproject_toml_in(temp_dir.path());
        assert!(pyproject.is_none());
    }

    #[test]
    fn test_parse_pyproject_toml_invalid_toml() {
        let temp_dir = TempDir::new().unwrap();
        let pyproject_path = temp_dir.path().join("pyproject.toml");

        let contents = r#"[project
name = "invalid"#;

        std::fs::write(&pyproject_path, contents).unwrap();

        let pyproject = parse_pyproject_toml_in(temp_dir.path());
        assert!(pyproject.is_none());
    }

    #[test]
    fn test_list_envs_in_directory_with_workspace() {
        let temp_dir = TempDir::new().unwrap();
        let project_path = temp_dir.path();

        // Create pyproject.toml with workspace
        let pyproject_path = project_path.join("pyproject.toml");
        let pyproject_contents = r#"[tool.uv.workspace]
members = ["packages/*"]"#;
        std::fs::write(&pyproject_path, pyproject_contents).unwrap();

        // Create .venv directory
        let venv_path = project_path.join(".venv");
        std::fs::create_dir_all(&venv_path).unwrap();

        // Create pyvenv.cfg
        let pyvenv_cfg_path = venv_path.join("pyvenv.cfg");
        let pyvenv_contents = r#"uv = 0.1.0
version_info = 3.11.0
prompt = workspace-env"#;
        std::fs::write(&pyvenv_cfg_path, pyvenv_contents).unwrap();

        // Create executables directory (Unix style for testing)
        let bin_path = venv_path.join("bin");
        std::fs::create_dir_all(&bin_path).unwrap();
        let python_path = bin_path.join("python");
        std::fs::File::create(&python_path).unwrap();

        let envs = list_envs_in_directory(project_path);
        assert_eq!(envs.len(), 1);
        assert_eq!(envs[0].kind, Some(PythonEnvironmentKind::UvWorkspace));
        assert_eq!(envs[0].name, Some("workspace-env".to_string()));
    }

    #[test]
    fn test_list_envs_in_directory_with_project() {
        let temp_dir = TempDir::new().unwrap();
        let project_path = temp_dir.path();

        // Create pyproject.toml with project (no workspace)
        let pyproject_path = project_path.join("pyproject.toml");
        let pyproject_contents = r#"[project]
name = "my-project"

[tool.uv]
dev-dependencies = []"#;
        std::fs::write(&pyproject_path, pyproject_contents).unwrap();

        // Create .venv directory
        let venv_path = project_path.join(".venv");
        std::fs::create_dir_all(&venv_path).unwrap();

        // Create pyvenv.cfg
        let pyvenv_cfg_path = venv_path.join("pyvenv.cfg");
        let pyvenv_contents = r#"uv = 0.1.0
version_info = 3.11.0
prompt = my-project"#;
        std::fs::write(&pyvenv_cfg_path, pyvenv_contents).unwrap();

        // Create executables directory
        let bin_path = venv_path.join("bin");
        std::fs::create_dir_all(&bin_path).unwrap();
        let python_path = bin_path.join("python");
        std::fs::File::create(&python_path).unwrap();

        let envs = list_envs_in_directory(project_path);
        assert_eq!(envs.len(), 1);
        assert_eq!(envs[0].kind, Some(PythonEnvironmentKind::Uv));
        assert_eq!(envs[0].display_name, Some("my-project".to_string()));
    }

    #[test]
    fn test_list_envs_in_directory_no_pyproject() {
        let temp_dir = TempDir::new().unwrap();
        let envs = list_envs_in_directory(temp_dir.path());
        assert_eq!(envs.len(), 0);
    }

    #[test]
    fn test_list_envs_in_directory_no_venv() {
        let temp_dir = TempDir::new().unwrap();
        let project_path = temp_dir.path();

        // Create pyproject.toml but no .venv
        let pyproject_path = project_path.join("pyproject.toml");
        let pyproject_contents = r#"[project]
name = "my-project""#;
        std::fs::write(&pyproject_path, pyproject_contents).unwrap();

        let envs = list_envs_in_directory(project_path);
        assert_eq!(envs.len(), 0);
    }
}