use serde::{de::DeserializeOwned, Serialize};
use serde_json::Value;

use crate::error::{Result, SdkError};

/// Trait for types that can be used as module state.
///
/// Any struct that implements `Serialize + DeserializeOwned + Clone + Send + Sync`
/// can serve as module state. The SDK uses serde to convert between the typed
/// state and JSON, and diffs JSON snapshots to detect which paths changed.
///
/// # Example
///
/// ```rust
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Clone, Default, Serialize, Deserialize)]
/// struct CounterState {
///     count: i32,
///     label: String,
/// }
/// ```
pub trait State: Serialize + DeserializeOwned + Clone + Send + Sync + 'static {}

/// Blanket implementation: any type meeting the bounds is automatically `State`.
impl<T> State for T where T: Serialize + DeserializeOwned + Clone + Send + Sync + 'static {}

/// Holds the typed state alongside its JSON representation for efficient diffing.
pub(crate) struct StateContainer<S: State> {
    /// The typed state value — what users interact with.
    value: S,
    /// JSON snapshot taken *before* the most recent handler ran.
    /// Used to compute changed paths after mutation.
    snapshot: Value,
}

impl<S: State> StateContainer<S> {
    /// Create a new container, serializing the initial state to JSON.
    pub fn new(initial: S) -> Result<Self> {
        let snapshot =
            serde_json::to_value(&initial).map_err(|e| SdkError::StateSerde(e.to_string()))?;
        Ok(Self {
            value: initial,
            snapshot,
        })
    }

    /// Borrow the current state immutably.
    pub fn get(&self) -> &S {
        &self.value
    }

    /// Borrow the current state mutably (for action handlers).
    pub fn get_mut(&mut self) -> &mut S {
        &mut self.value
    }

    /// Take a snapshot of the current state (call *before* a handler mutates it).
    pub fn take_snapshot(&mut self) -> Result<()> {
        self.snapshot =
            serde_json::to_value(&self.value).map_err(|e| SdkError::StateSerde(e.to_string()))?;
        Ok(())
    }

    /// Compare the current state against the last snapshot and return
    /// the list of dot-separated paths that changed.
    ///
    /// This is the core mechanism: handlers just do `state.count += 1`,
    /// and we diff before/after to find `["count"]`.
    pub fn changed_paths(&self) -> Result<Vec<String>> {
        let current =
            serde_json::to_value(&self.value).map_err(|e| SdkError::StateSerde(e.to_string()))?;
        let mut paths = Vec::new();
        diff_json("", &self.snapshot, &current, &mut paths);
        Ok(paths)
    }

    /// Get the current state as a JSON value (for the engine).
    pub fn to_json(&self) -> Result<Value> {
        serde_json::to_value(&self.value).map_err(|e| SdkError::StateSerde(e.to_string()))
    }

    /// Build a JSON patch object containing only the changed key-value pairs.
    pub fn diff_patch(&self) -> Result<Value> {
        let current =
            serde_json::to_value(&self.value).map_err(|e| SdkError::StateSerde(e.to_string()))?;
        let mut patch = serde_json::Map::new();
        collect_changed_values("", &self.snapshot, &current, &mut patch);
        Ok(Value::Object(patch))
    }
}

/// Recursively compare two JSON values and collect the dot-separated paths
/// of fields that differ.
fn diff_json(prefix: &str, old: &Value, new: &Value, paths: &mut Vec<String>) {
    match (old, new) {
        (Value::Object(old_map), Value::Object(new_map)) => {
            // Check keys present in old
            for (key, old_val) in old_map {
                let path = if prefix.is_empty() {
                    key.clone()
                } else {
                    format!("{prefix}.{key}")
                };
                match new_map.get(key) {
                    Some(new_val) => diff_json(&path, old_val, new_val, paths),
                    None => paths.push(path), // key was deleted
                }
            }
            // Check keys only in new
            for key in new_map.keys() {
                if !old_map.contains_key(key) {
                    let path = if prefix.is_empty() {
                        key.clone()
                    } else {
                        format!("{prefix}.{key}")
                    };
                    paths.push(path);
                }
            }
        }
        (Value::Array(old_arr), Value::Array(new_arr)) => {
            let max_len = old_arr.len().max(new_arr.len());
            for i in 0..max_len {
                let path = if prefix.is_empty() {
                    i.to_string()
                } else {
                    format!("{prefix}.{i}")
                };
                match (old_arr.get(i), new_arr.get(i)) {
                    (Some(old_val), Some(new_val)) => diff_json(&path, old_val, new_val, paths),
                    _ => paths.push(path), // added or removed element
                }
            }
        }
        (old_val, new_val) => {
            if old_val != new_val && !prefix.is_empty() {
                paths.push(prefix.to_string());
            }
        }
    }
}

