atproto-plc / src / validation.rs

//! Validation logic for operations and operation chains

use crate::crypto::VerifyingKey;
use crate::document::{PlcState, MAX_VERIFICATION_METHODS};
use crate::error::{PlcError, Result};
use crate::operations::Operation;
use chrono::{DateTime, Duration, Utc};
use std::collections::HashMap;

/// Recovery window duration (72 hours)
const RECOVERY_WINDOW_HOURS: i64 = 72;

/// Represents an operation with its server-assigned timestamp
#[derive(Debug, Clone)]
pub struct OperationWithTimestamp {
    /// The operation itself
    pub operation: Operation,
    /// Server-assigned timestamp when the operation was received
    pub timestamp: DateTime<Utc>,
}

/// Represents a fork point where multiple operations reference the same prev CID
#[derive(Debug)]
struct ForkPoint {
    /// The prev CID that multiple operations reference
    #[allow(dead_code)]
    prev_cid: String,
    /// Competing operations at this fork point, with their timestamps and signing key indices
    operations: Vec<(Operation, DateTime<Utc>, usize)>,
}

impl ForkPoint {
    /// Create a new fork point
    fn new(prev_cid: String) -> Self {
        Self {
            prev_cid,
            operations: Vec::new(),
        }
    }

    /// Add an operation to this fork point
    fn add_operation(&mut self, operation: Operation, timestamp: DateTime<Utc>, key_index: usize) {
        self.operations.push((operation, timestamp, key_index));
    }

    /// Resolve this fork point and return the canonical operation
    ///
    /// Resolution algorithm:
    /// 1. Sort operations by timestamp (first-received wins by default)
    /// 2. Check if any later-received operation with higher priority can invalidate within 72 hours
    /// 3. Return the winning operation
    fn resolve(&self) -> Result<Operation> {
        if self.operations.is_empty() {
            return Err(PlcError::ForkResolutionError(
                "Fork point has no operations".to_string(),
            ));
        }

        // If only one operation, it wins by default
        if self.operations.len() == 1 {
            return Ok(self.operations[0].0.clone());
        }

        // Sort by timestamp - first-received wins by default
        let mut sorted = self.operations.clone();
        sorted.sort_by_key(|(_, timestamp, _)| *timestamp);

        // The first operation in chronological order is the default winner
        let (mut canonical_op, mut canonical_ts, mut canonical_key_idx) = sorted[0].clone();

        // Check if any later-received operation can invalidate the current canonical
        for (competing_op, competing_ts, competing_key_idx) in &sorted[1..] {
            // Can only invalidate if the competing operation has higher priority (lower index)
            if *competing_key_idx < canonical_key_idx {
                // Check if within recovery window from the canonical operation
                let time_diff = *competing_ts - canonical_ts;
                if time_diff <= Duration::hours(RECOVERY_WINDOW_HOURS)
                    && time_diff >= Duration::zero()
                {
                    // This higher-priority operation invalidates the canonical one
                    // Update the canonical to this operation
                    canonical_op = competing_op.clone();
                    canonical_ts = *competing_ts;
                    canonical_key_idx = *competing_key_idx;
                }
            }
        }

        Ok(canonical_op)
    }
}

/// Builder for constructing the canonical chain from a set of operations with timestamps
struct CanonicalChainBuilder {
    /// All operations with timestamps, in the order they were received
    operations: Vec<OperationWithTimestamp>,
}

impl CanonicalChainBuilder {
    /// Create a new builder
    fn new(operations: Vec<OperationWithTimestamp>) -> Self {
        Self { operations }
    }

