lcpfs 2026.1.102

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

//! Branch operations: create, checkout, delete.
//!
//! This module implements the core branch operations that interact with
//! the underlying COW filesystem to create zero-copy branches.

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

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

// ═══════════════════════════════════════════════════════════════════════════════
// DATASET PROVIDER TRAIT
// ═══════════════════════════════════════════════════════════════════════════════

/// Trait for dataset operations required by branch management.
///
/// This abstraction allows the branch system to work with any COW filesystem
/// that supports cloning and snapshotting.
pub trait DatasetProvider {
    /// Clone a dataset (COW operation).
    ///
    /// Creates a zero-copy clone of the source dataset at the given TXG.
    /// Returns the GUID of the new dataset.
    fn clone_dataset(&mut self, source_guid: u64, name: &str, at_txg: u64) -> Result<u64, String>;

    /// Delete a dataset.
    fn delete_dataset(&mut self, guid: u64) -> Result<(), String>;

    /// Activate a dataset (make it the current working dataset).
    fn activate_dataset(&mut self, guid: u64) -> Result<(), String>;

    /// Get the current TXG of a dataset.
    fn get_dataset_txg(&self, guid: u64) -> Result<u64, String>;

    /// Sync all pending changes to a dataset.
    fn sync_dataset(&mut self, guid: u64) -> Result<u64, String>;

    /// Get the current timestamp.
    fn current_timestamp(&self) -> u64;
}

// ═══════════════════════════════════════════════════════════════════════════════
// BRANCH OPERATIONS
// ═══════════════════════════════════════════════════════════════════════════════

/// Branch operations manager.
///
/// Provides high-level branch operations that coordinate between
/// the branch registry and the underlying dataset provider.
pub struct BranchOps<'a, P: DatasetProvider> {
    /// The dataset provider.
    provider: &'a mut P,
    /// The branch registry.
    registry: &'a mut BranchRegistry,
}

impl<'a, P: DatasetProvider> BranchOps<'a, P> {
    /// Create a new BranchOps instance.
    pub fn new(provider: &'a mut P, registry: &'a mut BranchRegistry) -> Self {
        Self { provider, registry }
    }

    /// Get the registry.
    pub fn registry(&self) -> &BranchRegistry {
        self.registry
    }

    /// Get the registry mutably.
    pub fn registry_mut(&mut self) -> &mut BranchRegistry {
        self.registry
    }

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

