pub mod output;
pub mod render;

use std::collections::HashSet;

use itertools::Itertools;
use liwe::graph::{Graph, GraphContext};
use liwe::model::node::{NodeIter, NodePointer};
use liwe::model::{Key, NodeId};

use output::{BacklinkInfo, ChildDocumentInfo, DocumentOutput, ParentDocumentInfo, RetrieveOutput};

#[derive(Debug, Clone, Default)]
pub struct RetrieveOptions {
    pub depth: u8,
    pub context: u8,
    pub links: bool,
    pub backlinks: bool,
    pub exclude: HashSet<Key>,
    pub no_content: bool,
}

pub struct DocumentReader<'a> {
    graph: &'a Graph,
}

impl<'a> DocumentReader<'a> {
    pub fn new(graph: &'a Graph) -> Self {
        Self { graph }
    }

    pub fn retrieve(&self, key: &Key, options: &RetrieveOptions) -> RetrieveOutput {
        let mut documents = Vec::new();
        let mut seen_keys = HashSet::new();

        let keys_to_process = self.collect_document_keys(key, options);

        for doc_key in keys_to_process {
            if seen_keys.contains(&doc_key) || options.exclude.contains(&doc_key) {
                continue;
            }
            seen_keys.insert(doc_key.clone());

            let doc_output = self.build_document_output(&doc_key, options);
            documents.push(doc_output);
        }

        RetrieveOutput { documents }
    }

    pub fn retrieve_many(&self, keys: &[Key], options: &RetrieveOptions) -> RetrieveOutput {
        let mut documents = Vec::new();
        let mut seen_keys = HashSet::new();

        for key in keys {
            let keys_to_process = self.collect_document_keys(key, options);

            for doc_key in keys_to_process {
                if seen_keys.contains(&doc_key) || options.exclude.contains(&doc_key) {
                    continue;
                }
                seen_keys.insert(doc_key.clone());

                let doc_output = self.build_document_output(&doc_key, options);
                documents.push(doc_output);
            }
        }

        RetrieveOutput { documents }
    }

    fn collect_document_keys(&self, key: &Key, options: &RetrieveOptions) -> Vec<Key> {
        let mut result = vec![key.clone()];

        if options.depth > 0 {
            let expanded_keys = self.get_expanded_keys(key, options.depth);
            for expanded_key in expanded_keys {
                if expanded_key != *key {
                    result.push(expanded_key);
                }
            }
        }

        if options.context > 0 {
            let context_keys = self.collect_context(key, options.context);
            for context_key in context_keys {
                if context_key != *key && !result.contains(&context_key) {
                    result.push(context_key);
                }
            }

            if options.depth > 0 {
                let sub_doc_parents = self.collect_sub_document_parents(key, options.context);
                for parent_key in sub_doc_parents {
                    if parent_key != *key && !result.contains(&parent_key) {
                        result.push(parent_key);
                    }
                }
            }
        }

        if options.links {
            let linked_keys = self.get_inline_referenced_keys(key);
            for linked_key in linked_keys {
                if linked_key != *key && !result.contains(&linked_key) {
                    result.push(linked_key);
                }
            }
        }

        result
    }

    fn collect_context(&self, key: &Key, levels: u8) -> Vec<Key> {
        if levels == 0 {
            return vec![];
        }

        let mut parents = Vec::new();
        let refs = self.graph.get_block_references_to(key);

        for ref_id in refs {
            let node = self.graph.node(ref_id);

            if let Some(doc_node) = node.to_document() {
                if let Some(parent_key) = doc_node.document_key() {
                    if !parents.contains(&parent_key) {
                        parents.push(parent_key.clone());

                        if levels > 1 {
                            let grandparents = self.collect_context(&parent_key, levels - 1);
                            for grandparent_key in grandparents {
                                if !parents.contains(&grandparent_key) {
                                    parents.push(grandparent_key);
                                }
                            }
                        }
                    }
                }
            }
        }

        parents
    }

    fn collect_sub_document_parents(&self, key: &Key, levels: u8) -> Vec<Key> {
        let mut parents = Vec::new();
        let sub_docs = self.graph.get_block_references_in(key);

        for ref_id in sub_docs {
            if let Some(sub_key) = self.graph.graph_node(ref_id).ref_key() {
                let sub_parents = self.collect_context(&sub_key, levels);
                for parent_key in sub_parents {
                    if parent_key != *key && !parents.contains(&parent_key) {
                        parents.push(parent_key);
                    }
                }
            }
        }

        parents
    }

