use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};

use super::block::{Block, BlockId};
use super::inline::normalize_spans;
use crate::error::EditorError;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Document {
    blocks: HashMap<BlockId, Block>,
    root: Vec<BlockId>,
}

impl Document {
    pub fn new() -> Self {
        Self {
            blocks: HashMap::new(),
            root: Vec::new(),
        }
    }

    // --- Queries ---

    pub fn get_block(&self, id: &BlockId) -> Option<&Block> {
        self.blocks.get(id)
    }

    /// Look up a block by id, returning [`EditorError::BlockNotFound`] if it
    /// does not exist. Use this in transform / apply paths where a missing
    /// block is a hard error rather than a queryable condition.
    pub fn require_block(&self, id: &BlockId) -> Result<&Block, EditorError> {
        self.blocks
            .get(id)
            .ok_or_else(|| EditorError::BlockNotFound(id.clone()))
    }

    pub fn root_block_ids(&self) -> &[BlockId] {
        &self.root
    }

    pub fn root_block_count(&self) -> usize {
        self.root.len()
    }

    /// Walk all blocks in document order, including nested children, calling
    /// `visit` on each block with its indent-derived depth.
    pub fn walk_blocks(&self, mut visit: impl FnMut(&Block, usize)) {
        for id in &self.root {
            if let Some(block) = self.blocks.get(id) {
                let depth = usize::try_from(block.indent_level().max(0)).unwrap_or_default();
                visit(block, depth);
            }
        }
    }

    /// Get the children of `block_id` based on the indent-derived tree.
    /// Children are subsequent root blocks with indent = parent.indent + 1
    /// until interrupted by a block with indent ≤ parent.indent.
    pub fn children_of(&self, block_id: &BlockId) -> Vec<BlockId> {
        let Some(idx) = self.root.iter().position(|id| id == block_id) else {
            return Vec::new();
        };

        let Some(parent_block) = self.blocks.get(&self.root[idx]) else {
            return Vec::new();
        };

        let parent_indent = parent_block.indent_level();
        let target_child_indent = parent_indent + 1;
        let mut children = Vec::new();

        for sibling_id in &self.root[idx + 1..] {
            let Some(block) = self.blocks.get(sibling_id) else {
                continue;
            };

            let indent = block.indent_level();
            if indent <= parent_indent {
                break;
            }
            if indent == target_child_indent {
                children.push(sibling_id.clone());
            }
        }

        children
    }

    /// Get the parent of `block_id` from the indent-derived tree.
    pub fn parent_of(&self, block_id: &BlockId) -> Option<BlockId> {
        let idx = self.root.iter().position(|id| id == block_id)?;
        let block = self.blocks.get(&self.root[idx])?;
        let self_indent = block.indent_level();

        if self_indent == 0 {
            return None;
        }

        for prev_id in self.root[..idx].iter().rev() {
            let Some(prev_block) = self.blocks.get(prev_id) else {
                continue;
            };

            if prev_block.indent_level() < self_indent {
                return Some(prev_id.clone());
            }
        }

        None
    }

    /// All descendants of `block_id` in document (DFS) order — direct
    /// children first, then their children, etc. Reads the populated
    /// `Block.children` field, so callers must ensure tree metadata
    /// is up to date (call `rebuild_tree_metadata` first if needed).
    pub fn descendants_of(&self, block_id: &BlockId) -> Vec<BlockId> {
        let mut out = Vec::new();
        self.collect_descendants(block_id, &mut out);
        out
    }

    fn collect_descendants(&self, block_id: &BlockId, out: &mut Vec<BlockId>) {
        let Some(block) = self.blocks.get(block_id) else {
            return;
        };
        for child_id in &block.children {
            out.push(child_id.clone());
            self.collect_descendants(child_id, out);
        }
    }

