//! Document corpus source for ingesting markdown documentation.
//!
//! Reads a directory tree of markdown files, splits on ATX headings,
//! and emits one Record per section with breadcrumb paths attached.

use anyhow::Result;
use mem_chunk::RecordSource;
use mem_core::{Provenance, Record, Role};
use futures::stream::{self, Stream};
use sha2::{Sha256, Digest};
use std::path::{Path, PathBuf};
use time::OffsetDateTime;
use walkdir::WalkDir;

/// Maximum file size to process (10MB).
const MAX_FILE_SIZE: u64 = 10 * 1024 * 1024;

/// Supported file extensions.
const SUPPORTED_EXTENSIONS: &[&str] = &["md", "markdown", "txt"];

/// A document section extracted from a markdown file.
#[derive(Clone, Debug)]
pub struct DocSection {
    /// The heading path (e.g., "kubectl.md > Common Issues > CrashLoopBackOff")
    pub breadcrumb: String,
    /// The section content (heading + body)
    pub content: String,
    /// Source URI (absolute path or URL)
    pub source_uri: String,
    /// SHA256 hash of the original file
    pub file_hash: String,
    /// Whether this is a continuation of an over-long section
    pub continuation: bool,
    /// Section index within the file
    pub section_index: usize,
}

/// Document corpus source that reads markdown files from a directory.
pub struct DocCorpusSource {
    root_path: PathBuf,
    max_tokens: usize,
}

impl DocCorpusSource {
    /// Create a new document corpus source.
    ///
    /// # Arguments
    /// * `root_path` - Root directory to scan for markdown files
    /// * `max_tokens` - Maximum tokens per section (for splitting over-long sections)
    pub fn new(root_path: impl AsRef<Path>, max_tokens: usize) -> Self {
        DocCorpusSource {
            root_path: root_path.as_ref().to_path_buf(),
            max_tokens,
        }
    }

    /// Scan the directory and return all sections.
    pub fn scan(&self) -> Result<Vec<DocSection>> {
        let mut sections = Vec::new();

        for entry in WalkDir::new(&self.root_path)
            .follow_links(false) // Don't follow symlinks to avoid loops
            .into_iter()
            .filter_entry(|e| !is_hidden(e))
        {
            let entry = entry?;
            let path = entry.path();

            // Skip directories
            if path.is_dir() {
                continue;
            }

            // Check extension
            let ext = path.extension()
                .and_then(|e| e.to_str())
                .map(|e| e.to_lowercase());
            
            if !ext.map(|e| SUPPORTED_EXTENSIONS.contains(&e.as_str())).unwrap_or(false) {
                continue;
            }

            // Check file size
            let metadata = std::fs::metadata(path)?;
            if metadata.len() > MAX_FILE_SIZE {
                continue;
            }

            // Read and process the file
            let content = std::fs::read_to_string(path)?;
            let file_hash = compute_sha256(&content);
            let source_uri = path.canonicalize()?.to_string_lossy().to_string();
            let filename = path.file_name()
                .and_then(|n| n.to_str())
                .unwrap_or("unknown");

            let file_sections = split_on_headings(&content, filename, &source_uri, &file_hash, self.max_tokens);
            sections.extend(file_sections);
        }

        Ok(sections)
    }

    /// Dry-run: scan and print the plan without making any network calls.
    pub fn dry_run(&self) -> Result<DryRunReport> {
        let sections = self.scan()?;
        
        let mut files = std::collections::HashMap::new();
        for section in &sections {
            let entry = files.entry(section.source_uri.clone()).or_insert_with(|| FileReport {
                path: section.source_uri.clone(),
                sections: Vec::new(),
                total_tokens: 0,
            });
            let tokens = section.content.len() / 4; // Approximate
            entry.sections.push(SectionReport {
                breadcrumb: section.breadcrumb.clone(),
                tokens,
                continuation: section.continuation,
            });
            entry.total_tokens += tokens;
        }

        Ok(DryRunReport {
            total_files: files.len(),
            total_sections: sections.len(),
            files: files.into_values().collect(),
        })
    }
}

/// Check if a directory entry is hidden (starts with .)
fn is_hidden(entry: &walkdir::DirEntry) -> bool {
    entry.file_name()
        .to_str()
        .map(|s| s.starts_with('.'))
        .unwrap_or(false)
}