    fn get_inline_referenced_keys(&self, key: &Key) -> Vec<Key> {
        let mut keys = Vec::new();

        if let Some(node_id) = self.graph.get_node_id(key) {
            let sub_nodes = self.graph.node(node_id).get_all_sub_nodes();

            for sub_node_id in sub_nodes {
                if let Some(line_id) = self.graph.graph_node(sub_node_id).line_id() {
                    let line = self.graph.get_line(line_id);
                    for ref_key in line.ref_keys() {
                        if !keys.contains(&ref_key) {
                            keys.push(ref_key);
                        }
                    }
                }
            }
        }

        keys
    }

    fn get_expanded_keys(&self, key: &Key, depth: u8) -> Vec<Key> {
        let mut keys = Vec::new();

        let refs = self.graph.get_block_references_in(key);
        for ref_id in refs {
            if let Some(ref_key) = self.graph.graph_node(ref_id).ref_key() {
                if !keys.contains(&ref_key) {
                    keys.push(ref_key.clone());

                    if depth > 1 {
                        let sub_keys = self.get_expanded_keys(&ref_key, depth - 1);
                        for sub_key in sub_keys {
                            if !keys.contains(&sub_key) {
                                keys.push(sub_key);
                            }
                        }
                    }
                }
            }
        }

        keys
    }

    fn build_document_output(&self, key: &Key, options: &RetrieveOptions) -> DocumentOutput {
        let title = self
            .graph
            .get_key_title(key)
            .unwrap_or_else(|| key.to_string());

        let content = if options.no_content {
            String::new()
        } else {
            self.get_document_content(key)
        };
        let parent_documents = self.get_parent_documents(key);

        let child_documents = if options.no_content {
            self.get_child_documents(key)
        } else {
            Vec::new()
        };

        let backlinks = if options.backlinks {
            self.get_backlinks(key)
        } else {
            Vec::new()
        };

        DocumentOutput {
            key: key.to_string(),
            title,
            content,
            parent_documents,
            child_documents,
            backlinks,
        }
    }

    fn get_document_content(&self, key: &Key) -> String {
        self.graph.to_markdown_skip_frontmatter(key)
    }

    fn get_parent_documents(&self, key: &Key) -> Vec<ParentDocumentInfo> {
        let refs = self.graph.get_block_references_to(key);
        let mut parents = Vec::new();

        for ref_id in refs {
            let node = self.graph.node(ref_id);

            if let Some(doc_node) = node.to_document() {
                if let Some(doc_key) = doc_node.document_key() {
                    let title = self
                        .graph
                        .get_key_title(&doc_key)
                        .unwrap_or_else(|| doc_key.to_string());

                    let section_path = self.get_section_path(ref_id);

                    parents.push(ParentDocumentInfo {
                        key: doc_key.to_string(),
                        title,
                        section_path,
                    });
                }
            }
        }

        parents.into_iter().unique_by(|p| p.key.clone()).collect()
    }

    fn get_child_documents(&self, key: &Key) -> Vec<ChildDocumentInfo> {
        let refs = self.graph.get_block_references_in(key);
        let mut children = Vec::new();

        for ref_id in refs {
            if let Some(ref_key) = self.graph.graph_node(ref_id).ref_key() {
                let title = self
                    .graph
                    .get_key_title(&ref_key)
                    .unwrap_or_else(|| ref_key.to_string());

                children.push(ChildDocumentInfo {
                    key: ref_key.to_string(),
                    title,
                });
            }
        }

        children.into_iter().unique_by(|c| c.key.clone()).collect()
    }

    fn get_backlinks(&self, key: &Key) -> Vec<BacklinkInfo> {
        let inline_refs = self.graph.get_inline_references_to(key);

        let mut backlinks = Vec::new();
        let mut seen_keys = HashSet::new();

        for ref_id in inline_refs {
            let node = self.graph.node(ref_id);

            if let Some(doc_node) = node.to_document() {
                if let Some(doc_key) = doc_node.document_key() {
                    if seen_keys.contains(&doc_key) {
                        continue;
                    }
                    seen_keys.insert(doc_key.clone());

                    let title = self
                        .graph
                        .get_key_title(&doc_key)
                        .unwrap_or_else(|| doc_key.to_string());

                    let section_path = self.get_section_path(ref_id);

                    backlinks.push(BacklinkInfo {
                        key: doc_key.to_string(),
                        title,
                        section_path,
                    });
                }
            }
        }

        backlinks
    }

    fn get_section_path(&self, node_id: NodeId) -> Vec<String> {
        let mut path = Vec::new();
        let mut current = self.graph.node(node_id);

        while let Some(parent) = current.to_parent() {
            if parent.is_section() && parent.is_header() {
                if let Some(grandparent) = parent.to_parent() {
                    if !grandparent.is_document() {
                        let text = parent.plain_text().trim().to_string();
                        path.push(text);
                    }
                }
            }
            if parent.is_document() {
                break;
            }
            current = parent;
        }

        path.reverse();
        path
    }
}
