use crate::utils::{FastDashMap as DashMap, now_nanos, now_secs};
use foldhash::HashMap;
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

#[derive(Default)]
pub(crate) struct TrieNode {
    version: AtomicU64,
    last_accessed: AtomicU64,
    children: DashMap<String, Arc<TrieNode>>,
}

impl TrieNode {
    fn touch(&self, now: u64) {
        self.last_accessed.store(now, Ordering::Relaxed);
    }
}

#[derive(Clone)]
pub(crate) struct PrefixTrie {
    root: Arc<TrieNode>,
    global_version: Arc<AtomicU64>,
    min_pruned_version: Arc<AtomicU64>,
}

impl PrefixTrie {
    pub fn new() -> Self {
        Self {
            root: Arc::new(TrieNode::default()),
            global_version: Arc::new(AtomicU64::new(0)),
            min_pruned_version: Arc::new(AtomicU64::new(0)),
        }
    }

    #[inline]
    pub fn invalidate(&self, tag: &str) -> u64 {
        self.global_version.fetch_add(1, Ordering::SeqCst);
        let now_s = now_secs();
        let current = self.traverse_and_touch(tag, now_s);

        let now_n = now_nanos();
        let prev = current
            .version
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |v| {
                Some(v.max(now_n).max(v + 1))
            })
            .expect(
                "fetch_update on atomic u64 should never fail when update function returns Some",
            );

        prev.max(now_n).max(prev + 1)
    }

    #[inline]
    pub fn set_min_version(&self, tag: &str, version: u64) {
        self.global_version.fetch_add(1, Ordering::SeqCst);
        let now_s = now_secs();
        let current = self.traverse_and_touch(tag, now_s);
        current.version.fetch_max(version, Ordering::SeqCst);
    }

    pub fn get_path_versions<S: AsRef<str>>(&self, parts: &[S], now: u64) -> SmallVec<[u64; 8]> {
        let mut versions = SmallVec::with_capacity(parts.len() + 1);
        let mut current = Arc::clone(&self.root);
        current.touch(now);
        versions.push(current.version.load(Ordering::SeqCst));

        for part in parts {
            let next = match current.children.get(part.as_ref()) {
                Some(n) => Arc::clone(n.value()),
                None => break,
            };
            current = next;
            current.touch(now);
            versions.push(current.version.load(Ordering::SeqCst));
        }

        while versions.len() <= parts.len() {
            versions.push(0);
        }
        versions
    }

    #[inline]
    pub fn check_and_catch_up<S: AsRef<str>>(
        &self,
        parts: &[S],
        snapshot_versions: &[u64],
        now: u64,
        create: bool,
    ) -> bool {
        let mut current = Arc::clone(&self.root);
        current.touch(now);

        current
            .version
            .fetch_max(snapshot_versions[0], Ordering::SeqCst);

        let mut is_valid = current.version.load(Ordering::SeqCst) <= snapshot_versions[0];

        for (i, part) in parts.iter().enumerate() {
            let next_node = if !create && is_valid {
                match current.children.get(part.as_ref()) {
                    Some(n) => Arc::clone(n.value()),
                    None => {
                        let min_p = self.min_pruned_version.load(Ordering::SeqCst);
                        for v in &snapshot_versions[i + 1..] {
                            if *v > 0 && *v < min_p {
                                return false;
                            }
                        }
                        return true;
                    }
                }
            } else {
                self.get_or_create_child(&current, part.as_ref())
            };

            current = next_node;
            current.touch(now);
            current
                .version
                .fetch_max(snapshot_versions[i + 1], Ordering::SeqCst);

            if is_valid && current.version.load(Ordering::SeqCst) > snapshot_versions[i + 1] {
                is_valid = false;
            }
        }
        is_valid
    }

    pub fn get_tag_version(&self, tag: &str) -> u64 {
        let parts: SmallVec<[&str; 8]> = tag.split(':').filter(|s| !s.is_empty()).collect();
        let versions = self.get_path_versions(&parts, now_secs());
        *versions.last().unwrap_or(&0)
    }

    #[inline]
    pub fn get_global_version(&self) -> u64 {
        self.global_version.load(Ordering::SeqCst)
    }

    pub fn clear(&self) {
        self.root.children.clear();
        self.root.version.store(0, Ordering::SeqCst);
        self.global_version.fetch_add(1, Ordering::SeqCst);
        self.root.touch(now_secs());
    }

    pub fn prune(&self, max_age_secs: u64) {
        let now = now_secs();
        self.root.children.retain(|_, node| {
            !Self::should_prune(node, now, max_age_secs, &self.min_pruned_version)
        });
    }

    #[cfg(test)]
    pub fn force_prune_all(&self) {
        self.root
            .children
            .retain(|_, node| !Self::should_prune(node, u64::MAX, 0, &self.min_pruned_version));
    }

    fn should_prune(
        node: &Arc<TrieNode>,
        now: u64,
        max_age_secs: u64,
        min_pruned: &AtomicU64,
    ) -> bool {
        node.children
            .retain(|_, child| !Self::should_prune(child, now, max_age_secs, min_pruned));

        let last = node.last_accessed.load(Ordering::Relaxed);
        let age = now.saturating_sub(last);
        let version = node.version.load(Ordering::Relaxed);

        if node.children.is_empty() && age > max_age_secs {
            if version > 0 {
                min_pruned.fetch_max(version, Ordering::SeqCst);
            }
            true
        } else {
            false
        }
    }

    #[inline]
    fn traverse_and_touch(&self, tag: &str, now: u64) -> Arc<TrieNode> {
        let mut current = Arc::clone(&self.root);
        current.touch(now);
        for part in tag.split(':') {
            if !part.is_empty() {
                let next = self.get_or_create_child(&current, part);
                current = next;
                current.touch(now);
            }
        }
        current
    }

    fn get_or_create_child(&self, parent: &Arc<TrieNode>, part: &str) -> Arc<TrieNode> {
        if let Some(n) = parent.children.get(part) {
            Arc::clone(n.value())
        } else {
            let entry = parent
                .children
                .entry(part.to_string())
                .or_insert_with(|| Arc::new(TrieNode::default()));
            Arc::clone(entry.value())
        }
    }
}

