//! Entity lock system for collaborative editing
//!
//! Provides optimistic entity locking to prevent concurrent modifications.
//! Locks are acquired when entities are selected and released when deselected.
//!
//! # Lock Protocol
//!
//! 1. **Acquisition**: User selects entity → broadcast `LockRequest`
//! 2. **Optimistic Apply**: All peers apply lock locally
//! 3. **Confirm**: Holder broadcasts `LockAcquired`
//! 4. **Conflict Resolution**: If two nodes acquire simultaneously, higher node ID wins
//! 5. **Release**: User deselects entity → broadcast `LockReleased`
//! 6. **Timeout**: 5-second timeout as crash recovery fallback
//!
//! # Example
//!
//! ```no_run
//! use bevy::prelude::*;
//! use libmarathon::networking::{EntityLockRegistry, acquire_entity_lock, release_entity_lock};
//! use uuid::Uuid;
//!
//! fn my_system(world: &mut World) {
//!     let entity_id = Uuid::new_v4();
//!     let node_id = Uuid::new_v4();
//!
//!     let mut registry = world.resource_mut::<EntityLockRegistry>();
//!
//!     // Acquire lock when user selects entity
//!     registry.try_acquire(entity_id, node_id);
//!
//!     // Release lock when user deselects entity
//!     registry.release(entity_id, node_id);
//! }
//! ```

use std::{
    collections::HashMap,
    time::{
        Duration,
        Instant,
    },
};

use bevy::prelude::*;

use uuid::Uuid;

use crate::networking::{
    GossipBridge,
    NodeId,
    VersionedMessage,
    delta_generation::NodeVectorClock,
    messages::SyncMessage,
};

/// Duration before a lock automatically expires (crash recovery)
pub const LOCK_TIMEOUT: Duration = Duration::from_secs(5);

/// Maximum number of concurrent locks per node (rate limiting)
pub const MAX_LOCKS_PER_NODE: usize = 100;

/// Lock acquisition/release messages
#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize, PartialEq, Eq)]
pub enum LockMessage {
    /// Request to acquire a lock on an entity
    LockRequest {
        entity_id: Uuid,
        node_id: NodeId,
    },

    /// Confirmation that a lock was successfully acquired
    LockAcquired {
        entity_id: Uuid,
        holder: NodeId,
    },

    /// Lock acquisition failed (already locked by another node)
    LockRejected {
        entity_id: Uuid,
        requester: NodeId,
        current_holder: NodeId,
    },

    /// Heartbeat to renew a held lock (sent ~1/sec by holder)
    ///
    /// If no heartbeat is received for 5 seconds, the lock expires.
    /// This provides automatic crash recovery without explicit timeouts.
    LockHeartbeat {
        entity_id: Uuid,
        holder: NodeId,
    },

    /// Request to release a lock
    LockRelease {
        entity_id: Uuid,
        node_id: NodeId,
    },

    /// Confirmation that a lock was released
    LockReleased {
        entity_id: Uuid,
    },
}

/// Information about an active entity lock
#[derive(Debug, Clone)]
pub struct EntityLock {
    /// ID of the entity being locked
    pub entity_id: Uuid,

    /// Node that holds the lock
    pub holder: NodeId,

    /// When the last heartbeat was received (or when lock was acquired)
    pub last_heartbeat: Instant,

    /// Lock timeout duration (expires if no heartbeat for this long)
    pub timeout: Duration,
}

impl EntityLock {
    /// Create a new entity lock
    pub fn new(entity_id: Uuid, holder: NodeId) -> Self {
        Self {
            entity_id,
            holder,
            last_heartbeat: Instant::now(),
            timeout: LOCK_TIMEOUT,
        }
    }

    /// Renew the lock with a heartbeat
    pub fn renew(&mut self) {
        self.last_heartbeat = Instant::now();
    }

    /// Check if the lock has expired (no heartbeat for > timeout)
    pub fn is_expired(&self) -> bool {
        self.last_heartbeat.elapsed() >= self.timeout
    }