/// Collect the top-level changed key-value pairs into a JSON patch object.
fn collect_changed_values(
    prefix: &str,
    old: &Value,
    new: &Value,
    patch: &mut serde_json::Map<String, Value>,
) {
    match (old, new) {
        (Value::Object(old_map), Value::Object(new_map)) => {
            for (key, old_val) in old_map {
                let path = if prefix.is_empty() {
                    key.clone()
                } else {
                    format!("{prefix}.{key}")
                };
                match new_map.get(key) {
                    Some(new_val) if old_val != new_val => {
                        // For top-level keys, include the whole new subtree
                        if prefix.is_empty() {
                            patch.insert(key.clone(), new_val.clone());
                        } else {
                            collect_changed_values(&path, old_val, new_val, patch);
                        }
                    }
                    None => {
                        patch.insert(path, Value::Null);
                    }
                    _ => {}
                }
            }
            for (key, new_val) in new_map {
                if !old_map.contains_key(key) {
                    if prefix.is_empty() {
                        patch.insert(key.clone(), new_val.clone());
                    } else {
                        patch.insert(format!("{prefix}.{key}"), new_val.clone());
                    }
                }
            }
        }
        (_, new_val) => {
            if old != new_val && !prefix.is_empty() {
                patch.insert(prefix.to_string(), new_val.clone());
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde::{Deserialize, Serialize};
    use serde_json::json;

    #[derive(Clone, Default, Serialize, Deserialize, Debug, PartialEq)]
    struct TestState {
        count: i32,
        name: String,
        items: Vec<String>,
    }

    #[test]
    fn test_diff_no_change() {
        let container = StateContainer::new(TestState {
            count: 0,
            name: "Alice".into(),
            items: vec![],
        })
        .unwrap();

        let paths = container.changed_paths().unwrap();
        assert!(paths.is_empty());
    }

    #[test]
    fn test_diff_scalar_change() {
        let mut container = StateContainer::new(TestState {
            count: 0,
            name: "Alice".into(),
            items: vec![],
        })
        .unwrap();

        container.take_snapshot().unwrap();
        container.get_mut().count = 42;

        let paths = container.changed_paths().unwrap();
        assert_eq!(paths, vec!["count"]);
    }

    #[test]
    fn test_diff_multiple_changes() {
        let mut container = StateContainer::new(TestState {
            count: 0,
            name: "Alice".into(),
            items: vec![],
        })
        .unwrap();

        container.take_snapshot().unwrap();
        container.get_mut().count = 10;
        container.get_mut().name = "Bob".into();

        let mut paths = container.changed_paths().unwrap();
        paths.sort();
        assert_eq!(paths, vec!["count", "name"]);
    }

    #[test]
    fn test_diff_array_change() {
        let mut container = StateContainer::new(TestState {
            count: 0,
            name: "Alice".into(),
            items: vec!["a".into()],
        })
        .unwrap();

        container.take_snapshot().unwrap();
        container.get_mut().items.push("b".into());

        let paths = container.changed_paths().unwrap();
        assert!(paths.contains(&"items.1".to_string()));
    }

    #[test]
    fn test_diff_nested_struct() {
        #[derive(Clone, Default, Serialize, Deserialize)]
        struct Nested {
            user: User,
            count: i32,
        }

        #[derive(Clone, Default, Serialize, Deserialize)]
        struct User {
            name: String,
            age: i32,
        }

        let mut container = StateContainer::new(Nested {
            user: User {
                name: "Alice".into(),
                age: 30,
            },
            count: 0,
        })
        .unwrap();

        container.take_snapshot().unwrap();
        container.get_mut().user.age = 31;

        let paths = container.changed_paths().unwrap();
        assert_eq!(paths, vec!["user.age"]);
    }

    #[test]
    fn test_diff_json_helper() {
        let old = json!({"a": 1, "b": {"c": 2, "d": 3}});
        let new = json!({"a": 1, "b": {"c": 99, "d": 3}, "e": true});

        let mut paths = Vec::new();
        diff_json("", &old, &new, &mut paths);

        assert!(paths.contains(&"b.c".to_string()));
        assert!(paths.contains(&"e".to_string()));
        assert!(!paths.contains(&"a".to_string()));
        assert!(!paths.contains(&"b.d".to_string()));
    }

    #[test]
    fn test_diff_patch_output() {
        let mut container = StateContainer::new(TestState {
            count: 0,
            name: "Alice".into(),
            items: vec![],
        })
        .unwrap();

        container.take_snapshot().unwrap();
        container.get_mut().count = 5;

        let patch = container.diff_patch().unwrap();
        assert_eq!(patch, json!({"count": 5}));
    }

    #[test]
    fn test_to_json() {
        let container = StateContainer::new(TestState {
            count: 42,
            name: "Bob".into(),
            items: vec!["x".into()],
        })
        .unwrap();

        let json = container.to_json().unwrap();
        assert_eq!(json["count"], 42);
        assert_eq!(json["name"], "Bob");
    }
}