    /// Recompute `parent` and `children` fields for **root blocks only**
    /// from the indent-derived tree. Non-root blocks (table rows/cells)
    /// maintain their own parent/children relationships via explicit
    /// `insert_child_block`/`remove_child_block` calls and are **not
    /// touched** by this method.
    ///
    /// For root blocks that also have non-root children (e.g. a table block
    /// has table_row children that live outside root), those children are
    /// **preserved** — only indent-derived (root-to-root) children are reset.
    pub fn rebuild_tree_metadata(&mut self) {
        let root_set: std::collections::HashSet<BlockId> = self.root.iter().cloned().collect();

        // For root blocks: clear indent-based parent, and keep only
        // non-root children (explicitly managed via insert_child_block).
        for (id, block) in &mut self.blocks {
            if root_set.contains(id) {
                block.parent = None;
                // Keep children that are NOT in root (e.g. table rows).
                // Clear only indent-derived children (those that are in root).
                block
                    .children
                    .retain(|child_id| !root_set.contains(child_id));
            }
            // Non-root blocks are left entirely untouched.
        }

        // Rebuild indent-derived parent/children for root blocks.
        let root_ids = self.root.clone();
        for id in &root_ids {
            let Some(parent_id) = self.parent_of(id) else {
                continue;
            };

            if let Some(block) = self.blocks.get_mut(id) {
                block.parent = Some(parent_id.clone());
            }

            if let Some(parent) = self.blocks.get_mut(&parent_id) {
                parent.children.push(id.clone());
            }
        }
    }

    pub fn block_root_index(&self, id: &BlockId) -> Option<usize> {
        self.root.iter().position(|bid| bid == id)
    }

    pub fn prev_sibling(&self, id: &BlockId) -> Option<&BlockId> {
        let idx = self.block_root_index(id)?;
        if idx == 0 {
            None
        } else {
            Some(&self.root[idx - 1])
        }
    }

    pub fn next_sibling(&self, id: &BlockId) -> Option<&BlockId> {
        let idx = self.block_root_index(id)?;
        if idx + 1 < self.root.len() {
            Some(&self.root[idx + 1])
        } else {
            None
        }
    }

    pub fn text_len(&self, id: &BlockId) -> usize {
        self.blocks.get(id).map(|b| b.text_len()).unwrap_or(0)
    }

    pub fn plain_text(&self, id: &BlockId) -> String {
        self.blocks
            .get(id)
            .map(|b| b.plain_text())
            .unwrap_or_default()
    }

    /// Append a block at the end of the root list. Public for use by document
    /// builders (importers, tests). For mutating an existing engine, dispatch
    /// commands instead.
    pub fn push_root_block(&mut self, block: Block) {
        let idx = self.root.len();
        self.insert_block(block, idx);
    }

    // --- pub(crate) mutation helpers (only called by op::apply) ---

    pub(crate) fn insert_block(&mut self, block: Block, index: usize) {
        let id = block.id.clone();
        self.blocks.insert(id.clone(), block);
        let idx = index.min(self.root.len());
        self.root.insert(idx, id);
    }

    pub(crate) fn remove_block(&mut self, id: &BlockId) -> Option<Block> {
        self.root.retain(|bid| bid != id);
        self.blocks.remove(id)
    }

    /// Raw mutable lookup. Only used by tests — production code should go
    /// through [`Document::require_block_mut`] so that a missing block is a
    /// typed [`EditorError`] instead of a silent `None`.
    #[cfg(test)]
    pub(crate) fn get_block_mut(&mut self, id: &BlockId) -> Option<&mut Block> {
        self.blocks.get_mut(id)
    }

    /// Mutable counterpart of [`Document::require_block`]. Returns
    /// [`EditorError::BlockNotFound`] if the block does not exist instead of
    /// silently yielding `None`.
    pub(crate) fn require_block_mut(&mut self, id: &BlockId) -> Result<&mut Block, EditorError> {
        self.blocks
            .get_mut(id)
            .ok_or_else(|| EditorError::BlockNotFound(id.clone()))
    }

    pub(crate) fn set_block_type(
        &mut self,
        id: &BlockId,
        block_type: String,
        props: Map<String, Value>,
    ) -> Option<(String, Map<String, Value>)> {
        self.blocks.get_mut(id).map(|b| {
            let old_type = std::mem::replace(&mut b.ty, block_type);
            let old_props = std::mem::replace(&mut b.props, props);
            (old_type, old_props)
        })
    }

    pub(crate) fn normalize_block(&mut self, id: &BlockId) {
        if let Some(block) = self.blocks.get_mut(id) {
            normalize_spans(&mut block.content);
        }
    }

    pub(crate) fn root_mut(&mut self) -> &mut Vec<BlockId> {
        &mut self.root
    }

    // --- Nested block support (for table rows/cells) ---
    // These blocks live in `self.blocks` but NOT in `self.root`.
    // `rebuild_tree_metadata()` only touches root blocks, so these
    // parent/children relationships are preserved across rebuilds.

