use crate::args::{AnalyzeGraphArgs, ConfigArguments};
use crate::resolve::resolve;
use crate::{ExitStatus, resolve_default_files};
use anyhow::Result;
use indexmap::IndexSet;
use log::{debug, warn};
use path_absolutize::CWD;
use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
use ruff_graph::{
    Direction, ImportMap, ModuleDb, ModuleImports, ResolverEnvironment, resolve_search_paths,
};
use ruff_linter::package::PackageRoot;
use ruff_linter::source_kind::SourceKind;
use ruff_linter::{warn_user, warn_user_once};
use ruff_python_ast::SourceType;
use ruff_workspace::resolver::{ResolvedFile, match_exclusion, project_files_in_path};
use rustc_hash::{FxBuildHasher, FxHashMap};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

/// Generate an import map.
pub(crate) fn analyze_graph(
    args: AnalyzeGraphArgs,
    config_arguments: &ConfigArguments,
) -> Result<ExitStatus> {
    // Construct the "default" settings. These are used when no `pyproject.toml`
    // files are present, or files are injected from outside the hierarchy.
    let pyproject_config = resolve(config_arguments, None)?;
    if pyproject_config.settings.analyze.preview.is_disabled() {
        warn_user!("`ruff analyze graph` is experimental and may change without warning");
    }

    // Write all paths relative to the current working directory.
    let root =
        SystemPathBuf::from_path_buf(CWD.clone()).expect("Expected a UTF-8 working directory");

    // Find all Python files.
    let files = resolve_default_files(args.files, false);
    let (mut paths, resolver) = project_files_in_path(&files, &pyproject_config, config_arguments)?;

    // Filter to only Python files
    paths.retain(|path| {
        if let Ok(ResolvedFile::Root(path) | ResolvedFile::Nested(path)) = path {
            matches!(SourceType::from(path), SourceType::Python(_))
        } else {
            true
        }
    });

    if paths.is_empty() {
        warn_user_once!("No Python files found under the given path(s)");
        return Ok(ExitStatus::Success);
    }

    // Resolve all package roots.
    let package_roots = resolver
        .package_roots(
            &paths
                .iter()
                .flatten()
                .map(ResolvedFile::path)
                .collect::<Vec<_>>(),
        )
        .into_iter()
        .map(|(path, package)| {
            (
                path.to_path_buf(),
                package.map(PackageRoot::path).map(Path::to_path_buf),
            )
        })
        .collect::<FxHashMap<_, _>>();

    // Create a database from the source roots, combining configured `src` paths with detected
    // package roots. Configured paths are added first so they take precedence, and duplicates
    // are removed.
    let mut src_roots: IndexSet<SystemPathBuf, FxBuildHasher> = IndexSet::default();

    // Add configured `src` paths first (for precedence), filtering to only include existing
    // directories.
    src_roots.extend(
        pyproject_config
            .settings
            .linter
            .src
            .iter()
            .filter(|path| path.is_dir())
            .filter_map(|path| SystemPathBuf::from_path_buf(path.clone()).ok()),
    );

    // Add detected package roots.
    src_roots.extend(
        package_roots
            .values()
            .filter_map(|package| package.as_deref())
            .filter_map(|path| path.parent())
            .filter_map(|path| SystemPathBuf::from_path_buf(path.to_path_buf()).ok()),
    );

    let system = OsSystem::default();
    let search_paths = resolve_search_paths(
        &system,
        src_roots.into_iter().collect(),
        args.python
            .and_then(|python| SystemPathBuf::from_path_buf(python).ok()),
    )?;
    let db = ModuleDb::new(system);
    search_paths.try_register_static_roots(&db);

    let imports = {
        // Create a cache for resolved globs.
        let glob_resolver = Arc::new(Mutex::new(GlobResolver::default()));

        // Collect and resolve the imports for each file.
        let result = Arc::new(Mutex::new(Vec::new()));
        let inner_result = Arc::clone(&result);
        let db = db.clone();
        let search_paths = &search_paths;

        rayon::scope(move |scope| {
            for resolved_file in paths {
                let Ok(resolved_file) = resolved_file else {
                    continue;
                };

                let path = resolved_file.path();
                let package = path
                    .parent()
                    .and_then(|parent| package_roots.get(parent))
                    .and_then(Clone::clone);

                // Resolve the per-file settings.
                let settings = resolver.resolve(path);
                let string_imports = settings.analyze.string_imports;
                let include_dependencies = settings.analyze.include_dependencies.get(path).cloned();
                let type_checking_imports = settings.analyze.type_checking_imports;
                let python_version = settings.analyze.target_version;
                let source_type = settings.analyze.extension.get_source_type(path);

                // Skip excluded files.
                if (settings.file_resolver.force_exclude || !resolved_file.is_root())
                    && match_exclusion(
                        resolved_file.path(),
                        resolved_file.file_name(),
                        &settings.analyze.exclude,
                    )
                {
                    continue;
                }

                // Convert to system paths.
                let Ok(package) = package.map(SystemPathBuf::from_path_buf).transpose() else {
                    warn!("Failed to convert package to system path");
                    continue;
                };
                let Ok(path) = SystemPathBuf::from_path_buf(resolved_file.into_path()) else {
                    warn!("Failed to convert path to system path");
                    continue;
                };

                let db = db.clone();
                let glob_resolver = glob_resolver.clone();
                let root = root.clone();
                let result = inner_result.clone();
                scope.spawn(move |_| {
                    // Extract source code (handles both .py and .ipynb files)
                    let source_kind = match SourceKind::from_path(path.as_std_path(), source_type) {
                        Ok(Some(source_kind)) => source_kind,
                        Ok(None) => {
                            debug!("Skipping non-Python notebook: {path}");
                            return;
                        }
                        Err(err) => {
                            warn!("Failed to read source for {path}: {err}");
                            return;
                        }
                    };

                    let environment = ResolverEnvironment::new(&db, python_version, search_paths);

                    // Identify any imports via static analysis.
                    let mut imports = ModuleImports::detect(
                        &db,
                        environment,
                        &source_kind,
                        &path,
                        package.as_deref(),
                        string_imports,
                        type_checking_imports,
                    )
                    .unwrap_or_else(|err| {
                        warn!("Failed to generate import map for {path}: {err}");
                        ModuleImports::default()
                    });

                    debug!("Discovered {} imports for {}", imports.len(), path);

                    // Append any imports that were statically defined in the configuration.
                    if let Some((root, globs)) = include_dependencies {
                        let mut glob_resolver = glob_resolver.lock().unwrap();
                        imports.extend(glob_resolver.resolve(root, globs));
                    }

                    // Convert the path (and imports) to be relative to the working directory.
                    let path = path
                        .strip_prefix(&root)
                        .map(SystemPath::to_path_buf)
                        .unwrap_or(path);
                    let imports = imports.relative_to(&root);

                    result.lock().unwrap().push((path, imports));
                });
            }
        });

        // Collect the results.
        Arc::into_inner(result).unwrap().into_inner()?
    };

    // Generate the import map.
    let import_map = match args.direction {
        Direction::Dependencies => ImportMap::dependencies(imports),
        Direction::Dependents => ImportMap::dependents(imports),
    };

    // Print to JSON.
    writeln!(
        std::io::stdout(),
        "{}",
        serde_json::to_string_pretty(&import_map)?
    )?;

    std::mem::forget(db);

    Ok(ExitStatus::Success)
}

