lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! Branch registry for tracking and managing branches.
//!
//! The registry maintains metadata about all branches in a dataset,
//! including the current branch and branch relationships.

use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use alloc::vec::Vec;

use super::types::{Branch, BranchError};

// ═══════════════════════════════════════════════════════════════════════════════
// BRANCH REGISTRY
// ═══════════════════════════════════════════════════════════════════════════════

/// Registry of all branches in a dataset.
///
/// The registry stores branch metadata and tracks the current branch.
/// It can be serialized to/from JSON for persistent storage in dataset properties.
#[derive(Debug, Clone)]
pub struct BranchRegistry {
    /// All branches indexed by name.
    branches: BTreeMap<String, Branch>,
    /// Currently checked out branch.
    current_branch: String,
    /// Next GUID to assign to new branches.
    next_guid: u64,
}

impl BranchRegistry {
    /// Create a new registry with an initial "main" branch.
    pub fn new(initial_txg: u64, created: u64) -> Self {
        let main = Branch::main(1, initial_txg, created);
        let mut branches = BTreeMap::new();
        branches.insert("main".to_string(), main);

        Self {
            branches,
            current_branch: "main".to_string(),
            next_guid: 2,
        }
    }

    /// Create an empty registry (for deserialization).
    pub fn empty() -> Self {
        Self {
            branches: BTreeMap::new(),
            current_branch: String::new(),
            next_guid: 1,
        }
    }

    /// Get the current branch name.
    pub fn current(&self) -> &str {
        &self.current_branch
    }

    /// Get the current branch.
    pub fn current_branch(&self) -> Option<&Branch> {
        self.branches.get(&self.current_branch)
    }

    /// Get the current branch mutably.
    pub fn current_branch_mut(&mut self) -> Option<&mut Branch> {
        self.branches.get_mut(&self.current_branch)
    }

    /// Get a branch by name.
    pub fn get(&self, name: &str) -> Option<&Branch> {
        self.branches.get(name)
    }

    /// Get a branch mutably by name.
    pub fn get_mut(&mut self, name: &str) -> Option<&mut Branch> {
        self.branches.get_mut(name)
    }

    /// Check if a branch exists.
    pub fn contains(&self, name: &str) -> bool {
        self.branches.contains_key(name)
    }

    /// Get the number of branches.
    pub fn len(&self) -> usize {
        self.branches.len()
    }

    /// Check if the registry is empty.
    pub fn is_empty(&self) -> bool {
        self.branches.is_empty()
    }

    /// Iterate over all branches.
    pub fn iter(&self) -> impl Iterator<Item = &Branch> {
        self.branches.values()
    }

    /// Get all branch names.
    pub fn names(&self) -> impl Iterator<Item = &String> {
        self.branches.keys()
    }

    /// Get the default branch.
    pub fn default_branch(&self) -> Option<&Branch> {
        self.branches.values().find(|b| b.is_default)
    }

