use ruff_db::files::{File, FilePath, system_path_to_file};
use ruff_db::system::SystemPath;
use ty_python_semantic::{
    ModuleName, resolve_module, resolve_module_confident, resolve_real_module,
    resolve_real_module_confident,
};

use crate::ModuleDb;
use crate::collector::CollectedImport;

/// Collect all imports for a given Python file.
pub(crate) struct Resolver<'a> {
    db: &'a ModuleDb,
    file: Option<File>,
}

impl<'a> Resolver<'a> {
    /// Initialize a [`Resolver`] with a given [`ModuleDb`].
    pub(crate) fn new(db: &'a ModuleDb, path: &SystemPath) -> Self {
        // If we know the importing file we can potentially resolve more imports
        let file = system_path_to_file(db, path).ok();
        Self { db, file }
    }

    /// Resolve the [`CollectedImport`] into a [`FilePath`].
    pub(crate) fn resolve(&self, import: CollectedImport) -> impl Iterator<Item = &'a FilePath> {
        match import {
            CollectedImport::Import(import) => {
                // Attempt to resolve the module (e.g., given `import foo`, look for `foo`).
                let file = self.resolve_module(&import);

                // If the file is a stub, look for the corresponding source file.
                let source_file = file
                    .is_some_and(|file| file.extension() == Some("pyi"))
                    .then(|| self.resolve_real_module(&import))
                    .flatten();

                std::iter::once(file)
                    .chain(std::iter::once(source_file))
                    .flatten()
            }
            CollectedImport::ImportFrom(import) => {
                // Attempt to resolve the member (e.g., given `from foo import bar`, look for `foo.bar`).
                if let Some(file) = self.resolve_module(&import) {
                    // If the file is a stub, look for the corresponding source file.
                    let source_file = (file.extension() == Some("pyi"))
                        .then(|| self.resolve_real_module(&import))
                        .flatten();

                    return std::iter::once(Some(file))
                        .chain(std::iter::once(source_file))
                        .flatten();
                }

                // Attempt to resolve the module (e.g., given `from foo import bar`, look for `foo`).
                let parent = import.parent();
                let file = parent
                    .as_ref()
                    .and_then(|parent| self.resolve_module(parent));

                // If the file is a stub, look for the corresponding source file.
                let source_file = file
                    .is_some_and(|file| file.extension() == Some("pyi"))
                    .then(|| {
                        parent
                            .as_ref()
                            .and_then(|parent| self.resolve_real_module(parent))
                    })
                    .flatten();

                std::iter::once(file)
                    .chain(std::iter::once(source_file))
                    .flatten()
            }
        }
    }

    /// Resolves a module name to a module.
    pub(crate) fn resolve_module(&self, module_name: &ModuleName) -> Option<&'a FilePath> {
        let module = if let Some(file) = self.file {
            resolve_module(self.db, file, module_name)?
        } else {
            resolve_module_confident(self.db, module_name)?
        };
        Some(module.file(self.db)?.path(self.db))
    }

    /// Resolves a module name to a module (stubs not allowed).
    fn resolve_real_module(&self, module_name: &ModuleName) -> Option<&'a FilePath> {
        let module = if let Some(file) = self.file {
            resolve_real_module(self.db, file, module_name)?
        } else {
            resolve_real_module_confident(self.db, module_name)?
        };
        Some(module.file(self.db)?.path(self.db))
    }
}