/// A resolver for glob sets.
#[derive(Default, Debug)]
struct GlobResolver {
    cache: GlobCache,
}

impl GlobResolver {
    /// Resolve a set of globs, anchored at a given root.
    fn resolve(&mut self, root: PathBuf, globs: Vec<String>) -> Vec<SystemPathBuf> {
        if let Some(cached) = self.cache.get(&root, &globs) {
            return cached.clone();
        }

        let walker = match globwalk::GlobWalkerBuilder::from_patterns(&root, &globs)
            .file_type(globwalk::FileType::FILE)
            .build()
        {
            Ok(walker) => walker,
            Err(err) => {
                warn!("Failed to read glob walker: {err}");
                return Vec::new();
            }
        };

        let mut paths = Vec::new();
        for entry in walker {
            let entry = match entry {
                Ok(entry) => entry,
                Err(err) => {
                    warn!("Failed to read glob entry: {err}");
                    continue;
                }
            };
            let path = match SystemPathBuf::from_path_buf(entry.into_path()) {
                Ok(path) => path,
                Err(err) => {
                    warn!("Failed to convert path to system path: {}", err.display());
                    continue;
                }
            };
            paths.push(path);
        }

        self.cache.insert(root, globs, paths.clone());
        paths
    }
}

/// A cache for resolved globs.
#[derive(Default, Debug)]
struct GlobCache(FxHashMap<PathBuf, FxHashMap<Vec<String>, Vec<SystemPathBuf>>>);

impl GlobCache {
    /// Insert a resolved glob.
    fn insert(&mut self, root: PathBuf, globs: Vec<String>, paths: Vec<SystemPathBuf>) {
        self.0.entry(root).or_default().insert(globs, paths);
    }

    /// Get a resolved glob.
    fn get(&self, root: &Path, globs: &[String]) -> Option<&Vec<SystemPathBuf>> {
        self.0.get(root).and_then(|map| map.get(globs))
    }
}