    /// Create a new branch from the current branch.
    ///
    /// This is a zero-copy operation due to COW semantics.
    pub fn create_branch(&mut self, name: &str) -> Result<&Branch, BranchError> {
        let timestamp = self.provider.current_timestamp();

        // Get current branch info
        let (source_guid, source_txg) = {
            let current = self
                .registry
                .current_branch()
                .ok_or_else(|| BranchError::Internal("no current branch".into()))?;
            (current.guid, current.head_txg)
        };

        // Clone dataset (zero-copy via COW)
        let new_guid = self
            .provider
            .clone_dataset(source_guid, name, source_txg)
            .map_err(BranchError::IoError)?;

        // Register the branch
        self.registry.create_branch(name, timestamp)?;

        // Update the GUID in registry
        {
            let branch = self.registry.get_mut(name).unwrap();
            branch.guid = new_guid;
        }

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

    /// Create a branch from a specific parent.
    pub fn create_branch_from(&mut self, name: &str, parent: &str) -> Result<&Branch, BranchError> {
        let timestamp = self.provider.current_timestamp();

        // Get parent branch info
        let (source_guid, source_txg) = {
            let parent_branch = self
                .registry
                .get(parent)
                .ok_or_else(|| BranchError::BranchNotFound(parent.to_string()))?;
            (parent_branch.guid, parent_branch.head_txg)
        };

        // Clone dataset
        let new_guid = self
            .provider
            .clone_dataset(source_guid, name, source_txg)
            .map_err(BranchError::IoError)?;

        // Register the branch
        self.registry.create_branch_from(name, parent, timestamp)?;

        // Update the GUID
        {
            let branch = self.registry.get_mut(name).unwrap();
            branch.guid = new_guid;
        }

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

    /// Create a branch at a specific TXG (historical point).
    pub fn create_branch_at(
        &mut self,
        name: &str,
        parent: &str,
        txg: u64,
    ) -> Result<&Branch, BranchError> {
        let timestamp = self.provider.current_timestamp();

        // Get parent branch info
        let source_guid = {
            let parent_branch = self
                .registry
                .get(parent)
                .ok_or_else(|| BranchError::BranchNotFound(parent.to_string()))?;
            parent_branch.guid
        };

        // Clone dataset at specific TXG
        let new_guid = self
            .provider
            .clone_dataset(source_guid, name, txg)
            .map_err(BranchError::IoError)?;

        // Register the branch
        self.registry
            .create_branch_at_txg(name, parent, txg, timestamp)?;

        // Update the GUID
        {
            let branch = self.registry.get_mut(name).unwrap();
            branch.guid = new_guid;
        }

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

    // ═══════════════════════════════════════════════════════════════════════════
    // CHECKOUT
    // ═══════════════════════════════════════════════════════════════════════════

    /// Switch to a different branch.
    ///
    /// This activates the dataset associated with the target branch.
    pub fn checkout(&mut self, name: &str) -> Result<&Branch, BranchError> {
        // Get target branch info
        let target_guid = {
            let branch = self
                .registry
                .get(name)
                .ok_or_else(|| BranchError::BranchNotFound(name.to_string()))?;
            branch.guid
        };

        // Sync current branch before switching
        self.sync_current()?;

        // Activate the target dataset
        self.provider
            .activate_dataset(target_guid)
            .map_err(BranchError::IoError)?;

        // Update registry
        self.registry.checkout(name)
    }

    /// Create and checkout a new branch in one operation.
    pub fn checkout_new(&mut self, name: &str) -> Result<&Branch, BranchError> {
        self.create_branch(name)?;
        self.checkout(name)
    }

    /// Checkout or create a branch if it doesn't exist.
    pub fn checkout_or_create(&mut self, name: &str) -> Result<&Branch, BranchError> {
        if self.registry.contains(name) {
            self.checkout(name)
        } else {
            self.checkout_new(name)
        }
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // DELETE
    // ═══════════════════════════════════════════════════════════════════════════

    /// Delete a branch.
    ///
    /// Removes both the registry entry and the underlying dataset.
    pub fn delete_branch(&mut self, name: &str, force: bool) -> Result<Branch, BranchError> {
        // Get the branch GUID before deleting from registry
        let guid = {
            let branch = self
                .registry
                .get(name)
                .ok_or_else(|| BranchError::BranchNotFound(name.to_string()))?;
            branch.guid
        };

        // Delete from registry first (validates constraints)
        let branch = self.registry.delete_branch(name, force)?;

        // Delete the underlying dataset
        self.provider
            .delete_dataset(guid)
            .map_err(BranchError::IoError)?;

        Ok(branch)
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // SYNC
    // ═══════════════════════════════════════════════════════════════════════════

    /// Sync the current branch.
    pub fn sync_current(&mut self) -> Result<u64, BranchError> {
        let guid = {
            let current = self
                .registry
                .current_branch()
                .ok_or_else(|| BranchError::Internal("no current branch".into()))?;
            current.guid
        };

        let new_txg = self
            .provider
            .sync_dataset(guid)
            .map_err(BranchError::IoError)?;

        // Update head TXG
        self.registry.update_current_head(new_txg)?;

        Ok(new_txg)
    }

    /// Sync a specific branch.
    pub fn sync_branch(&mut self, name: &str) -> Result<u64, BranchError> {
        let guid = {
            let branch = self
                .registry
                .get(name)
                .ok_or_else(|| BranchError::BranchNotFound(name.to_string()))?;
            branch.guid
        };

        let new_txg = self
            .provider
            .sync_dataset(guid)
            .map_err(BranchError::IoError)?;

        // Update head TXG
        self.registry.update_head(name, new_txg)?;

        Ok(new_txg)
    }

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

    /// Get all branches.
    pub fn list(&self) -> Vec<&Branch> {
        self.registry.list_sorted()
    }

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

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

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

    /// Get branch status (ahead/behind counts).
    pub fn status(&self, name: &str) -> Result<BranchStatus, BranchError> {
        let branch = self
            .registry
            .get(name)
            .ok_or_else(|| BranchError::BranchNotFound(name.to_string()))?;

        let commits_since_fork = branch.head_txg.saturating_sub(branch.fork_txg);

        // Get parent status
        let (parent_name, parent_ahead) = if let Some(ref parent) = branch.parent {
            if let Some(parent_branch) = self.registry.get(parent) {
                let parent_ahead = parent_branch.head_txg.saturating_sub(branch.fork_txg);
                (Some(parent.clone()), parent_ahead)
            } else {
                (Some(parent.clone()), 0)
            }
        } else {
            (None, 0)
        };

        Ok(BranchStatus {
            name: branch.name.clone(),
            is_current: self.registry.current() == name,
            is_default: branch.is_default,
            parent: parent_name,
            ahead: commits_since_fork,
            behind: parent_ahead,
            head_txg: branch.head_txg,
        })
    }

    /// Rename a branch.
    pub fn rename(&mut self, old_name: &str, new_name: &str) -> Result<(), BranchError> {
        self.registry.rename(old_name, new_name)
    }

    /// Set the default branch.
    pub fn set_default(&mut self, name: &str) -> Result<(), BranchError> {
        self.registry.set_default(name)
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// BRANCH STATUS
// ═══════════════════════════════════════════════════════════════════════════════

/// Branch status information.
#[derive(Debug, Clone)]
pub struct BranchStatus {
    /// Branch name.
    pub name: String,
    /// Is this the current branch?
    pub is_current: bool,
    /// Is this the default branch?
    pub is_default: bool,
    /// Parent branch name.
    pub parent: Option<String>,
    /// TXGs ahead of fork point.
    pub ahead: u64,
    /// TXGs parent is ahead of fork point.
    pub behind: u64,
    /// Current head TXG.
    pub head_txg: u64,
}

impl BranchStatus {
    /// Check if branch is up to date with parent.
    pub fn is_up_to_date(&self) -> bool {
        self.behind == 0
    }

    /// Check if branch has changes.
    pub fn has_changes(&self) -> bool {
        self.ahead > 0
    }
}

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

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

    /// Mock dataset provider for testing.
    struct MockProvider {
        next_guid: u64,
        active_guid: Option<u64>,
        datasets: Vec<u64>,
        txg_map: alloc::collections::BTreeMap<u64, u64>,
    }

    impl MockProvider {
        fn new() -> Self {
            let mut provider = Self {
                next_guid: 1,
                active_guid: None,
                datasets: Vec::new(),
                txg_map: alloc::collections::BTreeMap::new(),
            };
            // Create initial "main" dataset
            provider.datasets.push(1);
            provider.txg_map.insert(1, 0);
            provider.active_guid = Some(1);
            provider.next_guid = 2;
            provider
        }
    }

    impl DatasetProvider for MockProvider {
        fn clone_dataset(
            &mut self,
            _source_guid: u64,
            _name: &str,
            at_txg: u64,
        ) -> Result<u64, String> {
            let guid = self.next_guid;
            self.next_guid += 1;
            self.datasets.push(guid);
            self.txg_map.insert(guid, at_txg);
            Ok(guid)
        }

        fn delete_dataset(&mut self, guid: u64) -> Result<(), String> {
            self.datasets.retain(|&g| g != guid);
            self.txg_map.remove(&guid);
            Ok(())
        }

        fn activate_dataset(&mut self, guid: u64) -> Result<(), String> {
            if self.datasets.contains(&guid) {
                self.active_guid = Some(guid);
                Ok(())
            } else {
                Err("dataset not found".into())
            }
        }

        fn get_dataset_txg(&self, guid: u64) -> Result<u64, String> {
            self.txg_map
                .get(&guid)
                .copied()
                .ok_or_else(|| "dataset not found".into())
        }

        fn sync_dataset(&mut self, guid: u64) -> Result<u64, String> {
            let txg = self
                .txg_map
                .get_mut(&guid)
                .ok_or_else(|| "dataset not found".to_string())?;
            *txg += 1;
            Ok(*txg)
        }

        fn current_timestamp(&self) -> u64 {
            1704067200
        }
    }

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

        let mut ops = BranchOps::new(&mut provider, &mut registry);

        let branch = ops.create_branch("feature").unwrap();
        assert_eq!(branch.name, "feature");
        assert!(provider.datasets.len() == 2);
    }

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

        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.create_branch("feature").unwrap();
        }

        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.checkout("feature").unwrap();
        }

        assert_eq!(registry.current(), "feature");
    }

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

        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.checkout_new("feature").unwrap();
        }

        assert_eq!(registry.current(), "feature");
        assert!(registry.contains("feature"));
    }

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

        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.create_branch("feature").unwrap();
        }

        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.delete_branch("feature", true).unwrap();
        }

        assert!(!registry.contains("feature"));
        assert_eq!(provider.datasets.len(), 1);
    }

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

        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            let txg = ops.sync_current().unwrap();
            assert_eq!(txg, 1);
        }

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

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

        // Sync main a few times
        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.sync_current().unwrap();
            ops.sync_current().unwrap();
        }

        // Create feature branch
        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.create_branch("feature").unwrap();
            ops.checkout("feature").unwrap();
            ops.sync_current().unwrap();
        }

        // Continue advancing main
        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.checkout("main").unwrap();
            ops.sync_current().unwrap();
            ops.sync_current().unwrap();
        }

        // Check feature status
        {
            let ops = BranchOps::new(&mut provider, &mut registry);
            let status = ops.status("feature").unwrap();
            assert_eq!(status.name, "feature");
            assert!(!status.is_current);
            assert_eq!(status.parent, Some("main".into()));
            assert!(status.ahead > 0);
            assert!(status.behind > 0);
        }
    }

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

        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.create_branch("feature").unwrap();
            ops.rename("feature", "feature-v2").unwrap();
        }

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

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

        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.create_branch("develop").unwrap();
            ops.set_default("develop").unwrap();
        }

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

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

        // Create develop and advance it
        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.create_branch("develop").unwrap();
            ops.checkout("develop").unwrap();
            ops.sync_current().unwrap();
            ops.sync_current().unwrap();
        }

        // Create feature from develop (not current main)
        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.checkout("main").unwrap();
            ops.create_branch_from("feature", "develop").unwrap();
        }

        let feature = registry.get("feature").unwrap();
        assert_eq!(feature.parent, Some("develop".into()));
        // fork_txg is develop's head_txg after 2 syncs
        // But we're using MockProvider, so the exact TXG depends on implementation
        // Main assertion: feature is a child of develop
        assert!(feature.fork_txg > 0);
    }

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

        // Advance main
        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            for _ in 0..5 {
                ops.sync_current().unwrap();
            }
        }

        // Create branch at historical TXG
        {
            let mut ops = BranchOps::new(&mut provider, &mut registry);
            ops.create_branch_at("historical", "main", 2).unwrap();
        }

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