use std::collections::BTreeMap;

use super::{Entry, Manifest, RelativePath, canonical_json_bytes};

const MAX_LEAF_ENTRIES: usize = 64;
const BLAKE3_HEX_LENGTH: usize = 64;
const ROOT_PREFIX: &str = "";
#[cfg(test)]
const ENTRIES_ABOVE_LEAF_LIMIT: usize = MAX_LEAF_ENTRIES + 1;

#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ManifestMerkleNode {
    Branch {
        prefix: String,
        hash: String,
        children: BTreeMap<String, String>,
    },
    Leaf {
        prefix: String,
        hash: String,
        entries: BTreeMap<RelativePath, Entry>,
    },
}

impl ManifestMerkleNode {
    pub fn hash(&self) -> &str {
        match self {
            Self::Branch { hash, .. } | Self::Leaf { hash, .. } => hash,
        }
    }

    pub fn prefix(&self) -> &str {
        match self {
            Self::Branch { prefix, .. } | Self::Leaf { prefix, .. } => prefix,
        }
    }
}

#[derive(Clone, Debug)]
pub struct ManifestMerkleIndex {
    nodes: BTreeMap<String, ManifestMerkleNode>,
}

impl ManifestMerkleIndex {
    pub fn build(manifest: &Manifest) -> anyhow::Result<Self> {
        Self::build_at(ROOT_PREFIX, manifest.entries.clone())
    }

    fn build_at(prefix: &str, entries: BTreeMap<RelativePath, Entry>) -> anyhow::Result<Self> {
        let mut nodes = BTreeMap::new();
        build_index_node(prefix, entries, &mut nodes)?;
        Ok(Self { nodes })
    }

    pub fn node(&self, prefix: &str) -> anyhow::Result<Option<ManifestMerkleNode>> {
        validate_merkle_prefix(prefix)?;
        Ok(self.nodes.get(prefix).cloned())
    }

    pub fn root_hash(&self) -> &str {
        self.nodes
            .get(ROOT_PREFIX)
            .expect("Merkle indexes always include a root node")
            .hash()
    }
}

pub fn manifest_root_hash(manifest: &Manifest) -> anyhow::Result<String> {
    Ok(ManifestMerkleIndex::build(manifest)?.root_hash().to_owned())
}

#[cfg(test)]
pub(super) fn manifest_merkle_node(
    manifest: &Manifest,
    prefix: &str,
) -> anyhow::Result<ManifestMerkleNode> {
    validate_merkle_prefix(prefix)?;
    let entries = manifest
        .entries
        .iter()
        .filter(|(path, _)| path_hash(path).starts_with(prefix))
        .map(|(path, entry)| (path.clone(), entry.clone()))
        .collect();
    ManifestMerkleIndex::build_at(prefix, entries)?
        .node(prefix)?
        .ok_or_else(|| anyhow::anyhow!("Merkle node is missing for prefix {prefix:?}"))
}

fn build_index_node(
    prefix: &str,
    entries: BTreeMap<RelativePath, Entry>,
    nodes: &mut BTreeMap<String, ManifestMerkleNode>,
) -> anyhow::Result<ManifestMerkleNode> {
    let node = if entries.len() <= MAX_LEAF_ENTRIES || prefix.len() == BLAKE3_HEX_LENGTH {
        ManifestMerkleNode::Leaf {
            prefix: prefix.to_owned(),
            hash: blake3::hash(&canonical_json_bytes(&entries)?)
                .to_hex()
                .to_string(),
            entries,
        }
    } else {
        let mut groups: BTreeMap<String, BTreeMap<RelativePath, Entry>> = BTreeMap::new();
        for (path, entry) in entries {
            // Hash paths before branching so unrelated filenames distribute evenly
            // without exposing their common textual prefixes in the trie shape.
            let path_hash = path_hash(&path);
            let label = path_hash
                .chars()
                .nth(prefix.len())
                .ok_or_else(|| anyhow::anyhow!("Merkle prefix exceeds a path hash"))?
                .to_string();
            groups.entry(label).or_default().insert(path, entry);
        }
        let mut children = BTreeMap::new();
        for (label, child_entries) in groups {
            let child = build_index_node(&format!("{prefix}{label}"), child_entries, nodes)?;
            children.insert(label, child.hash().to_owned());
        }
        ManifestMerkleNode::Branch {
            prefix: prefix.to_owned(),
            hash: blake3::hash(&canonical_json_bytes(&children)?)
                .to_hex()
                .to_string(),
            children,
        }
    };
    nodes.insert(prefix.to_owned(), node.clone());
    Ok(node)
}

fn validate_merkle_prefix(prefix: &str) -> anyhow::Result<()> {
    if prefix.len() > BLAKE3_HEX_LENGTH
        || !prefix
            .chars()
            .all(|character| character.is_ascii_hexdigit())
    {
        anyhow::bail!("invalid manifest Merkle prefix");
    }
    Ok(())
}

fn path_hash(path: &str) -> String {
    blake3::hash(path.as_bytes()).to_hex().to_string()
}

#[cfg(test)]
mod tests {
    use super::{ENTRIES_ABOVE_LEAF_LIMIT, ManifestMerkleIndex, manifest_merkle_node};
    use crate::domain::{Entry, EntryKind, Manifest};

    #[test]
    fn cached_merkle_nodes_match_the_on_demand_tree() -> anyhow::Result<()> {
        let folder_id = uuid::Uuid::new_v4();
        let mut manifest = Manifest::empty(folder_id);
        for number in 0..ENTRIES_ABOVE_LEAF_LIMIT {
            let path = format!("file-{number}.txt");
            manifest.entries.insert(
                path.clone(),
                Entry {
                    path,
                    kind: EntryKind::File,
                    blob_hash: Some(format!("blob-{number}")),
                    size_bytes: Some(1),
                    modified_at: time::OffsetDateTime::UNIX_EPOCH,
                    clock: Default::default(),
                    author_device_id: "device".to_owned(),
                },
            );
        }

        let index = ManifestMerkleIndex::build(&manifest)?;
        let cached_root = index.node("")?.expect("root node");
        assert_eq!(cached_root, manifest_merkle_node(&manifest, "")?);
        Ok(())
    }
}