use crate::types::{CodeAction, CompletionItem, Diagnostic, DocumentHighlight, FoldingRange, Hover, InitializeParams, InlayHint, LocationRange, SemanticTokens, SignatureHelp, StructureItem, TextEdit, WorkspaceEdit, WorkspaceSymbol};
use core::range::Range;
use oak_core::{
    language::{ElementRole, ElementType, Language},
    source::Source,
    tree::RedNode,
};
use oak_resolver::ModuleResolver;
use oak_vfs::{Vfs, WritableVfs};
use std::future::Future;

/// A trait that defines the capabilities and behavior of a language-specific service.
///
/// This trait is the primary interface for implementing Language Server Protocol (LSP)
/// features. It provides hooks for various IDE features like hover, completion,
/// diagnostics, and symbol navigation.
///
/// # Implementation
///
/// Implementors should provide language-specific logic for parsing, resolving symbols,
/// and generating IDE-specific data structures.
pub trait LanguageService: Send + Sync {
    /// The language type this service supports.
    type Lang: Language;
    /// The VFS type used for source management.
    type Vfs: WritableVfs;

    /// Returns a reference to the underlying Virtual File System.
    fn vfs(&self) -> &Self::Vfs;

    /// Returns a reference to the workspace manager.
    fn workspace(&self) -> &crate::workspace::WorkspaceManager;

    /// Retrieves the source content for a given URI from the VFS.
    fn get_source(&self, uri: &str) -> Option<<Self::Vfs as Vfs>::Source> {
        self.vfs().get_source(uri)
    }

    /// Retrieves the root red node of a file for the given URI.
    ///
    /// This method is responsible for parsing the source and providing a position-aware
    /// syntax tree. Implementations should typically use a cache to avoid re-parsing
    /// unchanged files.
    fn get_root(&self, _uri: &str) -> impl Future<Output = Option<RedNode<'_, Self::Lang>>> + Send + '_ {
        async { None }
    }

    /// Executes a closure with the root red node of a file.
    ///
    /// This is a convenience helper for running logic that requires the syntax tree.
    fn with_root<'a, R, F>(&'a self, uri: &'a str, f: F) -> impl Future<Output = Option<R>> + Send + 'a
    where
        R: Send,
        F: FnOnce(RedNode<'a, Self::Lang>) -> R + Send + 'a,
    {
        async move {
            let root = self.get_root(uri).await?;
            Some(f(root))
        }
    }

    /// Executes a closure with multiple root nodes in parallel.
    ///
    /// Useful for cross-file operations like workspace symbol search or global rename.
    fn with_roots<'a, R, F>(&'a self, uris: Vec<String>, f: F) -> impl Future<Output = Vec<R>> + Send + 'a
    where
        R: Send + 'static,
        F: Fn(RedNode<'a, Self::Lang>) -> R + Send + Sync + 'a,
    {
        let mut futures = Vec::new();
        let f = std::sync::Arc::new(f);

        for uri in uris {
            let f = f.clone();
            futures.push(async move { if let Some(root) = self.get_root(&uri).await { Some(f(root)) } else { None } })
        }

        async move { futures::future::join_all(futures).await.into_iter().flatten().collect() }
    }

    /// Provides hover information for a specific range in a file.
    ///
    /// # Arguments
    /// * `uri` - The URI of the file.
    /// * `range` - The byte range to provide information for.
    fn hover(&self, _uri: &str, _range: Range<usize>) -> impl Future<Output = Option<Hover>> + Send + '_ {
        async { None }
    }