    /// Build the canonical chain by detecting and resolving forks
    ///
    /// Algorithm:
    /// 1. Build a graph of operations by CID
    /// 2. Detect fork points (multiple operations with same prev CID)
    /// 3. Resolve each fork using rotation key priority and recovery window
    /// 4. Build the canonical chain from genesis to tip
    fn build(&self) -> Result<Vec<Operation>> {
        if self.operations.is_empty() {
            return Err(PlcError::EmptyChain);
        }

        // Build a map of CID -> (operation, timestamp) for quick lookup
        let mut operation_map: HashMap<String, (Operation, DateTime<Utc>)> = HashMap::new();
        for op_with_ts in &self.operations {
            let cid = op_with_ts.operation.cid()?;
            operation_map.insert(cid, (op_with_ts.operation.clone(), op_with_ts.timestamp));
        }

        // Find the genesis operation (prev = None)
        let genesis = self
            .operations
            .iter()
            .find(|op| op.operation.is_genesis())
            .ok_or_else(|| PlcError::FirstOperationNotGenesis)?;

        // Detect fork points - group operations by their prev CID
        let mut prev_to_operations: HashMap<String, Vec<(Operation, DateTime<Utc>)>> =
            HashMap::new();

        for op_with_ts in &self.operations {
            if let Some(prev_cid) = op_with_ts.operation.prev() {
                prev_to_operations
                    .entry(prev_cid.to_string())
                    .or_default()
                    .push((op_with_ts.operation.clone(), op_with_ts.timestamp));
            }
        }

        // Build fork points for any prev CID with multiple operations
        let mut fork_points: HashMap<String, ForkPoint> = HashMap::new();
        for (prev_cid, operations) in &prev_to_operations {
            if operations.len() > 1 {
                // This is a fork point - multiple operations reference the same prev
                let mut fork_point = ForkPoint::new(prev_cid.clone());

                // For each operation, determine which rotation key signed it
                // We need to look at the state at the prev operation
                for (operation, timestamp) in operations {
                    // Get the state at the prev operation to find rotation keys
                    let rotation_keys = if let Some((prev_op, _)) = operation_map.get(prev_cid) {
                        // Build state up to prev operation to get its rotation keys
                        self.get_state_at_operation(prev_op)?
                    } else {
                        PlcState::new()
                    };

                    // Find which rotation key signed this operation
                    let key_index = self.find_signing_key_index(operation, &rotation_keys)?;
                    fork_point.add_operation(operation.clone(), *timestamp, key_index);
                }

                fork_points.insert(prev_cid.clone(), fork_point);
            }
        }

        // Build the canonical chain by starting from genesis and following the canonical path
        let mut canonical_chain = vec![genesis.operation.clone()];
        let mut current_cid = genesis.operation.cid()?;

        // Follow the chain until we reach the tip
        loop {
            // Check if there's a fork at this point
            if let Some(fork_point) = fork_points.get(&current_cid) {
                // Resolve the fork and get the canonical operation
                let canonical_op = fork_point.resolve()?;
                let next_cid = canonical_op.cid()?;
                canonical_chain.push(canonical_op);
                current_cid = next_cid;
            } else if let Some(operations) = prev_to_operations.get(&current_cid) {
                // No fork, just a single operation
                if operations.len() == 1 {
                    let (operation, _) = &operations[0];
                    let next_cid = operation.cid()?;
                    canonical_chain.push(operation.clone());
                    current_cid = next_cid;
                } else {
                    // This shouldn't happen - we should have detected this as a fork
                    return Err(PlcError::ForkResolutionError(
                        "Unexpected multiple operations without fork point".to_string(),
                    ));
                }
            } else {
                // No more operations - we've reached the tip
                break;
            }
        }

        Ok(canonical_chain)
    }

    /// Find the index of the rotation key that signed this operation
    fn find_signing_key_index(
        &self,
        operation: &Operation,
        state: &PlcState,
    ) -> Result<usize> {
        if state.rotation_keys.is_empty() {
            return Err(PlcError::InvalidRotationKeys(
                "No rotation keys available for verification".to_string(),
            ));
        }

        // Try each rotation key and return the index of the first one that verifies
        for (index, key_did) in state.rotation_keys.iter().enumerate() {
            let verifying_key = VerifyingKey::from_did_key(key_did)?;
            if operation.verify(&[verifying_key]).is_ok() {
                return Ok(index);
            }
        }

        // No key verified the signature
        Err(PlcError::SignatureVerificationFailed)
    }

    /// Get the state at a specific operation
    ///
    /// This reconstructs the state by applying all operations from genesis up to
    /// and including the specified operation.
    fn get_state_at_operation(&self, target_operation: &Operation) -> Result<PlcState> {
        let mut state = PlcState::new();
        let target_cid = target_operation.cid()?;

        // Build the chain up to the target operation
        // We need to traverse from genesis to the target
        for op_with_ts in &self.operations {
            let op = &op_with_ts.operation;

            // Apply this operation to the state
            match op {
                Operation::PlcOperation {
                    rotation_keys,
                    verification_methods,
                    also_known_as,
                    services,
                    ..
                } => {
                    state.rotation_keys = rotation_keys.clone();
                    state.verification_methods = verification_methods.clone();
                    state.also_known_as = also_known_as.clone();
                    state.services = services.clone();
                }
                Operation::PlcTombstone { .. } => {
                    state = PlcState::new();
                }
                Operation::LegacyCreate { .. } => {
                    return Err(PlcError::InvalidOperationType(
                        "Legacy create operations not fully supported".to_string(),
                    ));
                }
            }

            // Check if this is the target operation
            if op.cid()? == target_cid {
                break;
            }
        }

        Ok(state)
    }
}