    /// Get the default branch name.
    pub fn default_branch_name(&self) -> Option<&str> {
        self.default_branch().map(|b| b.name.as_str())
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // BRANCH CREATION
    // ═══════════════════════════════════════════════════════════════════════════

    /// Validate a branch name.
    pub fn validate_name(name: &str) -> Result<(), BranchError> {
        if name.is_empty() {
            return Err(BranchError::InvalidBranchName(
                "branch name cannot be empty".into(),
            ));
        }

        if name.len() > 255 {
            return Err(BranchError::InvalidBranchName(
                "branch name too long (max 255 chars)".into(),
            ));
        }

        // Check for invalid characters
        let invalid_chars = ['/', '\\', ':', '*', '?', '"', '<', '>', '|', '\0', ' '];
        for c in invalid_chars {
            if name.contains(c) {
                return Err(BranchError::InvalidBranchName(alloc::format!(
                    "branch name contains invalid character '{}'",
                    if c == '\0' { '0' } else { c }
                )));
            }
        }

        // Cannot start with a dot or dash
        if name.starts_with('.') || name.starts_with('-') {
            return Err(BranchError::InvalidBranchName(
                "branch name cannot start with '.' or '-'".into(),
            ));
        }

        // Cannot end with a dot
        if name.ends_with('.') {
            return Err(BranchError::InvalidBranchName(
                "branch name cannot end with '.'".into(),
            ));
        }

        // Cannot contain ".." or "@{"
        if name.contains("..") {
            return Err(BranchError::InvalidBranchName(
                "branch name cannot contain '..'".into(),
            ));
        }

        if name.contains("@{") {
            return Err(BranchError::InvalidBranchName(
                "branch name cannot contain '@{'".into(),
            ));
        }

        Ok(())
    }

    /// Create a new branch from the current branch.
    ///
    /// Returns the GUID assigned to the new branch.
    pub fn create_branch(&mut self, name: &str, created: u64) -> Result<u64, BranchError> {
        // Validate name
        Self::validate_name(name)?;

        // Check if branch exists
        if self.branches.contains_key(name) {
            return Err(BranchError::BranchExists(name.to_string()));
        }

        // Get current branch info
        let (parent_name, fork_txg) = {
            let current = self
                .current_branch()
                .ok_or_else(|| BranchError::Internal("no current branch".into()))?;
            (current.name.clone(), current.head_txg)
        };

        // Allocate GUID
        let guid = self.next_guid;
        self.next_guid += 1;

        // Create the branch
        let branch = Branch::new(name.to_string(), guid, Some(parent_name), fork_txg, created);

        self.branches.insert(name.to_string(), branch);

        Ok(guid)
    }

    /// Create a branch from a specific parent branch.
    pub fn create_branch_from(
        &mut self,
        name: &str,
        parent: &str,
        created: u64,
    ) -> Result<u64, BranchError> {
        // Validate name
        Self::validate_name(name)?;

        // Check if branch exists
        if self.branches.contains_key(name) {
            return Err(BranchError::BranchExists(name.to_string()));
        }

        // Get parent branch info
        let fork_txg = {
            let parent_branch = self
                .branches
                .get(parent)
                .ok_or_else(|| BranchError::BranchNotFound(parent.to_string()))?;
            parent_branch.head_txg
        };

        // Allocate GUID
        let guid = self.next_guid;
        self.next_guid += 1;

        // Create the branch
        let branch = Branch::new(
            name.to_string(),
            guid,
            Some(parent.to_string()),
            fork_txg,
            created,
        );

        self.branches.insert(name.to_string(), branch);

        Ok(guid)
    }

    /// Create a branch at a specific TXG.
    pub fn create_branch_at_txg(
        &mut self,
        name: &str,
        parent: &str,
        txg: u64,
        created: u64,
    ) -> Result<u64, BranchError> {
        // Validate name
        Self::validate_name(name)?;

        // Check if branch exists
        if self.branches.contains_key(name) {
            return Err(BranchError::BranchExists(name.to_string()));
        }

        // Verify parent exists
        if !self.branches.contains_key(parent) {
            return Err(BranchError::BranchNotFound(parent.to_string()));
        }

        // Allocate GUID
        let guid = self.next_guid;
        self.next_guid += 1;

        // Create the branch
        let branch = Branch::new(
            name.to_string(),
            guid,
            Some(parent.to_string()),
            txg,
            created,
        );

        self.branches.insert(name.to_string(), branch);

        Ok(guid)
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // BRANCH SWITCHING
    // ═══════════════════════════════════════════════════════════════════════════

    /// Switch to a different branch.
    pub fn checkout(&mut self, name: &str) -> Result<&Branch, BranchError> {
        // Check if already on this branch
        if self.current_branch == name {
            return Err(BranchError::AlreadyOnBranch(name.to_string()));
        }

        // Check if branch exists
        if !self.branches.contains_key(name) {
            return Err(BranchError::BranchNotFound(name.to_string()));
        }

        self.current_branch = name.to_string();

        Ok(self.branches.get(name).unwrap())
    }

    /// Switch to a branch, creating it if it doesn't exist.
    pub fn checkout_or_create(&mut self, name: &str, created: u64) -> Result<&Branch, BranchError> {
        if !self.branches.contains_key(name) {
            self.create_branch(name, created)?;
        }

        self.current_branch = name.to_string();
        Ok(self.branches.get(name).unwrap())
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // BRANCH DELETION
    // ═══════════════════════════════════════════════════════════════════════════

    /// Delete a branch.
    pub fn delete_branch(&mut self, name: &str, force: bool) -> Result<Branch, BranchError> {
        // Cannot delete current branch
        if self.current_branch == name {
            return Err(BranchError::CannotDeleteCurrent(name.to_string()));
        }

        // Get the branch
        let branch = self
            .branches
            .get(name)
            .ok_or_else(|| BranchError::BranchNotFound(name.to_string()))?;

        // Cannot delete default branch
        if branch.is_default {
            return Err(BranchError::CannotDeleteDefault(name.to_string()));
        }

        // Check for unmerged changes (unless force)
        if !force && branch.txgs_since_fork() > 0 {
            return Err(BranchError::UnmergedChanges(name.to_string()));
        }

        Ok(self.branches.remove(name).unwrap())
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // BRANCH UPDATES
    // ═══════════════════════════════════════════════════════════════════════════

    /// Update the head TXG of a branch.
    pub fn update_head(&mut self, name: &str, txg: u64) -> Result<(), BranchError> {
        let branch = self
            .branches
            .get_mut(name)
            .ok_or_else(|| BranchError::BranchNotFound(name.to_string()))?;

        branch.head_txg = txg;
        Ok(())
    }

    /// Update the head TXG of the current branch.
    pub fn update_current_head(&mut self, txg: u64) -> Result<(), BranchError> {
        let name = self.current_branch.clone();
        self.update_head(&name, txg)
    }

    /// Set the default branch.
    pub fn set_default(&mut self, name: &str) -> Result<(), BranchError> {
        // Check if branch exists
        if !self.branches.contains_key(name) {
            return Err(BranchError::BranchNotFound(name.to_string()));
        }

        // Clear old default
        for branch in self.branches.values_mut() {
            branch.is_default = false;
        }

        // Set new default
        self.branches.get_mut(name).unwrap().is_default = true;

        Ok(())
    }

    /// Rename a branch.
    pub fn rename(&mut self, old_name: &str, new_name: &str) -> Result<(), BranchError> {
        // Validate new name
        Self::validate_name(new_name)?;

        // Check if old exists
        if !self.branches.contains_key(old_name) {
            return Err(BranchError::BranchNotFound(old_name.to_string()));
        }

        // Check if new already exists
        if self.branches.contains_key(new_name) {
            return Err(BranchError::BranchExists(new_name.to_string()));
        }

        // Remove and re-insert with new name
        let mut branch = self.branches.remove(old_name).unwrap();
        branch.name = new_name.to_string();
        self.branches.insert(new_name.to_string(), branch);

        // Update current branch reference if needed
        if self.current_branch == old_name {
            self.current_branch = new_name.to_string();
        }

        // Update parent references in child branches
        for branch in self.branches.values_mut() {
            if branch.parent.as_deref() == Some(old_name) {
                branch.parent = Some(new_name.to_string());
            }
        }

        Ok(())
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // QUERIES
    // ═══════════════════════════════════════════════════════════════════════════

    /// Get all children of a branch.
    pub fn children(&self, name: &str) -> Vec<&Branch> {
        self.branches
            .values()
            .filter(|b| b.parent.as_deref() == Some(name))
            .collect()
    }

    /// Get the ancestry chain from a branch back to root.
    pub fn ancestry(&self, name: &str) -> Vec<&Branch> {
        let mut result = Vec::new();
        let mut current = name;

        while let Some(branch) = self.branches.get(current) {
            result.push(branch);
            match &branch.parent {
                Some(parent) => current = parent,
                None => break,
            }
        }

        result
    }

    /// Find the common ancestor of two branches.
    pub fn common_ancestor(&self, branch_a: &str, branch_b: &str) -> Option<&Branch> {
        let ancestry_a: Vec<_> = self.ancestry(branch_a);
        let ancestry_b: Vec<_> = self.ancestry(branch_b);

        // Find first common branch in ancestry
        for ancestor_a in &ancestry_a {
            for ancestor_b in &ancestry_b {
                if ancestor_a.name == ancestor_b.name {
                    return Some(ancestor_a);
                }
            }
        }

        None
    }

    /// Get the merge base TXG for two branches.
    pub fn merge_base_txg(&self, branch_a: &str, branch_b: &str) -> Option<u64> {
        let a = self.branches.get(branch_a)?;
        let b = self.branches.get(branch_b)?;

        // If one is ancestor of the other, use fork_txg
        let ancestry_a: Vec<_> = self.ancestry(branch_a);
        for ancestor in &ancestry_a {
            if ancestor.name == branch_b {
                return Some(a.fork_txg);
            }
        }

        let ancestry_b: Vec<_> = self.ancestry(branch_b);
        for ancestor in &ancestry_b {
            if ancestor.name == branch_a {
                return Some(b.fork_txg);
            }
        }

        // Find common ancestor
        self.common_ancestor(branch_a, branch_b)
            .map(|ancestor| ancestor.head_txg)
    }

    /// List branches sorted by name.
    pub fn list_sorted(&self) -> Vec<&Branch> {
        let mut branches: Vec<_> = self.branches.values().collect();
        branches.sort_by(|a, b| a.name.cmp(&b.name));
        branches
    }

    /// List branches sorted by last activity (head_txg).
    pub fn list_by_activity(&self) -> Vec<&Branch> {
        let mut branches: Vec<_> = self.branches.values().collect();
        branches.sort_by(|a, b| b.head_txg.cmp(&a.head_txg));
        branches
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// TESTS
// ═══════════════════════════════════════════════════════════════════════════════

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

    #[test]
    fn test_new_registry() {
        let registry = BranchRegistry::new(0, 1704067200);

        assert_eq!(registry.len(), 1);
        assert_eq!(registry.current(), "main");
        assert!(registry.contains("main"));

        let main = registry.get("main").unwrap();
        assert!(main.is_default);
        assert!(main.is_root());
    }

    #[test]
    fn test_create_branch() {
        let mut registry = BranchRegistry::new(0, 1704067200);

        // Update main's head
        registry.update_head("main", 10).unwrap();

        // Create feature branch
        let guid = registry.create_branch("feature", 1704067200).unwrap();
        assert_eq!(guid, 2);
        assert!(registry.contains("feature"));

        let feature = registry.get("feature").unwrap();
        assert_eq!(feature.parent, Some("main".into()));
        assert_eq!(feature.fork_txg, 10);
        assert!(!feature.is_default);
    }

    #[test]
    fn test_create_duplicate_branch() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("feature", 1704067200).unwrap();

        let result = registry.create_branch("feature", 1704067200);
        assert!(matches!(result, Err(BranchError::BranchExists(_))));
    }

    #[test]
    fn test_validate_name() {
        // Valid names
        assert!(BranchRegistry::validate_name("feature").is_ok());
        assert!(BranchRegistry::validate_name("feature-123").is_ok());
        assert!(BranchRegistry::validate_name("my_branch").is_ok());

        // Invalid names
        assert!(BranchRegistry::validate_name("").is_err());
        assert!(BranchRegistry::validate_name(".hidden").is_err());
        assert!(BranchRegistry::validate_name("-dash").is_err());
        assert!(BranchRegistry::validate_name("has/slash").is_err());
        assert!(BranchRegistry::validate_name("has space").is_err());
        assert!(BranchRegistry::validate_name("double..dot").is_err());
    }

    #[test]
    fn test_checkout() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("feature", 1704067200).unwrap();

        registry.checkout("feature").unwrap();
        assert_eq!(registry.current(), "feature");

        // Already on branch
        let result = registry.checkout("feature");
        assert!(matches!(result, Err(BranchError::AlreadyOnBranch(_))));

        // Non-existent branch
        let result = registry.checkout("nonexistent");
        assert!(matches!(result, Err(BranchError::BranchNotFound(_))));
    }

    #[test]
    fn test_delete_branch() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("feature", 1704067200).unwrap();

        // Can delete feature (force=true)
        let deleted = registry.delete_branch("feature", true).unwrap();
        assert_eq!(deleted.name, "feature");
        assert!(!registry.contains("feature"));
    }

    #[test]
    fn test_cannot_delete_current() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("feature", 1704067200).unwrap();
        registry.checkout("feature").unwrap();

        let result = registry.delete_branch("feature", true);
        assert!(matches!(result, Err(BranchError::CannotDeleteCurrent(_))));
    }

    #[test]
    fn test_cannot_delete_default() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("feature", 1704067200).unwrap();
        registry.checkout("feature").unwrap();

        let result = registry.delete_branch("main", true);
        assert!(matches!(result, Err(BranchError::CannotDeleteDefault(_))));
    }

    #[test]
    fn test_rename() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("feature", 1704067200).unwrap();
        registry.create_branch("child", 1704067200).unwrap();

        // Create child from feature
        {
            let child = registry.get_mut("child").unwrap();
            child.parent = Some("feature".into());
        }

        // Rename feature
        registry.rename("feature", "feature-v2").unwrap();

        assert!(!registry.contains("feature"));
        assert!(registry.contains("feature-v2"));

        // Child's parent should be updated
        let child = registry.get("child").unwrap();
        assert_eq!(child.parent, Some("feature-v2".into()));
    }