    /// Check if this lock is held by the given node
    pub fn is_held_by(&self, node_id: NodeId) -> bool {
        self.holder == node_id
    }
}

/// Registry of all active entity locks
///
/// This resource tracks which entities are locked and by whom.
/// It's used to prevent concurrent modifications to the same entity.
#[derive(Resource, Default)]
pub struct EntityLockRegistry {
    /// Map of entity ID to lock info
    locks: HashMap<Uuid, EntityLock>,

    /// Count of locks held by each node (for rate limiting)
    locks_per_node: HashMap<NodeId, usize>,
}

impl EntityLockRegistry {
    /// Create a new empty lock registry
    pub fn new() -> Self {
        Self {
            locks: HashMap::new(),
            locks_per_node: HashMap::new(),
        }
    }

    /// Try to acquire a lock on an entity
    ///
    /// Returns Ok(()) if lock was acquired, Err with current holder if already locked.
    pub fn try_acquire(&mut self, entity_id: Uuid, node_id: NodeId) -> Result<(), NodeId> {
        // Check if already locked
        if let Some(existing_lock) = self.locks.get(&entity_id) {
            // If expired, allow re-acquisition
            if !existing_lock.is_expired() {
                return Err(existing_lock.holder);
            }

            // Remove expired lock
            self.remove_lock(entity_id);
        }

        // Check rate limit
        let node_lock_count = self.locks_per_node.get(&node_id).copied().unwrap_or(0);
        if node_lock_count >= MAX_LOCKS_PER_NODE {
            warn!(
                "Node {} at lock limit ({}/{}), rejecting acquisition",
                node_id, node_lock_count, MAX_LOCKS_PER_NODE
            );
            return Err(node_id); // Return self as "holder" to indicate rate limit
        }

        // Acquire the lock
        let lock = EntityLock::new(entity_id, node_id);
        self.locks.insert(entity_id, lock);

        // Update node lock count
        *self.locks_per_node.entry(node_id).or_insert(0) += 1;

        debug!("Lock acquired: entity {} by node {}", entity_id, node_id);
        Ok(())
    }

    /// Release a lock on an entity
    ///
    /// Only succeeds if the node currently holds the lock.
    pub fn release(&mut self, entity_id: Uuid, node_id: NodeId) -> bool {
        if let Some(lock) = self.locks.get(&entity_id) {
            if lock.holder == node_id {
                self.remove_lock(entity_id);
                debug!("Lock released: entity {} by node {}", entity_id, node_id);
                return true;
            } else {
                warn!(
                    "Node {} tried to release lock held by node {}",
                    node_id, lock.holder
                );
            }
        }
        false
    }

    /// Force release a lock (for timeout cleanup)
    pub fn force_release(&mut self, entity_id: Uuid) {
        if self.locks.remove(&entity_id).is_some() {
            debug!("Lock force-released: entity {}", entity_id);
        }
    }

    /// Check if an entity is locked by any node
    ///
    /// Takes the local node ID to properly handle expiration:
    /// - Our own locks are never considered expired (held exactly as long as selected)
    /// - Remote locks are subject to the 5-second timeout
    pub fn is_locked(&self, entity_id: Uuid, local_node_id: NodeId) -> bool {
        self.locks.get(&entity_id).map_or(false, |lock| {
            // Our own locks never expire
            lock.holder == local_node_id || !lock.is_expired()
        })
    }

    /// Check if an entity is locked by a specific node
    ///
    /// Takes the local node ID to properly handle expiration:
    /// - If checking our own lock, ignore expiration (held exactly as long as selected)
    /// - If checking another node's lock, apply 5-second timeout
    pub fn is_locked_by(&self, entity_id: Uuid, node_id: NodeId, local_node_id: NodeId) -> bool {
        self.locks.get(&entity_id).map_or(false, |lock| {
            if lock.holder != node_id {
                // Not held by the queried node
                false
            } else if lock.holder == local_node_id {
                // Checking our own lock - never expires
                true
            } else {
                // Checking remote lock - check expiration
                !lock.is_expired()
            }
        })
    }