/// Compute SHA256 hash of content.
fn compute_sha256(content: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(content.as_bytes());
    format!("{:x}", hasher.finalize())
}

/// Split markdown content on ATX headings.
fn split_on_headings(
    content: &str,
    filename: &str,
    source_uri: &str,
    file_hash: &str,
    max_tokens: usize,
) -> Vec<DocSection> {
    let mut sections = Vec::new();
    let mut heading_stack: Vec<(usize, String)> = Vec::new(); // (level, text)
    let mut current_content = String::new();
    let mut section_index = 0;

    // Helper to build breadcrumb from stack
    let build_breadcrumb = |stack: &[(usize, String)]| -> String {
        let mut parts = vec![filename.to_string()];
        parts.extend(stack.iter().map(|(_, text)| text.clone()));
        parts.join(" > ")
    };

    // Helper to flush current section
    let flush_section = |content: &mut String, stack: &[(usize, String)], sections: &mut Vec<DocSection>, index: &mut usize| {
        let trimmed = content.trim();
        // Skip sections that are only a heading (no body content)
        // A section with only one line that starts with # is just a heading
        let lines: Vec<&str> = trimmed.lines().collect();
        let has_body = lines.len() > 1 || (lines.len() == 1 && !lines[0].starts_with('#'));
        if !trimmed.is_empty() && has_body {
            let breadcrumb = build_breadcrumb(stack);
            let tokens = trimmed.len() / 4;

            // Check if section is over-long and needs splitting
            if tokens > max_tokens {
                let split_sections = split_overlong_section(trimmed, &breadcrumb, source_uri, file_hash, max_tokens, *index);
                *index += split_sections.len();
                sections.extend(split_sections);
            } else {
                // Prepend breadcrumb to content for embedding
                let content_with_breadcrumb = format!("[{}]\n\n{}", breadcrumb, trimmed);
                sections.push(DocSection {
                    breadcrumb,
                    content: content_with_breadcrumb,
                    source_uri: source_uri.to_string(),
                    file_hash: file_hash.to_string(),
                    continuation: false,
                    section_index: *index,
                });
                *index += 1;
            }
        }
        content.clear();
    };

    for line in content.lines() {
        // Check for ATX heading (# to ######)
        if let Some(heading) = parse_atx_heading(line) {
            // Flush previous section
            flush_section(&mut current_content, &heading_stack, &mut sections, &mut section_index);

            // Update heading stack
            // Pop all headings at same or deeper level
            while heading_stack.last().map(|(l, _)| *l >= heading.level).unwrap_or(false) {
                heading_stack.pop();
            }
            heading_stack.push((heading.level, heading.text.clone()));

            // Start new section with the heading
            current_content.push_str(line);
            current_content.push('\n');
        } else {
            // Regular content
            current_content.push_str(line);
            current_content.push('\n');
        }
    }

    // Flush final section
    flush_section(&mut current_content, &heading_stack, &mut sections, &mut section_index);

    // If no sections were created (file has no headings), create one section for entire file
    if sections.is_empty() && !content.trim().is_empty() {
        let breadcrumb = filename.to_string();
        let tokens = content.len() / 4;

        if tokens > max_tokens {
            sections = split_overlong_section(content.trim(), &breadcrumb, source_uri, file_hash, max_tokens, 0);
        } else {
            let content_with_breadcrumb = format!("[{}]\n\n{}", breadcrumb, content.trim());
            sections.push(DocSection {
                breadcrumb,
                content: content_with_breadcrumb,
                source_uri: source_uri.to_string(),
                file_hash: file_hash.to_string(),
                continuation: false,
                section_index: 0,
            });
        }
    }

    sections
}

/// Parse an ATX heading line.
struct AtxHeading {
    level: usize,
    text: String,
}

fn parse_atx_heading(line: &str) -> Option<AtxHeading> {
    let trimmed = line.trim_start();
    
    // Count leading #s
    let hashes = trimmed.chars().take_while(|&c| c == '#').count();
    
    // Must be 1-6 hashes followed by space or end of line
    if hashes == 0 || hashes > 6 {
        return None;
    }

    let rest = &trimmed[hashes..];
    if !rest.is_empty() && !rest.starts_with(' ') && !rest.starts_with('\t') {
        return None;
    }

    let text = rest.trim().trim_end_matches('#').trim().to_string();
    
    Some(AtxHeading { level: hashes, text })
}