    /// Insert a child block under `parent_id` at the given index.
    /// The child is added to `self.blocks` and to the parent's `children` vec,
    /// but NOT to `self.root`.
    pub(crate) fn insert_child_block(
        &mut self,
        parent_id: &BlockId,
        child: Block,
        index: usize,
    ) -> Result<(), crate::error::EditorError> {
        let child_id = child.id.clone();

        // Store the block in the global map
        self.blocks.insert(child_id.clone(), child);

        // Set parent reference on the child
        if let Some(child_block) = self.blocks.get_mut(&child_id) {
            child_block.parent = Some(parent_id.clone());
        }

        // Add to parent's children list
        let parent = self
            .blocks
            .get_mut(parent_id)
            .ok_or_else(|| crate::error::EditorError::BlockNotFound(parent_id.clone()))?;
        let idx = index.min(parent.children.len());
        parent.children.insert(idx, child_id);

        Ok(())
    }

    /// Remove a child block from its parent. Returns the removed block
    /// along with the index it occupied in the parent's children vec.
    pub(crate) fn remove_child_block(
        &mut self,
        parent_id: &BlockId,
        child_id: &BlockId,
    ) -> Result<(Block, usize), crate::error::EditorError> {
        // Find the index in parent's children
        let parent = self
            .blocks
            .get(parent_id)
            .ok_or_else(|| crate::error::EditorError::BlockNotFound(parent_id.clone()))?;
        let index = parent
            .children
            .iter()
            .position(|id| id == child_id)
            .ok_or_else(|| crate::error::EditorError::BlockNotFound(child_id.clone()))?;

        // Remove from parent's children
        let parent = self.blocks.get_mut(parent_id).unwrap();
        parent.children.remove(index);

        // Remove from global blocks map
        let block = self
            .blocks
            .remove(child_id)
            .ok_or_else(|| crate::error::EditorError::BlockNotFound(child_id.clone()))?;

        Ok((block, index))
    }

    /// Get direct children blocks of a parent (by explicit children vec,
    /// not indent-derived). Used for table → rows and row → cells.
    pub fn get_children(&self, parent_id: &BlockId) -> Vec<&Block> {
        let Some(parent) = self.blocks.get(parent_id) else {
            return Vec::new();
        };
        parent
            .children
            .iter()
            .filter_map(|id| self.blocks.get(id))
            .collect()
    }

    /// Build a 2D grid of cell BlockIds for a table.
    /// Returns `Vec<Vec<BlockId>>` where `grid[row][col]` is the cell at that position.
    /// Rows are taken from the table's children, cells from each row's children.
    pub fn table_cell_grid(&self, table_id: &BlockId) -> Vec<Vec<BlockId>> {
        let Some(table) = self.blocks.get(table_id) else {
            return Vec::new();
        };
        table
            .children
            .iter()
            .filter_map(|row_id| self.blocks.get(row_id).map(|row| row.children.clone()))
            .collect()
    }

    /// Check if a block is a non-root nested block (e.g. table_row, table_cell).
    pub fn is_nested_block(&self, id: &BlockId) -> bool {
        self.blocks.contains_key(id) && !self.root.contains(id)
    }
}