    /// Get the holder of a lock (if locked and not expired)
    ///
    /// Takes the local node ID to properly handle expiration:
    /// - Our own locks are never considered expired
    /// - Remote locks are subject to the 5-second timeout
    pub fn get_holder(&self, entity_id: Uuid, local_node_id: NodeId) -> Option<NodeId> {
        self.locks.get(&entity_id).and_then(|lock| {
            // Our own locks never expire
            if lock.holder == local_node_id || !lock.is_expired() {
                Some(lock.holder)
            } else {
                None
            }
        })
    }

    /// Renew a lock's heartbeat
    ///
    /// Returns true if the heartbeat was renewed, false if lock doesn't exist
    /// or is held by a different node.
    pub fn renew_heartbeat(&mut self, entity_id: Uuid, node_id: NodeId) -> bool {
        if let Some(lock) = self.locks.get_mut(&entity_id) {
            if lock.holder == node_id {
                lock.renew();
                return true;
            }
        }
        false
    }

    /// Get all expired locks
    pub fn get_expired_locks(&self) -> Vec<Uuid> {
        self.locks
            .iter()
            .filter(|(_, lock)| lock.is_expired())
            .map(|(entity_id, _)| *entity_id)
            .collect()
    }

    /// Get number of locks held by a node
    pub fn get_node_lock_count(&self, node_id: NodeId) -> usize {
        self.locks_per_node.get(&node_id).copied().unwrap_or(0)
    }

    /// Get total number of active locks
    pub fn total_locks(&self) -> usize {
        self.locks.len()
    }

    /// Remove a lock and update bookkeeping
    fn remove_lock(&mut self, entity_id: Uuid) {
        if let Some(lock) = self.locks.remove(&entity_id) {
            // Decrement node lock count
            if let Some(count) = self.locks_per_node.get_mut(&lock.holder) {
                *count = count.saturating_sub(1);
                if *count == 0 {
                    self.locks_per_node.remove(&lock.holder);
                }
            }
        }
    }

    /// Test helper: Manually expire a lock by setting its heartbeat timestamp to the past
    ///
    /// This is only intended for testing purposes to simulate lock expiration without waiting.
    pub fn expire_lock_for_testing(&mut self, entity_id: Uuid) {
        if let Some(lock) = self.locks.get_mut(&entity_id) {
            lock.last_heartbeat = Instant::now() - Duration::from_secs(10);
        }
    }
}

/// System to acquire locks when entities are selected
///
/// This system detects when entities are added to the global `LocalSelection`
/// resource and attempts to acquire locks on those entities, broadcasting
/// the request to other peers.
pub fn acquire_locks_on_selection_system(
    mut registry: ResMut<EntityLockRegistry>,
    node_clock: Res<NodeVectorClock>,
    bridge: Option<Res<GossipBridge>>,
    selection: Res<crate::networking::LocalSelection>,
) {
    // Only run when selection changes
    if !selection.is_changed() {
        return;
    }

    let node_id = node_clock.node_id;

    // Try to acquire locks for all selected entities
    for &entity_id in selection.iter() {
        let already_locked = registry.is_locked_by(entity_id, node_id, node_id);

        // Only try to acquire if we don't already hold the lock
        if !already_locked {
            match registry.try_acquire(entity_id, node_id) {
                Ok(()) => {
                    info!("Acquired lock on newly selected entity {}", entity_id);

                    // Broadcast LockRequest
                    if let Some(ref bridge) = bridge {
                        let msg = VersionedMessage::new(SyncMessage::Lock(LockMessage::LockRequest {
                            entity_id,
                            node_id,
                        }));

                        if let Err(e) = bridge.send(msg) {
                            error!("Failed to broadcast LockRequest on selection: {}", e);
                        } else {
                            debug!("LockRequest broadcast successful for entity {}", entity_id);
                        }
                    } else {
                        warn!("No GossipBridge available to broadcast LockRequest");
                    }
                }
                Err(holder) => {
                    warn!("Failed to acquire lock on selected entity {} (held by {})", entity_id, holder);
                }
            }
        }
    }
}

