//! Document chunking strategies for different file types

use std::path::Path;

use regex::Regex;

use crate::embeddings::DocumentChunk;
use crate::error::Result;

/// Document chunking engine
pub struct ChunkingEngine {
    chunk_size: usize,    // Tokens per chunk (approximate)
    chunk_overlap: usize, // Tokens of overlap between chunks
}

impl ChunkingEngine {
    /// Create a new chunking engine
    pub fn new(chunk_size: usize, chunk_overlap: usize) -> Self {
        Self {
            chunk_size,
            chunk_overlap,
        }
    }

    /// Create a document chunk with metadata
    fn create_chunk(
        &self,
        source_path: &str,
        chunk_id: usize,
        doc_type: &str,
        content: String,
        headings: &[String],
    ) -> DocumentChunk {
        let mut metadata = std::collections::HashMap::new();
        let heading_path = headings.join(" > ");
        if !heading_path.is_empty() {
            metadata.insert("heading_path".to_string(), heading_path);
        }
        DocumentChunk {
            id: format!("{}-{}", source_path, chunk_id),
            source_path: source_path.to_string(),
            doc_type: doc_type.to_string(),
            content: content.trim().to_string(),
            category: None,
            metadata,
        }
    }

    /// Chunk markdown document by heading hierarchy
    pub fn chunk_markdown(&self, content: &str, source_path: &str) -> Result<Vec<DocumentChunk>> {
        let mut chunks = Vec::new();
        let lines: Vec<&str> = content.lines().collect();

        let mut current_chunk = String::new();
        let mut current_headings: Vec<String> = Vec::new();
        let mut current_heading_levels: Vec<usize> = Vec::new();
        let mut chunk_id_counter = 0;

        for line in &lines {
            // Check if line is a heading
            if let Some(level) = get_heading_level(line) {
                // Update heading hierarchy
                while !current_heading_levels.is_empty()
                    && current_heading_levels.last().unwrap() >= &level
                {
                    current_heading_levels.pop();
                    current_headings.pop();
                }

                // Add new heading
                let heading_text = line.trim_start_matches('#').trim().to_string();
                current_headings.push(heading_text);
                current_heading_levels.push(level);

                // Add heading to current chunk
                if !current_chunk.is_empty() {
                    current_chunk.push('\n');
                }
                current_chunk.push_str(line);
            } else {
                // Regular content line
                current_chunk.push_str(line);
                current_chunk.push('\n');

                // Check if chunk is large enough
                let approx_tokens = approximate_token_count(&current_chunk);
                if approx_tokens >= self.chunk_size {
                    let chunk = self.create_chunk(
                        source_path,
                        chunk_id_counter,
                        "markdown",
                        current_chunk.clone(),
                        &current_headings,
                    );
                    chunks.push(chunk);
                    chunk_id_counter += 1;

                    // Start new chunk with overlap
                    let overlap_lines = (self.chunk_overlap / 4).max(1);
                    let lines_to_keep = lines
                        .iter()
                        .rev()
                        .take(overlap_lines)
                        .map(|l| l.to_string())
                        .collect::<Vec<_>>();
                    current_chunk = lines_to_keep
                        .iter()
                        .rev()
                        .cloned()
                        .collect::<Vec<_>>()
                        .join("\n");
                }
            }
        }

        // Add final chunk if not empty
        if !current_chunk.trim().is_empty() {
            let mut metadata = std::collections::HashMap::new();
            if !current_headings.is_empty() {
                metadata.insert("heading_path".to_string(), current_headings.join(" > "));
            }
            let chunk = DocumentChunk {
                id: format!("{}-{}", source_path, chunk_id_counter),
                source_path: source_path.to_string(),
                doc_type: "markdown".to_string(),
                content: current_chunk.trim().to_string(),
                category: None,
                metadata,
            };

            chunks.push(chunk);
        }

        tracing::debug!(
            "Chunked markdown file: {} into {} chunks",
            source_path,
            chunks.len()
        );

        Ok(chunks)
    }