/// Split an over-long section into smaller chunks.
fn split_overlong_section(
    content: &str,
    breadcrumb: &str,
    source_uri: &str,
    file_hash: &str,
    max_tokens: usize,
    start_index: usize,
) -> Vec<DocSection> {
    let mut sections = Vec::new();
    let max_chars = max_tokens * 4; // Approximate chars from tokens

    // First, try to split on paragraph boundaries
    let paragraphs: Vec<&str> = content.split("\n\n").collect();
    
    let mut current_content = String::new();
    let mut chunk_index = 0;

    for para in paragraphs {
        let para_with_sep = if current_content.is_empty() {
            para.to_string()
        } else {
            format!("\n\n{}", para)
        };

        if current_content.len() + para_with_sep.len() > max_chars && !current_content.is_empty() {
            // Flush current chunk
            let is_continuation = chunk_index > 0;
            let content_with_breadcrumb = if is_continuation {
                format!("[{} (continued)]\n\n{}", breadcrumb, current_content.trim())
            } else {
                format!("[{}]\n\n{}", breadcrumb, current_content.trim())
            };

            sections.push(DocSection {
                breadcrumb: breadcrumb.to_string(),
                content: content_with_breadcrumb,
                source_uri: source_uri.to_string(),
                file_hash: file_hash.to_string(),
                continuation: is_continuation,
                section_index: start_index + chunk_index,
            });

            current_content = para.to_string();
            chunk_index += 1;
        } else {
            if current_content.is_empty() {
                current_content = para.to_string();
            } else {
                current_content.push_str(&para_with_sep);
            }
        }
    }

    // Flush remaining content
    if !current_content.trim().is_empty() {
        let is_continuation = chunk_index > 0;
        let content_with_breadcrumb = if is_continuation {
            format!("[{} (continued)]\n\n{}", breadcrumb, current_content.trim())
        } else {
            format!("[{}]\n\n{}", breadcrumb, current_content.trim())
        };

        sections.push(DocSection {
            breadcrumb: breadcrumb.to_string(),
            content: content_with_breadcrumb,
            source_uri: source_uri.to_string(),
            file_hash: file_hash.to_string(),
            continuation: is_continuation,
            section_index: start_index + chunk_index,
        });
    }

    // If a single paragraph is still too long, hard split
    if sections.len() == 1 && sections[0].content.len() > max_chars {
        sections = hard_split(&sections[0].content, breadcrumb, source_uri, file_hash, max_chars, start_index);
    }

    sections
}

/// Hard split content at character boundaries when paragraph splitting isn't enough.
fn hard_split(
    content: &str,
    breadcrumb: &str,
    source_uri: &str,
    file_hash: &str,
    max_chars: usize,
    start_index: usize,
) -> Vec<DocSection> {
    let mut sections = Vec::new();
    let chars: Vec<char> = content.chars().collect();
    let mut chunk_index = 0;
    let mut start = 0;

    while start < chars.len() {
        let end = (start + max_chars).min(chars.len());
        let chunk: String = chars[start..end].iter().collect();
        
        let is_continuation = chunk_index > 0;
        let content_with_breadcrumb = if is_continuation {
            format!("[{} (continued)]\n\n{}", breadcrumb, chunk.trim())
        } else {
            chunk.to_string() // First chunk already has breadcrumb
        };

        sections.push(DocSection {
            breadcrumb: breadcrumb.to_string(),
            content: content_with_breadcrumb,
            source_uri: source_uri.to_string(),
            file_hash: file_hash.to_string(),
            continuation: is_continuation,
            section_index: start_index + chunk_index,
        });

        start = end;
        chunk_index += 1;
    }

    sections
}

/// Dry-run report.
#[derive(Debug)]
pub struct DryRunReport {
    pub total_files: usize,
    pub total_sections: usize,
    pub files: Vec<FileReport>,
}

/// Per-file report in dry-run.
#[derive(Debug)]
pub struct FileReport {
    pub path: String,
    pub sections: Vec<SectionReport>,
    pub total_tokens: usize,
}