/// System to release locks when entities are deselected
///
/// This system detects when entities are removed from the global `LocalSelection`
/// resource and releases any locks held on those entities, broadcasting the release
/// to other peers.
///
/// Add to your app as an Update system:
/// ```no_run
/// use bevy::prelude::*;
/// use libmarathon::networking::release_locks_on_deselection_system;
///
/// App::new().add_systems(Update, release_locks_on_deselection_system);
/// ```
pub fn release_locks_on_deselection_system(
    mut registry: ResMut<EntityLockRegistry>,
    node_clock: Res<NodeVectorClock>,
    bridge: Option<Res<GossipBridge>>,
    selection: Res<crate::networking::LocalSelection>,
) {
    // Only run when selection changes
    if !selection.is_changed() {
        return;
    }

    let node_id = node_clock.node_id;

    // Check all locks held by this node
    let locks_to_release: Vec<Uuid> = registry
        .locks
        .iter()
        .filter(|(entity_id, lock)| {
            // Release if held by us and not currently selected
            lock.holder == node_id && !selection.contains(**entity_id)
        })
        .map(|(entity_id, _)| *entity_id)
        .collect();

    if !locks_to_release.is_empty() {
        info!("Selection cleared, releasing {} locks", locks_to_release.len());
    }

    // Release each lock and broadcast
    for entity_id in locks_to_release {
        if registry.release(entity_id, node_id) {
            info!("Released lock on deselected entity {}", entity_id);

            // Broadcast LockRelease
            if let Some(ref bridge) = bridge {
                let msg = VersionedMessage::new(SyncMessage::Lock(LockMessage::LockRelease {
                    entity_id,
                    node_id,
                }));

                if let Err(e) = bridge.send(msg) {
                    error!("Failed to broadcast LockRelease on deselection: {}", e);
                } else {
                    info!("Lock released on deselection: entity {}", entity_id);
                }
            }
        }
    }
}

/// System to clean up expired locks (crash recovery)
///
/// This system periodically removes locks that have exceeded their timeout
/// duration (default 5 seconds). This provides crash recovery - if a **remote**
/// node crashes while holding a lock, it will eventually expire.
///
/// **Important**: Only remote locks are cleaned up. Local locks (held by this node)
/// are never timed out - they're held exactly as long as entities are selected,
/// and only released via deselection.
///
/// Add to your app as an Update system:
/// ```no_run
/// use bevy::prelude::*;
/// use libmarathon::networking::cleanup_expired_locks_system;
///
/// App::new().add_systems(Update, cleanup_expired_locks_system);
/// ```
pub fn cleanup_expired_locks_system(
    mut registry: ResMut<EntityLockRegistry>,
    node_clock: Res<NodeVectorClock>,
    bridge: Option<Res<GossipBridge>>,
) {
    let node_id = node_clock.node_id;

    // Only clean up REMOTE locks (locks held by other nodes)
    // Our own locks are managed by release_locks_on_deselection_system
    let expired: Vec<Uuid> = registry
        .locks
        .iter()
        .filter(|(_, lock)| {
            // Only expire locks held by OTHER nodes
            lock.is_expired() && lock.holder != node_id
        })
        .map(|(entity_id, _)| *entity_id)
        .collect();

    if !expired.is_empty() {
        info!("Cleaning up {} expired remote locks", expired.len());

        for entity_id in expired {
            debug!("Force-releasing expired remote lock on entity {}", entity_id);
            registry.force_release(entity_id);

            // Broadcast LockReleased
            if let Some(ref bridge) = bridge {
                let msg =
                    VersionedMessage::new(SyncMessage::Lock(LockMessage::LockReleased { entity_id }));

                if let Err(e) = bridge.send(msg) {
                    error!("Failed to broadcast LockReleased for expired lock: {}", e);
                } else {
                    info!("Expired remote lock cleaned up: entity {}", entity_id);
                }
            }
        }
    }
}