    /// Chunk KCL schema files by schema definition
    pub fn chunk_kcl(&self, content: &str, source_path: &str) -> Result<Vec<DocumentChunk>> {
        let mut chunks = Vec::new();
        let mut chunk_id_counter = 0;

        // Find all schema definitions
        let schema_pattern = Regex::new(r"^schema\s+(\w+)")?;
        let lines: Vec<&str> = content.lines().collect();

        let mut current_chunk = String::new();
        let mut current_schema_name = String::new();
        let mut brace_count = 0;

        for line in &lines {
            // Check if line starts a new schema
            if let Some(caps) = schema_pattern.captures(line) {
                // Save previous schema if exists
                if !current_chunk.is_empty() {
                    let mut metadata = std::collections::HashMap::new();
                    metadata.insert("schema_name".to_string(), current_schema_name.clone());
                    let chunk = DocumentChunk {
                        id: format!("{}-{}", source_path, chunk_id_counter),
                        source_path: source_path.to_string(),
                        doc_type: "kcl".to_string(),
                        content: current_chunk.trim().to_string(),
                        category: None,
                        metadata,
                    };

                    chunks.push(chunk);
                    chunk_id_counter += 1;
                }

                current_schema_name = caps.get(1).unwrap().as_str().to_string();
                current_chunk = line.to_string();
                brace_count = line.matches('{').count() as i32 - line.matches('}').count() as i32;
            } else {
                current_chunk.push('\n');
                current_chunk.push_str(line);
                brace_count += line.matches('{').count() as i32;
                brace_count -= line.matches('}').count() as i32;

                // Schema definition ends when braces are balanced
                if brace_count == 0 && !current_chunk.is_empty() {
                    let mut metadata = std::collections::HashMap::new();
                    metadata.insert("schema_name".to_string(), current_schema_name.clone());
                    let chunk = DocumentChunk {
                        id: format!("{}-{}", source_path, chunk_id_counter),
                        source_path: source_path.to_string(),
                        doc_type: "kcl".to_string(),
                        content: current_chunk.trim().to_string(),
                        category: None,
                        metadata,
                    };

                    chunks.push(chunk);
                    chunk_id_counter += 1;
                    current_chunk.clear();
                    current_schema_name.clear();
                }
            }
        }

        // Add final chunk if exists
        if !current_chunk.trim().is_empty() {
            let mut metadata = std::collections::HashMap::new();
            metadata.insert("schema_name".to_string(), current_schema_name);
            let chunk = DocumentChunk {
                id: format!("{}-{}", source_path, chunk_id_counter),
                source_path: source_path.to_string(),
                doc_type: "kcl".to_string(),
                content: current_chunk.trim().to_string(),
                category: None,
                metadata,
            };

            chunks.push(chunk);
        }

        tracing::debug!(
            "Chunked KCL file: {} into {} chunks",
            source_path,
            chunks.len()
        );

        Ok(chunks)
    }

    /// Chunk Nushell scripts by function definition
    pub fn chunk_nushell(&self, content: &str, source_path: &str) -> Result<Vec<DocumentChunk>> {
        let mut chunks = Vec::new();
        let mut chunk_id_counter = 0;

        // Find all function definitions
        let fn_pattern = Regex::new(r"^(?:export\s+)?def\s+([^\s\[{]+)")?;
        let lines: Vec<&str> = content.lines().collect();

        let mut current_chunk = String::new();
        let mut current_function_name = String::new();
        let mut brace_count = 0;

        for line in &lines {
            // Check if line starts a new function
            if let Some(caps) = fn_pattern.captures(line) {
                // Save previous function if exists
                if !current_chunk.is_empty() {
                    let mut metadata = std::collections::HashMap::new();
                    metadata.insert("function_name".to_string(), current_function_name.clone());
                    let chunk = DocumentChunk {
                        id: format!("{}-{}", source_path, chunk_id_counter),
                        source_path: source_path.to_string(),
                        doc_type: "nushell".to_string(),
                        content: current_chunk.trim().to_string(),
                        category: None,
                        metadata,
                    };

                    chunks.push(chunk);
                    chunk_id_counter += 1;
                }

                current_function_name = caps.get(1).unwrap().as_str().to_string();
                current_chunk = line.to_string();
                brace_count = line.matches('{').count() as i32 - line.matches('}').count() as i32;
            } else {
                current_chunk.push('\n');
                current_chunk.push_str(line);
                brace_count += line.matches('{').count() as i32;
                brace_count -= line.matches('}').count() as i32;

                // Function definition ends when braces are balanced (and we've seen some
                // content)
                if brace_count == 0 && !current_chunk.is_empty() && line.contains('}') {
                    let mut metadata = std::collections::HashMap::new();
                    metadata.insert("function_name".to_string(), current_function_name.clone());
                    let chunk = DocumentChunk {
                        id: format!("{}-{}", source_path, chunk_id_counter),
                        source_path: source_path.to_string(),
                        doc_type: "nushell".to_string(),
                        content: current_chunk.trim().to_string(),
                        category: None,
                        metadata,
                    };

                    chunks.push(chunk);
                    chunk_id_counter += 1;
                    current_chunk.clear();
                    current_function_name.clear();
                }
            }
        }

        // Add final chunk if exists
        if !current_chunk.trim().is_empty() {
            let mut metadata = std::collections::HashMap::new();
            metadata.insert("function_name".to_string(), current_function_name);
            let chunk = DocumentChunk {
                id: format!("{}-{}", source_path, chunk_id_counter),
                source_path: source_path.to_string(),
                doc_type: "nushell".to_string(),
                content: current_chunk.trim().to_string(),
                category: None,
                metadata,
            };

            chunks.push(chunk);
        }

        tracing::debug!(
            "Chunked Nushell file: {} into {} chunks",
            source_path,
            chunks.len()
        );

        Ok(chunks)
    }

