use std::{
    collections::BTreeMap,
    sync::{Mutex, MutexGuard},
};

use nana_domain::{RuntimeState, StoryNode};
use thiserror::Error;

#[derive(Debug, Error, PartialEq, Eq)]
pub enum StoreError {
    #[error("story not found: {0}")]
    StoryNotFound(String),
    #[error("branch not found: {story_id}/{branch_id}")]
    BranchNotFound { story_id: String, branch_id: String },
    #[error("node already exists: {0}")]
    NodeAlreadyExists(String),
    #[error("parent node not found: {0}")]
    ParentNotFound(String),
    #[error("branch head changed: expected {expected}, got {actual}")]
    StaleBranchHead { expected: String, actual: String },
    #[error("node and materialized state disagree: {0}")]
    StateMismatch(&'static str),
    #[error("storage lock is poisoned")]
    Poisoned,
}

pub trait StoryStore: Send + Sync {
    fn append_node(&self, node: &StoryNode, state: &RuntimeState) -> Result<(), StoreError>;

    fn load_state(&self, story_id: &str, branch_id: &str) -> Result<RuntimeState, StoreError>;
}

#[derive(Debug, Default)]
struct MemoryData {
    nodes: BTreeMap<(String, String), StoryNode>,
    states: BTreeMap<(String, String), RuntimeState>,
    branch_heads: BTreeMap<(String, String), String>,
}

/// Deterministic test and development store.
///
/// Production persistence will implement the same append-only semantics with
/// SQLite. Keeping this implementation in the crate lets the engine/runtime
/// integration tests exercise branch rules without depending on a database.
#[derive(Debug, Default)]
pub struct InMemoryStoryStore {
    data: Mutex<MemoryData>,
}

impl InMemoryStoryStore {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    pub fn branch_head(
        &self,
        story_id: &str,
        branch_id: &str,
    ) -> Result<Option<String>, StoreError> {
        let data = self.lock()?;
        Ok(data
            .branch_heads
            .get(&(story_id.to_owned(), branch_id.to_owned()))
            .cloned())
    }

    pub fn load_state_at_node(
        &self,
        story_id: &str,
        node_id: &str,
    ) -> Result<RuntimeState, StoreError> {
        let data = self.lock()?;
        data.states
            .get(&(story_id.to_owned(), node_id.to_owned()))
            .cloned()
            .ok_or_else(|| StoreError::ParentNotFound(node_id.to_owned()))
    }

    fn lock(&self) -> Result<MutexGuard<'_, MemoryData>, StoreError> {
        self.data.lock().map_err(|_| StoreError::Poisoned)
    }
}

impl StoryStore for InMemoryStoryStore {
    fn append_node(&self, node: &StoryNode, state: &RuntimeState) -> Result<(), StoreError> {
        validate_materialized_state(node, state)?;

        let node_key = (node.story_id.clone(), node.id.clone());
        let branch_key = (node.story_id.clone(), node.branch_id.clone());
        let mut data = self.lock()?;

        if data.nodes.contains_key(&node_key) {
            return Err(StoreError::NodeAlreadyExists(node.id.clone()));
        }

        if let Some(parent_id) = &node.parent_id {
            if !data
                .nodes
                .contains_key(&(node.story_id.clone(), parent_id.clone()))
            {
                return Err(StoreError::ParentNotFound(parent_id.clone()));
            }
        } else if data
            .nodes
            .keys()
            .any(|(story_id, _)| story_id == &node.story_id)
        {
            return Err(StoreError::ParentNotFound(
                "a non-root node must name its parent".to_owned(),
            ));
        }

        if let Some(current_head) = data.branch_heads.get(&branch_key) {
            let expected_parent = node.parent_id.as_deref().unwrap_or("<none>");
            if current_head != expected_parent {
                return Err(StoreError::StaleBranchHead {
                    expected: current_head.clone(),
                    actual: expected_parent.to_owned(),
                });
            }
        }

        // All validation happens before any map is changed. This mirrors the
        // all-or-nothing transaction boundary required from the SQLite store.
        data.nodes.insert(node_key.clone(), node.clone());
        data.states.insert(node_key, state.clone());
        data.branch_heads.insert(branch_key, node.id.clone());
        Ok(())
    }

    fn load_state(&self, story_id: &str, branch_id: &str) -> Result<RuntimeState, StoreError> {
        let data = self.lock()?;
        let story_exists = data
            .nodes
            .keys()
            .any(|(stored_story_id, _)| stored_story_id == story_id);
        if !story_exists {
            return Err(StoreError::StoryNotFound(story_id.to_owned()));
        }

        let head = data
            .branch_heads
            .get(&(story_id.to_owned(), branch_id.to_owned()))
            .ok_or_else(|| StoreError::BranchNotFound {
                story_id: story_id.to_owned(),
                branch_id: branch_id.to_owned(),
            })?;
        data.states
            .get(&(story_id.to_owned(), head.clone()))
            .cloned()
            .ok_or(StoreError::StateMismatch("branch head has no state"))
    }
}

