//! Tree-sitter based source code parser.

use std::collections::HashMap;

use tree_sitter::{Language as TsLanguage, Parser, Tree};

use crate::chunker::{chunk_type_from_node_kind, extract_symbol_name};
use crate::config::ChunkingConfig;
use crate::error::{Error, Result};
use crate::types::{Chunk, ChunkId, ChunkType, Document, Language};

/// Tree-sitter based parser for multiple languages.
pub struct TreeSitterParser {
    parsers: HashMap<Language, Parser>,
}

impl TreeSitterParser {
    /// Create a new tree-sitter parser with all supported languages.
    pub fn new() -> Result<Self> {
        let mut parsers = HashMap::new();

        // Initialize parsers for each language
        #[cfg(feature = "rust")]
        {
            let mut parser = Parser::new();
            parser
                .set_language(&tree_sitter_rust::language())
                .map_err(|e| Error::parse("rust", e.to_string()))?;
            parsers.insert(Language::Rust, parser);
        }

        #[cfg(feature = "python")]
        {
            let mut parser = Parser::new();
            parser
                .set_language(&tree_sitter_python::language())
                .map_err(|e| Error::parse("python", e.to_string()))?;
            parsers.insert(Language::Python, parser);
        }

        #[cfg(feature = "javascript")]
        {
            let mut parser = Parser::new();
            parser
                .set_language(&tree_sitter_javascript::language())
                .map_err(|e| Error::parse("javascript", e.to_string()))?;
            parsers.insert(Language::JavaScript, parser);
        }

        #[cfg(feature = "typescript")]
        {
            let mut parser = Parser::new();
            parser
                .set_language(&tree_sitter_typescript::language_typescript())
                .map_err(|e| Error::parse("typescript", e.to_string()))?;
            parsers.insert(Language::TypeScript, parser);
        }

        #[cfg(feature = "go")]
        {
            let mut parser = Parser::new();
            parser
                .set_language(&tree_sitter_go::language())
                .map_err(|e| Error::parse("go", e.to_string()))?;
            parsers.insert(Language::Go, parser);
        }

        #[cfg(feature = "c")]
        {
            let mut parser = Parser::new();
            parser
                .set_language(&tree_sitter_c::language())
                .map_err(|e| Error::parse("c", e.to_string()))?;
            parsers.insert(Language::C, parser);
        }

        #[cfg(feature = "cpp")]
        {
            let mut parser = Parser::new();
            parser
                .set_language(&tree_sitter_cpp::language())
                .map_err(|e| Error::parse("cpp", e.to_string()))?;
            parsers.insert(Language::Cpp, parser);
        }

        #[cfg(feature = "java")]
        {
            let mut parser = Parser::new();
            parser
                .set_language(&tree_sitter_java::language())
                .map_err(|e| Error::parse("java", e.to_string()))?;
            parsers.insert(Language::Java, parser);
        }

        Ok(Self { parsers })
    }

    /// Check if a language is supported.
    pub fn supports_language(&self, language: Language) -> bool {
        self.parsers.contains_key(&language)
    }

    /// Parse source code and extract chunks.
    pub fn parse_and_chunk(
        &self,
        document: &Document,
        content: &str,
        config: &ChunkingConfig,
    ) -> Result<Vec<Chunk>> {
        let parser = self
            .parsers
            .get(&document.language)
            .ok_or_else(|| Error::parse(&document.relative_path, "Unsupported language"))?;

        // Parse with tree-sitter (need mutable borrow)
        let tree = {
            let mut parser = Parser::new();
            parser
                .set_language(self.get_language(document.language)?)
                .map_err(|e| Error::parse(&document.relative_path, e.to_string()))?;
            parser
                .parse(content, None)
                .ok_or_else(|| Error::parse(&document.relative_path, "Failed to parse"))?
        };

        // Extract chunks from the AST
        let mut chunks = Vec::new();
        self.extract_chunks(
            &tree,
            content,
            document,
            config,
            &mut chunks,
            None, // No parent symbol at root
        );

        // Filter and validate chunks
        let chunks: Vec<Chunk> = chunks
            .into_iter()
            .filter(|c| {
                c.content.len() >= config.min_chunk_size && c.content.len() <= config.max_chunk_size
            })
            .collect();

        Ok(chunks)
    }