    /// Provides folding ranges for a file.
    fn folding_ranges(&self, _uri: &str) -> impl Future<Output = Vec<FoldingRange>> + Send + '_ {
        async { vec![] }
    }

    /// Provides document symbols (structure) for a file.
    ///
    /// This method extracts structural elements like classes, functions, and variables
    /// from the source file. It first tries to query the workspace symbol index,
    /// falling back to parsing the file if necessary.
    fn document_symbols<'a>(&'a self, uri: &'a str) -> impl Future<Output = Vec<StructureItem>> + Send + 'a {
        let uri = uri.to_string();
        async move {
            let _source = match self.get_source(&uri) {
                Some(s) => s,
                None => return vec![],
            };
            let _root = match self.get_root(&uri).await {
                Some(r) => r,
                None => return vec![],
            };
            let symbols = self.workspace().symbols.query_file(&uri);
            if !symbols.is_empty() {
                return symbols.into_iter().map(StructureItem::from).collect();
            }
            vec![]
        }
    }

    /// Provides workspace-wide symbol search based on a query string.
    ///
    /// This method searches across all files in the workspace for symbols
    /// that match the given query string.
    fn workspace_symbols<'a>(&'a self, query: String) -> impl Future<Output = Vec<WorkspaceSymbol>> + Send + 'a {
        async move { self.workspace().symbols.query(&query).into_iter().map(|s| WorkspaceSymbol::from(s)).collect() }
    }

    /// Recursively lists all files in the VFS starting from the given root URI.
    ///
    /// This is used to discover all relevant source files in a workspace or directory.
    fn list_all_files(&self, root_uri: &str) -> impl Future<Output = Vec<String>> + Send + '_ {
        let root_uri: oak_core::Arc<str> = root_uri.into();
        async move {
            let mut files = Vec::new();
            let mut stack = vec![root_uri];

            while let Some(uri) = stack.pop() {
                if self.vfs().is_file(&uri) {
                    files.push(uri.to_string());
                }
                else if self.vfs().is_dir(&uri) {
                    if let Some(entries) = self.vfs().read_dir(&uri) {
                        stack.extend(entries);
                    }
                }
            }
            files
        }
    }

    /// Finds the definition(s) of a symbol at the specified range.
    ///
    /// This method attempts to resolve the symbol under the cursor to its
    /// original definition. It handles:
    /// 1. Local symbol resolution (language-specific).
    /// 2. Global symbol lookup via the workspace index.
    /// 3. Module/file import resolution.
    fn definition<'a>(&'a self, uri: &'a str, range: Range<usize>) -> impl Future<Output = Vec<LocationRange>> + Send + 'a {
        let uri = uri.to_string();
        async move {
            let root = match self.get_root(&uri).await {
                Some(r) => r,
                None => return vec![],
            };
            let source = match self.get_source(&uri) {
                Some(s) => s,
                None => return vec![],
            };

            // 1. Identify token at range
            use oak_core::tree::RedTree;
            let node = match root.child_at_offset(range.start) {
                Some(RedTree::Node(n)) => n,
                Some(RedTree::Leaf(l)) => return vec![LocationRange { uri: uri.clone().into(), range: l.span }],
                None => root,
            };

            // 2. If it's a reference, try to resolve it
            let role = node.green.kind.role();
            if role.universal() == oak_core::language::UniversalElementRole::Reference {
                let name = &source.get_text_in(node.span());

                // Try local symbols first (not implemented here, should be done by lang-specific logic)

                // Try global symbols
                if let Some(sym) = self.workspace().symbols.lookup(name) {
                    return vec![LocationRange { uri: sym.uri, range: sym.range }];
                }

                // Try as a module import
                if let Some(resolved_uri) = self.workspace().resolver.resolve(&uri, name) {
                    return vec![LocationRange { uri: resolved_uri.into(), range: (0..0).into() }];
                }

                // Try local symbols (TODO)
            }

            vec![]
        }
    }

    /// Provides document highlights for a symbol at the specified range.
    fn document_highlight<'a>(&'a self, _uri: &'a str, _range: Range<usize>) -> impl Future<Output = Vec<DocumentHighlight>> + Send + 'a {
        async { vec![] }
    }

    /// Provides code actions for a specific range in a file.
    fn code_action<'a>(&'a self, _uri: &'a str, _range: Range<usize>) -> impl Future<Output = Vec<CodeAction>> + Send + 'a {
        async { vec![] }
    }

    /// Provides formatting edits for a file.
    fn formatting<'a>(&'a self, _uri: &'a str) -> impl Future<Output = Vec<TextEdit>> + Send + 'a {
        async { vec![] }
    }

    /// Provides range formatting edits for a file.
    fn range_formatting<'a>(&'a self, _uri: &'a str, _range: Range<usize>) -> impl Future<Output = Vec<TextEdit>> + Send + 'a {
        async { vec![] }
    }

    /// Provides rename edits for a symbol at the specified range.
    fn rename<'a>(&'a self, _uri: &'a str, _range: Range<usize>, _new_name: String) -> impl Future<Output = Option<WorkspaceEdit>> + Send + 'a {
        async { None }
    }

    /// Provides semantic tokens for a file.
    fn semantic_tokens<'a>(&'a self, _uri: &'a str) -> impl Future<Output = Option<SemanticTokens>> + Send + 'a {
        async { None }
    }

    /// Provides completion items for a file at the specified position.
    ///
    /// Implementations should use the syntax tree and symbol index to suggest
    /// relevant keywords, identifiers, and snippets based on the cursor context.
    fn completion<'a>(&'a self, _uri: &'a str, _offset: usize) -> impl Future<Output = Vec<CompletionItem>> + Send + 'a {
        async { vec![] }
    }

    /// Provides signature help for a symbol at the specified range.
    ///
    /// Typically used for function or method calls to show parameter information
    /// and documentation as the user types.
    fn signature_help<'a>(&'a self, _uri: &'a str, _range: Range<usize>) -> impl Future<Output = Option<SignatureHelp>> + Send + 'a {
        async { None }
    }

    /// Provides inlay hints for a file.
    ///
    /// Inlay hints are small pieces of information shown inline with the code,
    /// such as parameter names or inferred types.
    fn inlay_hint<'a>(&'a self, _uri: &'a str, _range: Range<usize>) -> impl Future<Output = Vec<InlayHint>> + Send + 'a {
        async { vec![] }
    }

    /// Finds all references to a symbol at the specified range.
    fn references<'a>(&'a self, _uri: &'a str, _range: Range<usize>) -> impl Future<Output = Vec<LocationRange>> + Send + 'a {
        async { vec![] }
    }

    /// Finds the type definition of a symbol at the specified range.
    fn type_definition<'a>(&'a self, _uri: &'a str, _range: Range<usize>) -> impl Future<Output = Vec<LocationRange>> + Send + 'a {
        async { vec![] }
    }

    /// Finds the implementation(s) of a symbol at the specified range.
    fn implementation<'a>(&'a self, _uri: &'a str, _range: Range<usize>) -> impl Future<Output = Vec<LocationRange>> + Send + 'a {
        async { vec![] }
    }

    /// Handles an LSP initialize request.
    fn initialize<'a>(&'a self, _params: InitializeParams) -> impl Future<Output = ()> + Send + 'a {
        async {}
    }

    /// Called when the language server is fully initialized.
    fn initialized<'a>(&'a self) -> impl Future<Output = ()> + Send + 'a {
        async {}
    }

    /// Called when the language server is shut down.
    fn shutdown<'a>(&'a self) -> impl Future<Output = ()> + Send + 'a {
        async {}
    }

    /// Called when a file is saved in the editor.
    fn did_save<'a>(&'a self, _uri: &'a str) -> impl Future<Output = ()> + Send + 'a {
        async {}
    }

    /// Called when a file is closed in the editor.
    fn did_close<'a>(&'a self, _uri: &'a str) -> impl Future<Output = ()> + Send + 'a {
        async {}
    }

    /// Provides diagnostics for a file.
    fn diagnostics<'a>(&'a self, _uri: &'a str) -> impl Future<Output = Vec<Diagnostic>> + Send + 'a {
        async { vec![] }
    }
}