fn validate_materialized_state(
    node: &StoryNode,
    state: &RuntimeState,
) -> Result<(), StoreError> {
    if node.story_id != state.story_id {
        return Err(StoreError::StateMismatch("story_id"));
    }
    if node.id != state.current_node {
        return Err(StoreError::StateMismatch("current_node"));
    }
    if node.branch_id != state.current_branch {
        return Err(StoreError::StateMismatch("current_branch"));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use nana_domain::{RuntimeState, StateDelta, StoryNode};

    use super::{InMemoryStoryStore, StoreError, StoryStore};

    fn state(node: &str, branch: &str) -> RuntimeState {
        RuntimeState {
            story_id: "story_demo".to_owned(),
            current_node: node.to_owned(),
            current_branch: branch.to_owned(),
            world_flags: BTreeMap::new(),
            relationships: BTreeMap::new(),
            relationship_states: Vec::new(),
            promises: Vec::new(),
            knowledge: Vec::new(),
            items: Vec::new(),
            clocks: Vec::new(),
            checks: Vec::new(),
        }
    }

    fn node(id: &str, parent_id: Option<&str>, branch: &str) -> StoryNode {
        StoryNode {
            id: id.to_owned(),
            story_id: "story_demo".to_owned(),
            branch_id: branch.to_owned(),
            parent_id: parent_id.map(ToOwned::to_owned),
            action_id: format!("action_{id}"),
            user_input: String::new(),
            beats: Vec::new(),
            delta: StateDelta { ops: Vec::new() },
            state_hash: format!("hash_{id}"),
        }
    }

    #[test]
    fn appends_and_loads_the_branch_head() {
        let store = InMemoryStoryStore::new();
        store
            .append_node(&node("node_001", None, "branch_main"), &state("node_001", "branch_main"))
            .expect("root append");
        store
            .append_node(
                &node("node_002", Some("node_001"), "branch_main"),
                &state("node_002", "branch_main"),
            )
            .expect("child append");

        assert_eq!(
            store
                .load_state("story_demo", "branch_main")
                .expect("stored state")
                .current_node,
            "node_002"
        );
    }

    #[test]
    fn forks_from_an_old_node_without_moving_the_source_head() {
        let store = InMemoryStoryStore::new();
        store
            .append_node(&node("node_001", None, "branch_main"), &state("node_001", "branch_main"))
            .expect("root append");
        store
            .append_node(
                &node("node_002", Some("node_001"), "branch_main"),
                &state("node_002", "branch_main"),
            )
            .expect("main append");
        store
            .append_node(
                &node("node_003", Some("node_001"), "branch_rewind"),
                &state("node_003", "branch_rewind"),
            )
            .expect("fork append");

        assert_eq!(
            store
                .branch_head("story_demo", "branch_main")
                .expect("main head"),
            Some("node_002".to_owned())
        );
        assert_eq!(
            store
                .branch_head("story_demo", "branch_rewind")
                .expect("fork head"),
            Some("node_003".to_owned())
        );
    }

    #[test]
    fn rejects_a_stale_append_without_moving_the_head() {
        let store = InMemoryStoryStore::new();
        store
            .append_node(&node("node_001", None, "branch_main"), &state("node_001", "branch_main"))
            .expect("root append");
        store
            .append_node(
                &node("node_002", Some("node_001"), "branch_main"),
                &state("node_002", "branch_main"),
            )
            .expect("main append");

        assert_eq!(
            store.append_node(
                &node("node_stale", Some("node_001"), "branch_main"),
                &state("node_stale", "branch_main"),
            ),
            Err(StoreError::StaleBranchHead {
                expected: "node_002".to_owned(),
                actual: "node_001".to_owned(),
            })
        );
        assert_eq!(
            store
                .branch_head("story_demo", "branch_main")
                .expect("main head"),
            Some("node_002".to_owned())
        );
    }

    #[test]
    fn rejects_state_that_does_not_match_the_node() {
        let store = InMemoryStoryStore::new();
        assert_eq!(
            store.append_node(
                &node("node_001", None, "branch_main"),
                &state("different_node", "branch_main"),
            ),
            Err(StoreError::StateMismatch("current_node"))
        );
        assert_eq!(
            store
                .branch_head("story_demo", "branch_main")
                .expect("empty store"),
            None
        );
    }
}