    fn get_language(&self, language: Language) -> Result<&TsLanguage> {
        match language {
            #[cfg(feature = "rust")]
            Language::Rust => Ok(tree_sitter_rust::language()),
            #[cfg(feature = "python")]
            Language::Python => Ok(tree_sitter_python::language()),
            #[cfg(feature = "javascript")]
            Language::JavaScript => Ok(tree_sitter_javascript::language()),
            #[cfg(feature = "typescript")]
            Language::TypeScript => Ok(tree_sitter_typescript::language_typescript()),
            #[cfg(feature = "go")]
            Language::Go => Ok(tree_sitter_go::language()),
            #[cfg(feature = "c")]
            Language::C => Ok(tree_sitter_c::language()),
            #[cfg(feature = "cpp")]
            Language::Cpp => Ok(tree_sitter_cpp::language()),
            #[cfg(feature = "java")]
            Language::Java => Ok(tree_sitter_java::language()),
            _ => Err(Error::parse("unknown", "Unsupported language")),
        }
    }

    fn extract_chunks(
        &self,
        tree: &Tree,
        content: &str,
        document: &Document,
        config: &ChunkingConfig,
        chunks: &mut Vec<Chunk>,
        parent_symbol: Option<String>,
    ) {
        let root = tree.root_node();
        self.visit_node(root, content, document, config, chunks, parent_symbol);
    }

    fn visit_node(
        &self,
        node: tree_sitter::Node,
        content: &str,
        document: &Document,
        config: &ChunkingConfig,
        chunks: &mut Vec<Chunk>,
        parent_symbol: Option<String>,
    ) {
        let kind = node.kind();
        let chunk_type = chunk_type_from_node_kind(kind, document.language);

        // Determine if this node should be a chunk
        let should_chunk = match chunk_type {
            ChunkType::Function | ChunkType::Class | ChunkType::Struct | ChunkType::Enum
            | ChunkType::Interface | ChunkType::Implementation | ChunkType::Module => true,
            ChunkType::Imports => config.chunk_imports,
            ChunkType::Comment => config.extract_comments && node.byte_range().len() > 50,
            ChunkType::Constant | ChunkType::TypeDef => true,
            ChunkType::Block => false,
        };

        if should_chunk {
            let start_byte = node.start_byte();
            let end_byte = node.end_byte();

            if let Some(node_content) = content.get(start_byte..end_byte) {
                // Only create chunk if content is within size limits
                if node_content.len() >= config.min_chunk_size {
                    let symbol_name = extract_symbol_name(node_content, chunk_type, document.language);

                    let chunk = Chunk {
                        id: ChunkId::new(),
                        document_id: document.id.clone(),
                        content: node_content.to_string(),
                        chunk_type,
                        start_line: node.start_position().row as u32 + 1,
                        end_line: node.end_position().row as u32 + 1,
                        start_byte,
                        end_byte,
                        symbol_name: symbol_name.clone(),
                        parent_symbol: parent_symbol.clone(),
                    };

                    // If chunk is too large, we might want to break it down
                    if node_content.len() > config.max_chunk_size {
                        // For very large constructs, visit children instead
                        let new_parent = symbol_name.or(parent_symbol);
                        let mut cursor = node.walk();
                        for child in node.children(&mut cursor) {
                            self.visit_node(
                                child,
                                content,
                                document,
                                config,
                                chunks,
                                new_parent.clone(),
                            );
                        }
                    } else {
                        chunks.push(chunk);
                    }
                    return;
                }
            }
        }

        // Visit children
        let new_parent = if should_chunk {
            extract_symbol_name(
                content.get(node.byte_range()).unwrap_or(""),
                chunk_type,
                document.language,
            )
            .or(parent_symbol)
        } else {
            parent_symbol
        };

        let mut cursor = node.walk();
        for child in node.children(&mut cursor) {
            self.visit_node(child, content, document, config, chunks, new_parent.clone());
        }
    }
}

impl Default for TreeSitterParser {
    fn default() -> Self {
        Self::new().expect("Failed to initialize tree-sitter parser")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    fn make_document(language: Language) -> Document {
        Document {
            id: crate::types::DocumentId::new(),
            relative_path: PathBuf::from("test.rs"),
            absolute_path: PathBuf::from("/test/test.rs"),
            language,
            content_hash: "test".to_string(),
            size_bytes: 0,
            modified_at: chrono::Utc::now(),
        }
    }

    #[test]
    #[cfg(feature = "rust")]
    fn test_parse_rust_function() {
        let parser = TreeSitterParser::new().unwrap();
        let doc = make_document(Language::Rust);
        let content = r#"
fn hello_world() {
    println!("Hello, world!");
}

pub fn add(a: i32, b: i32) -> i32 {
    a + b
}
"#;

        let config = ChunkingConfig::default();
        let chunks = parser.parse_and_chunk(&doc, content, &config).unwrap();

        assert!(!chunks.is_empty());
        assert!(chunks.iter().any(|c| c.symbol_name == Some("hello_world".to_string())));
        assert!(chunks.iter().any(|c| c.symbol_name == Some("add".to_string())));
    }
}