/// Operation chain validator
pub struct OperationChainValidator;

impl OperationChainValidator {
    /// Validate a complete operation chain and return the final state
    ///
    /// # Errors
    ///
    /// Returns errors if:
    /// - Chain is empty
    /// - First operation is not genesis
    /// - Any operation has invalid prev reference
    /// - Any signature is invalid
    /// - Any operation violates constraints
    pub fn validate_chain(operations: &[Operation]) -> Result<PlcState> {
        if operations.is_empty() {
            return Err(PlcError::EmptyChain);
        }

        // First operation must be genesis
        if !operations[0].is_genesis() {
            return Err(PlcError::FirstOperationNotGenesis);
        }

        let mut current_state = PlcState::new();
        let mut prev_cid: Option<String> = None;

        for (i, operation) in operations.iter().enumerate() {
            // Verify prev field matches expected CID
            if i == 0 {
                // Genesis operation must have prev = None
                if operation.prev().is_some() {
                    return Err(PlcError::InvalidPrev(
                        "Genesis operation must have prev = null".to_string(),
                    ));
                }
            } else {
                // Non-genesis operations must reference previous CID
                let expected_prev = prev_cid.as_ref().ok_or_else(|| {
                    PlcError::ChainValidationFailed("Missing previous CID".to_string())
                })?;

                let actual_prev = operation.prev().ok_or_else(|| {
                    PlcError::InvalidPrev("Non-genesis operation must have prev field".to_string())
                })?;

                if actual_prev != expected_prev {
                    return Err(PlcError::InvalidPrev(format!(
                        "Expected prev = {}, got {}",
                        expected_prev, actual_prev
                    )));
                }
            }

            // Verify signature using current rotation keys
            if !current_state.rotation_keys.is_empty() {
                let verifying_keys: Result<Vec<VerifyingKey>> = current_state
                    .rotation_keys
                    .iter()
                    .map(|k| VerifyingKey::from_did_key(k))
                    .collect();

                let verifying_keys = verifying_keys?;
                operation.verify(&verifying_keys)?;
            } else if i > 0 {
                // After genesis, we must have rotation keys
                return Err(PlcError::InvalidRotationKeys(
                    "No rotation keys available for verification".to_string(),
                ));
            }

            // Apply operation to state
            match operation {
                Operation::PlcOperation {
                    rotation_keys,
                    verification_methods,
                    also_known_as,
                    services,
                    ..
                } => {
                    current_state.rotation_keys = rotation_keys.clone();
                    current_state.verification_methods = verification_methods.clone();
                    current_state.also_known_as = also_known_as.clone();
                    current_state.services = services.clone();

                    // Validate the state
                    current_state.validate()?;
                }
                Operation::PlcTombstone { .. } => {
                    // Tombstone marks the DID as deleted
                    // Clear all state
                    current_state = PlcState::new();
                }
                Operation::LegacyCreate { .. } => {
                    // Legacy create format - convert to modern format
                    // This is for backwards compatibility
                    return Err(PlcError::InvalidOperationType(
                        "Legacy create operations not fully supported".to_string(),
                    ));
                }
            }

            // Update prev CID for next iteration
            prev_cid = Some(operation.cid()?);
        }

        Ok(current_state)
    }