impl Default for Document {
    fn default() -> Self {
        Self::new()
    }
}

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

    fn make_doc() -> Document {
        let mut doc = Document::new();
        let b1 = Block::new_paragraph(BlockId::new("b1"));
        let b2 = Block::new_paragraph(BlockId::new("b2"));
        let b3 = Block::new_paragraph(BlockId::new("b3"));
        doc.insert_block(b1, 0);
        doc.insert_block(b2, 1);
        doc.insert_block(b3, 2);
        doc
    }

    #[test]
    fn insert_and_query() {
        let doc = make_doc();
        assert_eq!(doc.root_block_count(), 3);
        assert!(doc.get_block(&BlockId::new("b1")).is_some());
        assert!(doc.get_block(&BlockId::new("b99")).is_none());
    }

    #[test]
    fn sibling_navigation() {
        let doc = make_doc();
        assert_eq!(doc.prev_sibling(&BlockId::new("b1")), None);
        assert_eq!(
            doc.prev_sibling(&BlockId::new("b2")),
            Some(&BlockId::new("b1"))
        );
        assert_eq!(
            doc.next_sibling(&BlockId::new("b2")),
            Some(&BlockId::new("b3"))
        );
        assert_eq!(doc.next_sibling(&BlockId::new("b3")), None);
    }

    #[test]
    fn remove_block() {
        let mut doc = make_doc();
        let removed = doc.remove_block(&BlockId::new("b2"));
        assert!(removed.is_some());
        assert_eq!(doc.root_block_count(), 2);
        assert!(doc.get_block(&BlockId::new("b2")).is_none());
    }

    #[test]
    fn children_and_parent_are_derived_from_indent() {
        let mut doc = Document::new();

        let a = Block::new_paragraph(BlockId::new("a"));

        let mut b = Block::new_paragraph(BlockId::new("b"));
        b.props.insert("indent".into(), Value::from(1));

        let mut c = Block::new_paragraph(BlockId::new("c"));
        c.props.insert("indent".into(), Value::from(2));

        let mut d = Block::new_paragraph(BlockId::new("d"));
        d.props.insert("indent".into(), Value::from(1));

        let e = Block::new_paragraph(BlockId::new("e"));

        doc.push_root_block(a);
        doc.push_root_block(b);
        doc.push_root_block(c);
        doc.push_root_block(d);
        doc.push_root_block(e);

        assert_eq!(
            doc.children_of(&BlockId::new("a")),
            vec![BlockId::new("b"), BlockId::new("d")]
        );
        assert_eq!(doc.children_of(&BlockId::new("b")), vec![BlockId::new("c")]);
        assert_eq!(doc.parent_of(&BlockId::new("b")), Some(BlockId::new("a")));
        assert_eq!(doc.parent_of(&BlockId::new("c")), Some(BlockId::new("b")));
        assert_eq!(doc.parent_of(&BlockId::new("d")), Some(BlockId::new("a")));
        assert_eq!(doc.parent_of(&BlockId::new("e")), None);
    }

    #[test]
    fn rebuild_tree_metadata_sets_parent_and_children() {
        let mut doc = Document::new();

        let a = Block::new_paragraph(BlockId::new("a"));

        let mut b = Block::new_paragraph(BlockId::new("b"));
        b.props.insert("indent".into(), Value::from(1));

        let mut c = Block::new_paragraph(BlockId::new("c"));
        c.props.insert("indent".into(), Value::from(2));

        let mut d = Block::new_paragraph(BlockId::new("d"));
        d.props.insert("indent".into(), Value::from(1));

        doc.push_root_block(a);
        doc.push_root_block(b);
        doc.push_root_block(c);
        doc.push_root_block(d);

        doc.rebuild_tree_metadata();

        let a = doc.get_block(&BlockId::new("a")).unwrap();
        let b = doc.get_block(&BlockId::new("b")).unwrap();
        let c = doc.get_block(&BlockId::new("c")).unwrap();
        let d = doc.get_block(&BlockId::new("d")).unwrap();

        assert_eq!(a.parent, None);
        assert_eq!(a.children, vec![BlockId::new("b"), BlockId::new("d")]);

        assert_eq!(b.parent, Some(BlockId::new("a")));
        assert_eq!(b.children, vec![BlockId::new("c")]);

        assert_eq!(c.parent, Some(BlockId::new("b")));
        assert!(c.children.is_empty());

        assert_eq!(d.parent, Some(BlockId::new("a")));
        assert!(d.children.is_empty());
    }

    #[test]
    fn rebuild_tree_metadata_after_indent_change() {
        let mut doc = Document::new();
        let p1 = Block::new_paragraph(BlockId::new("p1"));
        let p2 = Block::new_paragraph(BlockId::new("p2"));

        doc.push_root_block(p1);
        doc.push_root_block(p2);
        doc.rebuild_tree_metadata();

        assert!(
            doc.get_block(&BlockId::new("p1"))
                .unwrap()
                .children
                .is_empty()
        );

        let p2 = doc.get_block_mut(&BlockId::new("p2")).unwrap();
        p2.props.insert("indent".into(), Value::from(1));
        doc.rebuild_tree_metadata();

        assert_eq!(
            doc.get_block(&BlockId::new("p1")).unwrap().children,
            vec![BlockId::new("p2")]
        );
        assert_eq!(
            doc.get_block(&BlockId::new("p2")).unwrap().parent,
            Some(BlockId::new("p1"))
        );
    }
}