/// System to broadcast heartbeats for all locks we currently hold
///
/// This system runs periodically (~1/sec) and broadcasts a heartbeat for each
/// lock this node holds. This keeps locks alive and provides crash detection -
/// if a node crashes, heartbeats stop and locks expire after 5 seconds.
///
/// Add to your app as an Update system with a run condition to throttle it:
/// ```no_run
/// use bevy::prelude::*;
/// use bevy::time::common_conditions::on_timer;
/// use std::time::Duration;
/// use libmarathon::networking::broadcast_lock_heartbeats_system;
///
/// App::new().add_systems(Update,
///     broadcast_lock_heartbeats_system.run_if(on_timer(Duration::from_secs(1)))
/// );
/// ```
pub fn broadcast_lock_heartbeats_system(
    mut registry: ResMut<EntityLockRegistry>,
    node_clock: Res<NodeVectorClock>,
    bridge: Option<Res<GossipBridge>>,
) {
    let node_id = node_clock.node_id;

    // Find all locks held by this node
    let our_locks: Vec<Uuid> = registry
        .locks
        .iter()
        .filter(|(_, lock)| lock.holder == node_id && !lock.is_expired())
        .map(|(entity_id, _)| *entity_id)
        .collect();

    if our_locks.is_empty() {
        return;
    }

    debug!("Broadcasting {} lock heartbeats", our_locks.len());

    // Renew local locks and broadcast heartbeat for each lock
    for entity_id in &our_locks {
        // Renew the lock locally first (don't rely on network loopback)
        registry.renew_heartbeat(*entity_id, node_id);
    }

    // Broadcast heartbeat messages to peers
    if let Some(ref bridge) = bridge {
        for entity_id in our_locks {
            let msg = VersionedMessage::new(SyncMessage::Lock(LockMessage::LockHeartbeat {
                entity_id,
                holder: node_id,
            }));

            if let Err(e) = bridge.send(msg) {
                error!(
                    "Failed to broadcast heartbeat for entity {}: {}",
                    entity_id, e
                );
            } else {
                trace!("Heartbeat sent for locked entity {}", entity_id);
            }
        }
    }
}

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

    #[test]
    fn test_lock_acquisition() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node_id = Uuid::new_v4();

        // Should acquire successfully
        assert!(registry.try_acquire(entity_id, node_id).is_ok());
        assert!(registry.is_locked(entity_id, node_id));
        assert!(registry.is_locked_by(entity_id, node_id, node_id));
        assert_eq!(registry.get_holder(entity_id, node_id), Some(node_id));
    }

    #[test]
    fn test_lock_conflict() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node1 = Uuid::new_v4();
        let node2 = Uuid::new_v4();

        // Node 1 acquires
        assert!(registry.try_acquire(entity_id, node1).is_ok());

        // Node 2 should be rejected
        assert_eq!(registry.try_acquire(entity_id, node2), Err(node1));
    }

    #[test]
    fn test_lock_release() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node_id = Uuid::new_v4();

        // Acquire and release
        registry.try_acquire(entity_id, node_id).unwrap();
        assert!(registry.release(entity_id, node_id));
        assert!(!registry.is_locked(entity_id, node_id));
    }

    #[test]
    fn test_wrong_node_cannot_release() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node1 = Uuid::new_v4();
        let node2 = Uuid::new_v4();

        // Node 1 acquires
        registry.try_acquire(entity_id, node1).unwrap();

        // Node 2 cannot release
        assert!(!registry.release(entity_id, node2));
        assert!(registry.is_locked(entity_id, node2));
        assert!(registry.is_locked_by(entity_id, node1, node2));
    }

    #[test]
    fn test_lock_timeout() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node_id = Uuid::new_v4();

        // Acquire with very short timeout
        registry.try_acquire(entity_id, node_id).unwrap();

        // Manually set timeout to 0 for testing
        if let Some(lock) = registry.locks.get_mut(&entity_id) {
            lock.timeout = Duration::from_secs(0);
        }

        // Should be detected as expired
        let expired = registry.get_expired_locks();
        assert_eq!(expired.len(), 1);
        assert_eq!(expired[0], entity_id);
    }

    #[test]
    fn test_force_release() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node_id = Uuid::new_v4();

        registry.try_acquire(entity_id, node_id).unwrap();
        registry.force_release(entity_id);
        assert!(!registry.is_locked(entity_id, node_id));
    }

    #[test]
    fn test_rate_limiting() {
        let mut registry = EntityLockRegistry::new();
        let node_id = Uuid::new_v4();

        // Acquire MAX_LOCKS_PER_NODE locks
        for _ in 0..MAX_LOCKS_PER_NODE {
            let entity_id = Uuid::new_v4();
            assert!(registry.try_acquire(entity_id, node_id).is_ok());
        }

        // Next acquisition should fail (rate limit)
        let entity_id = Uuid::new_v4();
        assert!(registry.try_acquire(entity_id, node_id).is_err());
    }

    #[test]
    fn test_node_lock_count() {
        let mut registry = EntityLockRegistry::new();
        let node_id = Uuid::new_v4();

        assert_eq!(registry.get_node_lock_count(node_id), 0);

        // Acquire 3 locks
        for _ in 0..3 {
            let entity_id = Uuid::new_v4();
            registry.try_acquire(entity_id, node_id).unwrap();
        }

        assert_eq!(registry.get_node_lock_count(node_id), 3);
        assert_eq!(registry.total_locks(), 3);
    }

    #[test]
    fn test_lock_message_serialization() {
        let entity_id = Uuid::new_v4();
        let node_id = Uuid::new_v4();

        let messages = vec![
            LockMessage::LockRequest { entity_id, node_id },
            LockMessage::LockAcquired {
                entity_id,
                holder: node_id,
            },
            LockMessage::LockRejected {
                entity_id,
                requester: node_id,
                current_holder: Uuid::new_v4(),
            },
            LockMessage::LockHeartbeat {
                entity_id,
                holder: node_id,
            },
            LockMessage::LockRelease { entity_id, node_id },
            LockMessage::LockReleased { entity_id },
        ];

        for message in messages {
            let bytes = rkyv::to_bytes::<rkyv::rancor::Failure>(&message).map(|b| b.to_vec()).unwrap();
            let deserialized: LockMessage = rkyv::from_bytes::<LockMessage, rkyv::rancor::Failure>(&bytes).unwrap();
            assert_eq!(message, deserialized);
        }
    }

    #[test]
    fn test_heartbeat_renewal() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node_id = Uuid::new_v4();

        // Acquire lock
        registry.try_acquire(entity_id, node_id).unwrap();

        // Get initial heartbeat time
        let initial_heartbeat = registry.locks.get(&entity_id).unwrap().last_heartbeat;

        // Sleep a bit to ensure time difference
        std::thread::sleep(std::time::Duration::from_millis(10));

        // Renew heartbeat
        assert!(registry.renew_heartbeat(entity_id, node_id));

        // Check that heartbeat was updated
        let updated_heartbeat = registry.locks.get(&entity_id).unwrap().last_heartbeat;
        assert!(updated_heartbeat > initial_heartbeat);
    }

    #[test]
    fn test_heartbeat_wrong_node() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node1 = Uuid::new_v4();
        let node2 = Uuid::new_v4();

        // Node 1 acquires
        registry.try_acquire(entity_id, node1).unwrap();

        // Node 2 tries to renew heartbeat - should fail
        assert!(!registry.renew_heartbeat(entity_id, node2));
    }

    #[test]
    fn test_heartbeat_expiration() {
        let mut registry = EntityLockRegistry::new();
        let entity_id = Uuid::new_v4();
        let node_id = Uuid::new_v4();

        // Acquire with very short timeout
        registry.try_acquire(entity_id, node_id).unwrap();

        // Manually set timeout to 0 for testing
        if let Some(lock) = registry.locks.get_mut(&entity_id) {
            lock.timeout = Duration::from_secs(0);
        }

        // Should be detected as expired
        let expired = registry.get_expired_locks();
        assert_eq!(expired.len(), 1);
        assert_eq!(expired[0], entity_id);
    }
}