    /// Validate a chain with fork resolution
    ///
    /// This handles the recovery mechanism where operations signed by higher-priority
    /// rotation keys can invalidate later operations if submitted within 72 hours.
    ///
    /// # Arguments
    ///
    /// * `operations` - All operations in the audit log (may contain forks)
    /// * `timestamps` - Server-assigned timestamps for each operation
    ///
    /// # Returns
    ///
    /// The final state after applying the canonical chain (after fork resolution)
    ///
    /// # Errors
    ///
    /// Returns errors if:
    /// - Operations and timestamps arrays have different lengths
    /// - Chain is empty
    /// - First operation is not genesis
    /// - Fork resolution fails
    /// - Any signature is invalid
    /// - Any operation violates constraints
    pub fn validate_chain_with_forks(
        operations: &[Operation],
        timestamps: &[DateTime<Utc>],
    ) -> Result<PlcState> {
        if operations.len() != timestamps.len() {
            return Err(PlcError::ChainValidationFailed(
                "Operations and timestamps length mismatch".to_string(),
            ));
        }

        if operations.is_empty() {
            return Err(PlcError::EmptyChain);
        }

        // Build operations with timestamps
        let operations_with_timestamps: Vec<OperationWithTimestamp> = operations
            .iter()
            .zip(timestamps.iter())
            .map(|(op, ts)| OperationWithTimestamp {
                operation: op.clone(),
                timestamp: *ts,
            })
            .collect();

        // Use the canonical chain builder to resolve forks
        let builder = CanonicalChainBuilder::new(operations_with_timestamps);
        let canonical_chain = builder.build()?;

        // Validate the canonical chain
        Self::validate_chain(&canonical_chain)
    }

    /// Check if an operation is within the recovery window relative to another operation
    ///
    /// Returns true if the time difference is less than 72 hours
    pub fn is_within_recovery_window(
        fork_timestamp: DateTime<Utc>,
        current_timestamp: DateTime<Utc>,
    ) -> bool {
        let diff = current_timestamp - fork_timestamp;
        diff < Duration::hours(RECOVERY_WINDOW_HOURS) && diff >= Duration::zero()
    }
}

/// Validate rotation keys
///
/// # Errors
///
/// Returns errors if:
/// - Not 1-5 keys
/// - Contains duplicates
/// - Invalid did:key format
/// - Unsupported key type
pub fn validate_rotation_keys(keys: &[String]) -> Result<()> {
    if keys.is_empty() {
        return Err(PlcError::InvalidRotationKeys(
            "At least one rotation key is required".to_string(),
        ));
    }

    if keys.len() > 5 {
        return Err(PlcError::TooManyEntries {
            field: "rotation_keys".to_string(),
            max: 5,
            actual: keys.len(),
        });
    }

    // Check for duplicates
    let mut seen = std::collections::HashSet::new();
    for key in keys {
        if !seen.insert(key) {
            return Err(PlcError::DuplicateEntry {
                field: "rotation_keys".to_string(),
                value: key.clone(),
            });
        }

        // Validate format
        if !key.starts_with("did:key:") {
            return Err(PlcError::InvalidRotationKeys(format!(
                "Rotation key must be in did:key format: {}",
                key
            )));
        }

        // Try to parse to ensure it's valid
        VerifyingKey::from_did_key(key)?;
    }

    Ok(())
}

/// Validate verification methods
///
/// # Errors
///
/// Returns errors if:
/// - More than 10 methods
/// - Invalid did:key format
pub fn validate_verification_methods(
    methods: &std::collections::HashMap<String, String>,
) -> Result<()> {
    if methods.len() > MAX_VERIFICATION_METHODS {
        return Err(PlcError::TooManyEntries {
            field: "verification_methods".to_string(),
            max: MAX_VERIFICATION_METHODS,
            actual: methods.len(),
        });
    }

    for (name, key) in methods {
        if !key.starts_with("did:key:") {
            return Err(PlcError::InvalidVerificationMethods(format!(
                "Verification method '{}' must be in did:key format: {}",
                name, key
            )));
        }

        // Try to parse to ensure it's valid
        VerifyingKey::from_did_key(key)?;
    }

    Ok(())
}

/// Validate also-known-as URIs
///
/// # Errors
///
/// Returns errors if any URI is invalid
pub fn validate_also_known_as(uris: &[String]) -> Result<()> {
    for uri in uris {
        if uri.is_empty() {
            return Err(PlcError::InvalidAlsoKnownAs(
                "URI cannot be empty".to_string(),
            ));
        }

        // Basic URI validation - should start with a scheme
        if !uri.contains(':') {
            return Err(PlcError::InvalidAlsoKnownAs(format!(
                "URI must contain a scheme: {}",
                uri
            )));
        }
    }

    Ok(())
}