/// Per-section report in dry-run.
#[derive(Debug)]
pub struct SectionReport {
    pub breadcrumb: String,
    pub tokens: usize,
    pub continuation: bool,
}

impl std::fmt::Display for DryRunReport {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "=== Document Corpus Dry Run ===")?;
        writeln!(f, "Total files: {}", self.total_files)?;
        writeln!(f, "Total sections: {}", self.total_sections)?;
        writeln!(f)?;

        for file in &self.files {
            writeln!(f, "📄 {}", file.path)?;
            writeln!(f, "   Total tokens: ~{}", file.total_tokens)?;
            for section in &file.sections {
                let cont = if section.continuation { " (cont)" } else { "" };
                writeln!(f, "   └─ {} (~{} tokens){}", section.breadcrumb, section.tokens, cont)?;
            }
            writeln!(f)?;
        }

        Ok(())
    }
}

impl RecordSource for DocCorpusSource {
    fn records(self) -> Box<dyn Stream<Item = Result<Record, String>> + Unpin> {
        // Scan sections synchronously (blocking is okay for file I/O)
        let sections = match self.scan() {
            Ok(s) => s,
            Err(e) => {
                return Box::new(stream::iter(vec![
                    Err(format!("Failed to scan corpus: {}", e))
                ]));
            }
        };

        let records: Vec<Result<Record, String>> = sections
            .into_iter()
            .enumerate()
            .map(|(i, section)| {
                Ok(Record {
                    role: Role::User, // Documentation is always "user" content
                    text: section.content,
                    timestamp: OffsetDateTime::now_utc(),
                    provenance: Provenance {
                        source_id: section.source_uri,
                        offset: i as u64,
                    },
                })
            })
            .collect();

        Box::new(stream::iter(records))
    }
}

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

    #[test]
    fn test_parse_atx_heading() {
        assert!(parse_atx_heading("# Heading").is_some());
        assert!(parse_atx_heading("## Sub").is_some());
        assert!(parse_atx_heading("###### Deep").is_some());
        assert!(parse_atx_heading("####### Too deep").is_none());
        assert!(parse_atx_heading("#NoSpace").is_none());
        assert!(parse_atx_heading("Not a heading").is_none());
        
        let h = parse_atx_heading("## My Heading ##").unwrap();
        assert_eq!(h.level, 2);
        assert_eq!(h.text, "My Heading");
    }

    #[test]
    fn test_split_on_headings_simple() {
        let content = r#"# Title

Some intro text.

## Section 1

Content 1.

## Section 2

Content 2.
"#;
        let sections = split_on_headings(content, "test.md", "/test.md", "abc123", 5000);
        
        assert_eq!(sections.len(), 3);
        assert_eq!(sections[0].breadcrumb, "test.md > Title");
        assert_eq!(sections[1].breadcrumb, "test.md > Title > Section 1");
        assert_eq!(sections[2].breadcrumb, "test.md > Title > Section 2");
    }

    #[test]
    fn test_split_on_headings_nested() {
        let content = r#"# Top

## A

### A1

Content A1.

### A2

Content A2.

## B

Content B.
"#;
        let sections = split_on_headings(content, "nested.md", "/nested.md", "xyz", 5000);
        
        // Sections without body content (only headings) are skipped
        // So we get: A1, A2, B (3 sections, not 4)
        assert_eq!(sections.len(), 3);
        assert_eq!(sections[0].breadcrumb, "nested.md > Top > A > A1");
        assert_eq!(sections[1].breadcrumb, "nested.md > Top > A > A2");
        assert_eq!(sections[2].breadcrumb, "nested.md > Top > B");
    }

    #[test]
    fn test_no_headings() {
        let content = "Just some text without any headings.";
        let sections = split_on_headings(content, "plain.md", "/plain.md", "hash", 5000);
        
        assert_eq!(sections.len(), 1);
        assert_eq!(sections[0].breadcrumb, "plain.md");
        assert!(!sections[0].continuation);
    }

    #[test]
    fn test_overlong_section_splits() {
        let content = format!("# Big Section\n\n{}", "word ".repeat(2000)); // ~10000 chars
        let sections = split_on_headings(&content, "big.md", "/big.md", "hash", 500); // Force split
        
        assert!(sections.len() > 1);
        assert!(!sections[0].continuation);
        assert!(sections[1].continuation);
    }
}