    /// Detect file type and chunk appropriately
    pub fn chunk_file(&self, content: &str, file_path: &str) -> Result<Vec<DocumentChunk>> {
        let path = Path::new(file_path);
        let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");

        match ext {
            "md" => self.chunk_markdown(content, file_path),
            "k" => self.chunk_kcl(content, file_path),
            "nu" => self.chunk_nushell(content, file_path),
            _ => {
                // Default: treat as markdown-like
                tracing::warn!("Unknown file type: {}, treating as text", ext);
                Ok(vec![DocumentChunk {
                    id: format!("{}-0", file_path),
                    source_path: file_path.to_string(),
                    doc_type: "unknown".to_string(),
                    content: content.to_string(),
                    category: None,
                    metadata: std::collections::HashMap::new(),
                }])
            }
        }
    }
}

/// Get heading level (1-6 for H1-H6)
fn get_heading_level(line: &str) -> Option<usize> {
    let trimmed = line.trim_start();
    if trimmed.starts_with('#') {
        let level = trimmed.chars().take_while(|c| *c == '#').count();
        if level <= 6 && trimmed.chars().nth(level) == Some(' ') {
            return Some(level);
        }
    }
    None
}

/// Approximate token count (rough heuristic: 1 token ≈ 4 chars)
fn approximate_token_count(text: &str) -> usize {
    (text.len() / 4).max(1)
}

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

    #[test]
    fn test_heading_detection() {
        assert_eq!(get_heading_level("# Title"), Some(1));
        assert_eq!(get_heading_level("## Subtitle"), Some(2));
        assert_eq!(get_heading_level("### Section"), Some(3));
        assert_eq!(get_heading_level("Not a heading"), None);
        assert_eq!(get_heading_level("#NoSpace"), None);
    }

    #[test]
    fn test_chunk_markdown() {
        let engine = ChunkingEngine::new(500, 100);
        let content = r#"
# Main Title

This is introduction text.

## Subsection

More content here.

### Sub-subsection

Even more content.
"#;

        let chunks = engine
            .chunk_markdown(content, "test.md")
            .expect("chunking failed");

        assert!(!chunks.is_empty());
        assert!(chunks.iter().all(|c| c.doc_type == "markdown"));
    }

    #[test]
    fn test_chunk_kcl() {
        let engine = ChunkingEngine::new(500, 100);
        let content = r#"schema Storage:
    device: str
    size: int

schema Server:
    name: str
    cpu: int
"#;

        let chunks = engine
            .chunk_kcl(content, "test.k")
            .expect("chunking failed");

        // At least one chunk should be created
        assert!(!chunks.is_empty());
        // Verify all chunks have metadata
        assert!(chunks
            .iter()
            .all(|c| c.doc_type == "kcl" && c.metadata.contains_key("schema_name")));
    }

    #[test]
    fn test_chunk_nushell() {
        let engine = ChunkingEngine::new(500, 100);
        let content = r#"
def get-version [] {
    "1.0.0"
}

export def main [name: string] {
    print $"Hello, ($name)!"
}
"#;

        let chunks = engine
            .chunk_nushell(content, "test.nu")
            .expect("chunking failed");

        assert!(!chunks.is_empty());
        assert!(chunks
            .iter()
            .all(|c| c.metadata.contains_key("function_name")));
    }

    #[test]
    fn test_token_approximation() {
        let text = "Hello world";
        let tokens = approximate_token_count(text);
        assert!(tokens > 0);
    }
}