/// Validate service endpoints
///
/// # Errors
///
/// Returns errors if any service is invalid
pub fn validate_services(
    services: &std::collections::HashMap<String, crate::document::ServiceEndpoint>,
) -> Result<()> {
    for (name, service) in services {
        if name.is_empty() {
            return Err(PlcError::InvalidService(
                "Service name cannot be empty".to_string(),
            ));
        }

        service.validate()?;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::crypto::SigningKey;
    use crate::document::ServiceEndpoint;
    use std::collections::HashMap;

    #[test]
    fn test_validate_rotation_keys() {
        let key1 = SigningKey::generate_p256();
        let key2 = SigningKey::generate_k256();

        let keys = vec![key1.to_did_key(), key2.to_did_key()];
        assert!(validate_rotation_keys(&keys).is_ok());

        // Empty keys
        assert!(validate_rotation_keys(&[]).is_err());

        // Too many keys
        let many_keys: Vec<String> = (0..6).map(|_| SigningKey::generate_p256().to_did_key()).collect();
        assert!(validate_rotation_keys(&many_keys).is_err());

        // Duplicate keys
        let dup_key = key1.to_did_key();
        let dup_keys = vec![dup_key.clone(), dup_key];
        assert!(validate_rotation_keys(&dup_keys).is_err());
    }

    #[test]
    fn test_validate_verification_methods() {
        let mut methods = HashMap::new();
        let key = SigningKey::generate_p256();
        methods.insert("atproto".to_string(), key.to_did_key());

        assert!(validate_verification_methods(&methods).is_ok());

        // Too many methods
        let mut many_methods = HashMap::new();
        for i in 0..11 {
            let key = SigningKey::generate_p256();
            many_methods.insert(format!("key{}", i), key.to_did_key());
        }
        assert!(validate_verification_methods(&many_methods).is_err());
    }

    #[test]
    fn test_validate_also_known_as() {
        let uris = vec![
            "at://alice.example.com".to_string(),
            "https://example.com".to_string(),
        ];
        assert!(validate_also_known_as(&uris).is_ok());

        // Empty URI
        assert!(validate_also_known_as(&[String::new()]).is_err());

        // Invalid URI (no scheme)
        assert!(validate_also_known_as(&["not-a-uri".to_string()]).is_err());
    }

    #[test]
    fn test_recovery_window() {
        let base = Utc::now();
        let within = base + Duration::hours(24);
        let outside = base + Duration::hours(100);

        assert!(OperationChainValidator::is_within_recovery_window(base, within));
        assert!(!OperationChainValidator::is_within_recovery_window(base, outside));
    }

    #[test]
    fn test_validate_chain_genesis() {
        let key = SigningKey::generate_p256();
        let did_key = key.to_did_key();

        let unsigned = Operation::new_genesis(
            vec![did_key],
            HashMap::new(),
            vec![],
            HashMap::new(),
        );

        let signed = unsigned.sign(&key).unwrap();

        // Single genesis operation should validate
        let state = OperationChainValidator::validate_chain(&[signed]).unwrap();
        assert_eq!(state.rotation_keys.len(), 1);
    }

    #[test]
    fn test_validate_chain_empty() {
        assert!(OperationChainValidator::validate_chain(&[]).is_err());
    }

    // ========================================================================
    // Fork Resolution Tests
    // ========================================================================

    /// Test simple fork resolution where higher priority key wins within recovery window
    #[test]
    fn test_fork_resolution_priority_within_window() {
        // Create two rotation keys with different priorities
        let primary_key = SigningKey::generate_p256(); // Index 0 - highest priority
        let backup_key = SigningKey::generate_k256(); // Index 1 - lower priority

        let rotation_keys = vec![primary_key.to_did_key(), backup_key.to_did_key()];

        // Create genesis operation
        let genesis = Operation::new_genesis(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            HashMap::new(),
        )
        .sign(&primary_key)
        .unwrap();

        let genesis_cid = genesis.cid().unwrap();
        let genesis_time = Utc::now();

        // Create two competing operations that both reference genesis
        // Operation A: signed by backup key (lower priority)
        let mut services_a = HashMap::new();
        services_a.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://pds-a.example.com".to_string(),
            },
        );

        let op_a = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_a,
            genesis_cid.clone(),
        )
        .sign(&backup_key)
        .unwrap();

        let op_a_time = genesis_time + Duration::hours(1);

        // Operation B: signed by primary key (higher priority), arrives 24 hours after A
        let mut services_b = HashMap::new();
        services_b.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://pds-b.example.com".to_string(),
            },
        );

        let op_b = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_b,
            genesis_cid,
        )
        .sign(&primary_key)
        .unwrap();

        let op_b_time = op_a_time + Duration::hours(24); // Within 72-hour window

        // Build operations and timestamps arrays
        let operations = vec![genesis.clone(), op_a.clone(), op_b.clone()];
        let timestamps = vec![genesis_time, op_a_time, op_b_time];

        // Validate with fork resolution
        let result = OperationChainValidator::validate_chain_with_forks(&operations, &timestamps);
        assert!(result.is_ok());

        let state = result.unwrap();

        // Operation B should win because it has higher priority (lower index)
        // and was received within the 72-hour recovery window
        assert_eq!(
            state.services.get("pds").unwrap().endpoint,
            "https://pds-b.example.com"
        );
    }

    /// Test fork resolution where lower priority operation wins after recovery window expires
    #[test]
    fn test_fork_resolution_priority_outside_window() {
        let primary_key = SigningKey::generate_p256(); // Index 0
        let backup_key = SigningKey::generate_k256(); // Index 1

        let rotation_keys = vec![primary_key.to_did_key(), backup_key.to_did_key()];

        let genesis = Operation::new_genesis(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            HashMap::new(),
        )
        .sign(&primary_key)
        .unwrap();

        let genesis_cid = genesis.cid().unwrap();
        let genesis_time = Utc::now();

        // Operation A: signed by backup key (lower priority)
        let mut services_a = HashMap::new();
        services_a.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://pds-a.example.com".to_string(),
            },
        );

        let op_a = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_a,
            genesis_cid.clone(),
        )
        .sign(&backup_key)
        .unwrap();

        let op_a_time = genesis_time + Duration::hours(1);

        // Operation B: signed by primary key, arrives 100 hours after A (outside 72-hour window)
        let mut services_b = HashMap::new();
        services_b.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://pds-b.example.com".to_string(),
            },
        );

        let op_b = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_b,
            genesis_cid,
        )
        .sign(&primary_key)
        .unwrap();

        let op_b_time = op_a_time + Duration::hours(100); // Outside 72-hour window

        let operations = vec![genesis.clone(), op_a.clone(), op_b.clone()];
        let timestamps = vec![genesis_time, op_a_time, op_b_time];

        let result = OperationChainValidator::validate_chain_with_forks(&operations, &timestamps);
        assert!(result.is_ok());

        let state = result.unwrap();

        // Operation A should win because even though B has higher priority,
        // it arrived outside the 72-hour recovery window
        assert_eq!(
            state.services.get("pds").unwrap().endpoint,
            "https://pds-a.example.com"
        );
    }

    /// Test multiple forks at different points in the chain
    #[test]
    fn test_fork_resolution_multiple_forks() {
        let key1 = SigningKey::generate_p256();
        let key2 = SigningKey::generate_k256();
        let key3 = SigningKey::generate_p256();

        let rotation_keys = vec![key1.to_did_key(), key2.to_did_key(), key3.to_did_key()];

        // Genesis
        let genesis = Operation::new_genesis(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            HashMap::new(),
        )
        .sign(&key1)
        .unwrap();

        let genesis_cid = genesis.cid().unwrap();
        let genesis_time = Utc::now();

        // First fork at genesis
        let mut services_1a = HashMap::new();
        services_1a.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://fork1a.example.com".to_string(),
            },
        );

        let op_1a = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_1a,
            genesis_cid.clone(),
        )
        .sign(&key2)
        .unwrap();

        let op_1a_time = genesis_time + Duration::hours(1);

        // Second operation at same fork (higher priority, within window)
        let mut services_1b = HashMap::new();
        services_1b.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://fork1b.example.com".to_string(),
            },
        );

        let op_1b = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_1b,
            genesis_cid,
        )
        .sign(&key1)
        .unwrap();

        let op_1b_time = op_1a_time + Duration::hours(2);

        // Next operation continues from op_1b (the winner)
        let op_1b_cid = op_1b.cid().unwrap();

        let mut services_2 = HashMap::new();
        services_2.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://continuation.example.com".to_string(),
            },
        );

        let op_2 = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_2,
            op_1b_cid,
        )
        .sign(&key1)
        .unwrap();

        let op_2_time = op_1b_time + Duration::hours(1);

        let operations = vec![
            genesis.clone(),
            op_1a.clone(),
            op_1b.clone(),
            op_2.clone(),
        ];
        let timestamps = vec![genesis_time, op_1a_time, op_1b_time, op_2_time];

        let result = OperationChainValidator::validate_chain_with_forks(&operations, &timestamps);
        assert!(result.is_ok());

        let state = result.unwrap();

        // The final state should be from op_2, which continues from op_1b
        assert_eq!(
            state.services.get("pds").unwrap().endpoint,
            "https://continuation.example.com"
        );
    }

    /// Test fork with three competing operations
    #[test]
    fn test_fork_resolution_three_way_fork() {
        let key1 = SigningKey::generate_p256(); // Highest priority
        let key2 = SigningKey::generate_k256(); // Medium priority
        let key3 = SigningKey::generate_p256(); // Lowest priority

        let rotation_keys = vec![key1.to_did_key(), key2.to_did_key(), key3.to_did_key()];

        let genesis = Operation::new_genesis(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            HashMap::new(),
        )
        .sign(&key1)
        .unwrap();

        let genesis_cid = genesis.cid().unwrap();
        let genesis_time = Utc::now();

        // Three competing operations
        let mut services_a = HashMap::new();
        services_a.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://op-a.example.com".to_string(),
            },
        );

        let op_a = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_a,
            genesis_cid.clone(),
        )
        .sign(&key3)
        .unwrap(); // Lowest priority, first to arrive

        let op_a_time = genesis_time + Duration::hours(1);

        let mut services_b = HashMap::new();
        services_b.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://op-b.example.com".to_string(),
            },
        );

        let op_b = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_b,
            genesis_cid.clone(),
        )
        .sign(&key2)
        .unwrap(); // Medium priority

        let op_b_time = op_a_time + Duration::hours(2);

        let mut services_c = HashMap::new();
        services_c.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://op-c.example.com".to_string(),
            },
        );

        let op_c = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_c,
            genesis_cid,
        )
        .sign(&key1)
        .unwrap(); // Highest priority, within window

        let op_c_time = op_a_time + Duration::hours(10);

        let operations = vec![
            genesis.clone(),
            op_a.clone(),
            op_b.clone(),
            op_c.clone(),
        ];
        let timestamps = vec![genesis_time, op_a_time, op_b_time, op_c_time];

        let result = OperationChainValidator::validate_chain_with_forks(&operations, &timestamps);
        assert!(result.is_ok());

        let state = result.unwrap();

        // Operation C should win (highest priority, within window)
        assert_eq!(
            state.services.get("pds").unwrap().endpoint,
            "https://op-c.example.com"
        );
    }

    /// Test no fork - linear chain should work as before
    #[test]
    fn test_fork_resolution_no_fork() {
        let key = SigningKey::generate_p256();
        let rotation_keys = vec![key.to_did_key()];

        let genesis = Operation::new_genesis(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            HashMap::new(),
        )
        .sign(&key)
        .unwrap();

        let genesis_cid = genesis.cid().unwrap();
        let genesis_time = Utc::now();

        let mut services = HashMap::new();
        services.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://pds.example.com".to_string(),
            },
        );

        let op1 = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services.clone(),
            genesis_cid,
        )
        .sign(&key)
        .unwrap();

        let op1_time = genesis_time + Duration::hours(1);

        let operations = vec![genesis.clone(), op1.clone()];
        let timestamps = vec![genesis_time, op1_time];

        let result = OperationChainValidator::validate_chain_with_forks(&operations, &timestamps);
        assert!(result.is_ok());

        let state = result.unwrap();
        assert_eq!(
            state.services.get("pds").unwrap().endpoint,
            "https://pds.example.com"
        );
    }

    /// Test fork resolution with rotation key changes
    #[test]
    fn test_fork_resolution_with_key_rotation() {
        let key1 = SigningKey::generate_p256();
        let key2 = SigningKey::generate_k256();
        let key3 = SigningKey::generate_p256();

        // Initial rotation keys
        let rotation_keys_v1 = vec![key1.to_did_key(), key2.to_did_key()];

        let genesis = Operation::new_genesis(
            rotation_keys_v1.clone(),
            HashMap::new(),
            vec![],
            HashMap::new(),
        )
        .sign(&key1)
        .unwrap();

        let genesis_cid = genesis.cid().unwrap();
        let genesis_time = Utc::now();

        // Update rotation keys in first operation
        let rotation_keys_v2 = vec![key1.to_did_key(), key3.to_did_key()];

        let op1 = Operation::new_update(
            rotation_keys_v2.clone(),
            HashMap::new(),
            vec![],
            HashMap::new(),
            genesis_cid,
        )
        .sign(&key1)
        .unwrap();

        let op1_cid = op1.cid().unwrap();
        let op1_time = genesis_time + Duration::hours(1);

        // Create a fork at op1 - both operations use the new rotation keys
        let mut services_a = HashMap::new();
        services_a.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://op-a.example.com".to_string(),
            },
        );

        let op_a = Operation::new_update(
            rotation_keys_v2.clone(),
            HashMap::new(),
            vec![],
            services_a,
            op1_cid.clone(),
        )
        .sign(&key3)
        .unwrap(); // Index 1 in new keys

        let op_a_time = op1_time + Duration::hours(1);

        let mut services_b = HashMap::new();
        services_b.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://op-b.example.com".to_string(),
            },
        );

        let op_b = Operation::new_update(
            rotation_keys_v2.clone(),
            HashMap::new(),
            vec![],
            services_b,
            op1_cid,
        )
        .sign(&key1)
        .unwrap(); // Index 0 in new keys (higher priority)

        let op_b_time = op_a_time + Duration::hours(2);

        let operations = vec![
            genesis.clone(),
            op1.clone(),
            op_a.clone(),
            op_b.clone(),
        ];
        let timestamps = vec![genesis_time, op1_time, op_a_time, op_b_time];

        let result = OperationChainValidator::validate_chain_with_forks(&operations, &timestamps);
        assert!(result.is_ok());

        let state = result.unwrap();

        // Operation B should win (signed by key1 which is index 0 in the new rotation keys)
        assert_eq!(
            state.services.get("pds").unwrap().endpoint,
            "https://op-b.example.com"
        );
    }

    /// Test that operations with mismatched timestamps and operations fail
    #[test]
    fn test_fork_resolution_mismatched_lengths() {
        let key = SigningKey::generate_p256();
        let rotation_keys = vec![key.to_did_key()];

        let genesis = Operation::new_genesis(
            rotation_keys,
            HashMap::new(),
            vec![],
            HashMap::new(),
        )
        .sign(&key)
        .unwrap();

        let operations = vec![genesis];
        let timestamps = vec![Utc::now(), Utc::now()]; // Different length

        let result = OperationChainValidator::validate_chain_with_forks(&operations, &timestamps);
        assert!(result.is_err());
    }

    /// Test recovery window boundary (exactly 72 hours)
    #[test]
    fn test_fork_resolution_recovery_window_boundary() {
        let primary_key = SigningKey::generate_p256();
        let backup_key = SigningKey::generate_k256();

        let rotation_keys = vec![primary_key.to_did_key(), backup_key.to_did_key()];

        let genesis = Operation::new_genesis(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            HashMap::new(),
        )
        .sign(&primary_key)
        .unwrap();

        let genesis_cid = genesis.cid().unwrap();
        let genesis_time = Utc::now();

        // Operation A: signed by backup key
        let mut services_a = HashMap::new();
        services_a.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://op-a.example.com".to_string(),
            },
        );

        let op_a = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_a,
            genesis_cid.clone(),
        )
        .sign(&backup_key)
        .unwrap();

        let op_a_time = genesis_time + Duration::hours(1);

        // Operation B: exactly at 72-hour boundary (should still be within window)
        let mut services_b = HashMap::new();
        services_b.insert(
            "pds".to_string(),
            ServiceEndpoint {
                service_type: "AtprotoPersonalDataServer".to_string(),
                endpoint: "https://op-b.example.com".to_string(),
            },
        );

        let op_b = Operation::new_update(
            rotation_keys.clone(),
            HashMap::new(),
            vec![],
            services_b,
            genesis_cid,
        )
        .sign(&primary_key)
        .unwrap();

        // Exactly 72 hours after op_a
        let op_b_time = op_a_time + Duration::hours(72);

        let operations = vec![genesis.clone(), op_a.clone(), op_b.clone()];
        let timestamps = vec![genesis_time, op_a_time, op_b_time];

        let result = OperationChainValidator::validate_chain_with_forks(&operations, &timestamps);
        assert!(result.is_ok());

        let state = result.unwrap();

        // At exactly 72 hours, the higher priority operation should still win
        assert_eq!(
            state.services.get("pds").unwrap().endpoint,
            "https://op-b.example.com"
        );
    }
}