    #[test]
    fn test_ancestry() {
        let mut registry = BranchRegistry::new(0, 1704067200);

        // main -> feature -> sub-feature
        registry.create_branch("feature", 1704067200).unwrap();
        registry.checkout("feature").unwrap();
        registry.create_branch("sub-feature", 1704067200).unwrap();

        let ancestry = registry.ancestry("sub-feature");
        assert_eq!(ancestry.len(), 3);
        assert_eq!(ancestry[0].name, "sub-feature");
        assert_eq!(ancestry[1].name, "feature");
        assert_eq!(ancestry[2].name, "main");
    }

    #[test]
    fn test_common_ancestor() {
        let mut registry = BranchRegistry::new(0, 1704067200);

        // main -> feature1
        //      -> feature2
        registry.create_branch("feature1", 1704067200).unwrap();
        registry.create_branch("feature2", 1704067200).unwrap();

        let ancestor = registry.common_ancestor("feature1", "feature2");
        assert!(ancestor.is_some());
        assert_eq!(ancestor.unwrap().name, "main");
    }

    #[test]
    fn test_children() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("feature1", 1704067200).unwrap();
        registry.create_branch("feature2", 1704067200).unwrap();

        let children = registry.children("main");
        assert_eq!(children.len(), 2);
    }

    #[test]
    fn test_set_default() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("develop", 1704067200).unwrap();

        // Set develop as default
        registry.set_default("develop").unwrap();

        let develop = registry.get("develop").unwrap();
        assert!(develop.is_default);

        let main = registry.get("main").unwrap();
        assert!(!main.is_default);
    }

    #[test]
    fn test_update_head() {
        let mut registry = BranchRegistry::new(0, 1704067200);

        registry.update_head("main", 100).unwrap();
        assert_eq!(registry.get("main").unwrap().head_txg, 100);

        registry.update_current_head(200).unwrap();
        assert_eq!(registry.get("main").unwrap().head_txg, 200);
    }

    #[test]
    fn test_checkout_or_create() {
        let mut registry = BranchRegistry::new(0, 1704067200);

        // Create new branch
        registry.checkout_or_create("feature", 1704067200).unwrap();
        assert_eq!(registry.current(), "feature");
        assert!(registry.contains("feature"));

        // Checkout existing
        registry.checkout("main").unwrap();
        registry.checkout_or_create("feature", 1704067200).unwrap();
        assert_eq!(registry.current(), "feature");
    }

    #[test]
    fn test_list_sorted() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("zebra", 1704067200).unwrap();
        registry.create_branch("alpha", 1704067200).unwrap();

        let sorted = registry.list_sorted();
        assert_eq!(sorted[0].name, "alpha");
        assert_eq!(sorted[1].name, "main");
        assert_eq!(sorted[2].name, "zebra");
    }

    #[test]
    fn test_list_by_activity() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("old", 1704067200).unwrap();
        registry.create_branch("new", 1704067200).unwrap();

        registry.update_head("new", 100).unwrap();
        registry.update_head("old", 50).unwrap();

        let by_activity = registry.list_by_activity();
        assert_eq!(by_activity[0].name, "new");
        assert_eq!(by_activity[1].name, "old");
    }

    #[test]
    fn test_create_branch_from() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.create_branch("develop", 1704067200).unwrap();
        registry.update_head("develop", 50).unwrap();

        // Create from develop, not current (main)
        let guid = registry
            .create_branch_from("feature", "develop", 1704067200)
            .unwrap();
        assert!(guid > 0);

        let feature = registry.get("feature").unwrap();
        assert_eq!(feature.parent, Some("develop".into()));
        assert_eq!(feature.fork_txg, 50);
    }

    #[test]
    fn test_create_branch_at_txg() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.update_head("main", 100).unwrap();

        // Create at specific TXG (historical)
        let guid = registry
            .create_branch_at_txg("historical", "main", 50, 1704067200)
            .unwrap();
        assert!(guid > 0);

        let historical = registry.get("historical").unwrap();
        assert_eq!(historical.fork_txg, 50);
        assert_eq!(historical.head_txg, 50);
    }

    #[test]
    fn test_merge_base_txg() {
        let mut registry = BranchRegistry::new(0, 1704067200);
        registry.update_head("main", 100).unwrap();

        registry.create_branch("feature1", 1704067200).unwrap();
        registry.update_head("feature1", 150).unwrap();

        // Create feature2 directly from main (we're already on main)
        registry.create_branch("feature2", 1704067200).unwrap();
        registry.update_head("feature2", 120).unwrap();

        // Both branched from main at TXG 100
        let base = registry.merge_base_txg("feature1", "feature2");
        assert!(base.is_some());
        // The common ancestor is main with head_txg=100
        assert_eq!(base.unwrap(), 100);
    }
}