#[derive(Serialize, Deserialize, Clone)]
pub(crate) struct DepSnapshot {
    pub parts: SmallVec<[String; 8]>,
    pub path_versions: SmallVec<[u64; 8]>,
}

#[inline]
pub(crate) fn validate_dependencies(
    trie: &PrefixTrie,
    deps: &HashMap<String, DepSnapshot>,
    now: u64,
) -> bool {
    if deps.is_empty() {
        return true;
    }
    for snapshot in deps.values() {
        if !trie.check_and_catch_up(&snapshot.parts, &snapshot.path_versions, now, true) {
            return false;
        }
    }
    true
}

#[inline]
pub(crate) fn build_dependency_snapshots(
    trie: &PrefixTrie,
    dependencies: Vec<String>,
    now: u64,
) -> Arc<HashMap<String, DepSnapshot>> {
    let mut snapshots = HashMap::with_capacity_and_hasher(dependencies.len(), Default::default());
    for tag in dependencies {
        let parts: SmallVec<[String; 8]> = tag
            .split(':')
            .filter(|s| !s.is_empty())
            .map(|s| s.to_string())
            .collect();
        let path_versions = trie.get_path_versions(&parts, now);
        snapshots.insert(
            tag,
            DepSnapshot {
                parts,
                path_versions,
            },
        );
    }
    Arc::new(snapshots)
}

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

    #[test]
    fn test_trie_basic() {
        let trie = PrefixTrie::new();
        let parts = vec!["user".to_string(), "1".to_string()];
        let now = now_secs();
        assert_eq!(
            trie.get_path_versions(&parts, now).into_vec(),
            vec![0, 0, 0]
        );

        trie.invalidate("user:1");
        let v1 = trie.get_path_versions(&parts, now);
        assert!(v1[2] > 0);

        trie.invalidate("user:1");
        let v2 = trie.get_path_versions(&parts, now);
        assert!(v2[2] > v1[2]);
    }

    #[test]
    fn test_trie_prefix_invalidation() {
        let trie = PrefixTrie::new();
        let tag_parts = vec![
            "org".to_string(),
            "1".to_string(),
            "user".to_string(),
            "1".to_string(),
        ];

        let now = now_secs();
        let v0 = trie.get_path_versions(&tag_parts, now);
        assert!(trie.check_and_catch_up(&tag_parts, &v0, now, false));

        trie.invalidate("org:1:user:1");
        let v1 = trie.get_path_versions(&tag_parts, now);
        assert!(!trie.check_and_catch_up(&tag_parts, &v0, now, false));
        assert!(trie.check_and_catch_up(&tag_parts, &v1, now, false));

        trie.invalidate("org:1");
        assert!(!trie.check_and_catch_up(&tag_parts, &v1, now, false));
    }

    #[test]
    fn test_trie_deep_hierarchy() {
        let trie = PrefixTrie::new();
        let deep_tag = "a:b:c:d:e:f:g:h:i:j";
        let parts: Vec<String> = deep_tag.split(':').map(|s| s.to_string()).collect();

        let now = now_secs();
        let path_v0 = trie.get_path_versions(&parts, now);
        trie.invalidate(deep_tag);
        let path_v1 = trie.get_path_versions(&parts, now);

        assert!(!trie.check_and_catch_up(&parts, &path_v0, now, false));
        assert!(trie.check_and_catch_up(&parts, &path_v1, now, false));

        trie.invalidate("a:b:c");
        assert!(!trie.check_and_catch_up(&parts, &path_v1, now, false));
    }

    #[test]
    fn test_trie_clear() {
        let trie = PrefixTrie::new();
        trie.invalidate("user:1");
        trie.invalidate("user:2");

        let parts = vec!["user".to_string(), "1".to_string()];
        let now = now_secs();
        let v1 = trie.get_path_versions(&parts, now);
        assert!(v1[2] > 0);

        trie.clear();
        let v2 = trie.get_path_versions(&parts, now);
        assert_eq!(v2.into_vec(), vec![0, 0, 0]);
    }

    #[test]
    fn test_validate_dependencies_empty() {
        let trie = PrefixTrie::new();
        let deps = HashMap::default();
        assert!(validate_dependencies(&trie, &deps, now_secs()));
    }

    #[test]
    fn test_validate_dependencies_valid() {
        let trie = PrefixTrie::new();
        let now = now_secs();
        let deps = build_dependency_snapshots(&trie, vec!["user:1".to_string()], now);
        assert!(validate_dependencies(&trie, &deps, now));
    }

    #[test]
    fn test_validate_dependencies_invalidated() {
        let trie = PrefixTrie::new();
        let now = now_secs();
        let deps = build_dependency_snapshots(&trie, vec!["user:1".to_string()], now);
        trie.invalidate("user:1");
        assert!(!validate_dependencies(&trie, &deps, now));
    }

    #[test]
    fn test_validate_dependencies_parent_invalidated() {
        let trie = PrefixTrie::new();
        let now = now_secs();
        let deps = build_dependency_snapshots(&trie, vec!["org:1:user:5".to_string()], now);
        trie.invalidate("org:1");
        assert!(!validate_dependencies(&trie, &deps, now));
    }

    #[test]
    fn test_build_dependency_snapshots() {
        let trie = PrefixTrie::new();
        trie.invalidate("user:1");

        let now = now_secs();
        let snapshots = build_dependency_snapshots(
            &trie,
            vec!["user:1".to_string(), "user:2".to_string()],
            now,
        );

        assert_eq!(snapshots.len(), 2);
        assert_eq!(snapshots["user:1"].parts.as_slice(), &["user", "1"]);
        assert!(snapshots["user:1"].path_versions[2] > 0);
        assert_eq!(snapshots["user:2"].path_versions.as_slice(), &[0, 0, 0]);
    }

    #[test]
    fn test_cold_start_simulation() {
        let trie_old = PrefixTrie::new();
        trie_old.invalidate("user:1");
        let now = now_secs();
        let v_old = trie_old.get_path_versions(&["user", "1"], now)[2];

        let trie_new = PrefixTrie::new();
        let snapshot_v = vec![0, 0, v_old];
        assert!(trie_new.check_and_catch_up(&["user", "1"], &snapshot_v, now, false));

        trie_new.invalidate("user:1");
        let v_new = trie_new.get_path_versions(&["user", "1"], now)[2];

        assert!(v_new >= v_old);
        assert!(!trie_new.check_and_catch_up(&["user", "1"], &snapshot_v, now, false));
    }

    #[test]
    fn test_catch_up() {
        let trie = PrefixTrie::new();
        let parts = vec!["org", "acme", "user", "42"];
        let snapshot_versions = vec![0, 0, 100, 0, 500];

        let now = now_secs();
        trie.check_and_catch_up(&parts, &snapshot_versions, now, true);

        let current_versions = trie.get_path_versions(&parts, now);
        assert_eq!(current_versions[2], 100);
        assert_eq!(current_versions[4], 500);
        assert_eq!(current_versions[0], 0);
    }

    #[test]
    fn test_pruning() {
        let trie = PrefixTrie::new();

        trie.traverse_and_touch("user:1", now_secs());

        trie.prune(100);
        assert_eq!(trie.root.children.len(), 1);

        std::thread::sleep(std::time::Duration::from_secs(2));
        trie.prune(0);

        assert_eq!(trie.root.children.len(), 0);
    }

    #[test]
    fn test_prune_barrier_validation() {
        let trie = PrefixTrie::new();
        let tag = "user:1";

        let version = trie.invalidate(tag);
        assert!(version > 0);

        let snapshot_v = vec![0, 0, version];
        assert!(trie.check_and_catch_up(&["user", "1"], &snapshot_v, now_secs(), false));

        trie.force_prune_all();
        assert!(trie.check_and_catch_up(&["user", "1"], &snapshot_v, now_secs(), false));

        let old_snapshot_v = vec![0, 0, version - 1];
        assert!(!trie.check_and_catch_up(&["user", "1"], &old_snapshot_v, now_secs(), false));

        assert_eq!(trie.root.children.len(), 0, "Nodes should be pruned if old");
    }

    #[test]
    fn test_hlc_ratchet() {
        let trie = PrefixTrie::new();
        let tag = "test:hlc";

        let v1 = trie.invalidate(tag);
        assert!(v1 > 0);

        let future_ver = v1 + 1_000_000_000;
        trie.set_min_version(tag, future_ver);

        let current_ver = trie.get_tag_version(tag);
        assert!(current_ver >= future_ver);

        let v2 = trie.invalidate(tag);
        assert!(v2 > future_ver);
    }

    #[test]
    fn test_hlc_fast_forward() {
        let trie = PrefixTrie::new();
        let tag = "test:fast_forward";

        let _v1 = trie.invalidate(tag);

        let far_future = u64::MAX - 1000;
        trie.set_min_version(tag, far_future);

        assert_eq!(trie.get_tag_version(tag), far_future);

        let v_new = trie.invalidate(tag);

        assert!(v_new > far_future);
    }

    #[test]
    fn test_prune_version_leak_fixed() {
        let trie = PrefixTrie::new();
        let tag = "leak:me";

        trie.invalidate(tag);
        trie.force_prune_all();

        let parts = vec!["leak", "me"];
        let path_versions = trie.get_path_versions(&parts, now_secs());

        assert_eq!(
            path_versions[2], 0,
            "Fix verification: Node with version > 0 was correctly pruned."
        );
        assert!(
            !trie.root.children.contains_key("leak"),
            "Root child 'leak' should be gone if it had no other children"
        );
    }

    #[test]
    fn test_prune_barrier_security() {
        let trie = PrefixTrie::new();
        let tag = "secure:me";

        let v1 = trie.invalidate(tag);
        assert!(v1 > 0);

        let snapshot_v = vec![0, 0, v1 - 1];
        let parts = vec!["secure", "me"];
        assert!(!trie.check_and_catch_up(&parts, &snapshot_v, now_secs(), false));

        trie.force_prune_all();

        let snapshot_v_new = vec![0, 0, v1 + 1];
        assert!(trie.check_and_catch_up(&parts, &snapshot_v_new, now_secs(), false));

        let snapshot_v_zero = vec![0, 0, 0];
        assert!(trie.check_and_catch_up(&parts, &snapshot_v_zero, now_secs(), false));
    }
}
