rivven_core/

consumer_group.rs

//! Consumer Group Coordination
//!
//! Kafka-compatible consumer groups with automatic partition rebalancing.
//!
//! ## Features
//!
//! - Multiple assignment strategies (range, round-robin, sticky)
//! - Automatic rebalancing on member join/leave
//! - Offset commit/fetch with Raft-based durability
//! - Heartbeat-based failure detection
//! - Generation-based protocol (prevents split-brain)
//! - **Static Membership (KIP-345)**: Stable consumer identity for K8s
//! - **Cooperative Rebalancing (KIP-429)**: Incremental partition reassignment
//!
//! ## Static Membership (KIP-345)
//!
//! Static membership allows consumers to maintain their identity across restarts,
//! preventing unnecessary rebalances in Kubernetes environments.
//!
//! ```text
//! Without static membership:
//!   Pod restart → new member_id → REBALANCE (stop-the-world)
//!
//! With static membership:
//!   Pod restart → same group.instance.id → REJOIN (keep assignments)
//! ```
//!
//! ### Configuration
//!
//! Set `group.instance.id` on the consumer to enable static membership:
//! - Instance ID must be unique within the group
//! - On rejoin, previous member is fenced and new member inherits assignments
//! - Session timeout still applies for failure detection
//!
//! ## Cooperative Rebalancing (KIP-429)
//!
//! Cooperative rebalancing eliminates "stop-the-world" rebalances by using
//! incremental partition reassignment:
//!
//! ```text
//! Eager (traditional):
//!   All consumers STOP → Reassign ALL → All consumers RESTART
//!   (service interruption during entire rebalance)
//!
//! Cooperative (KIP-429):
//!   1st rebalance: Revoke ONLY partitions that need moving
//!   2nd rebalance: Assign revoked partitions to new owners
//!   (consumers keep processing non-moving partitions)
//! ```
//!
//! ### Rebalance Protocol Selection
//!
//! - **Eager**: All partitions revoked immediately (Kafka default pre-2.4)
//! - **Cooperative**: Only changed partitions revoked (requires all members support)
//! - The group uses the highest common protocol (cooperative if all support it)
//!
//! ## Protocol
//!
//! 1. **JoinGroup**: Consumer joins group, triggers rebalance
//! 2. **SyncGroup**: Leader assigns partitions, followers receive assignments
//! 3. **Heartbeat**: Keep-alive messages, detect failures
//! 4. **LeaveGroup**: Graceful departure, triggers rebalance
//! 5. **OffsetCommit**: Persist consumer offsets
//! 6. **OffsetFetch**: Retrieve committed offsets
//!
//! ## State Machine
//!
//! ```text
//! Empty → PreparingRebalance → CompletingRebalance → Stable
//!   ↑                                                   │
//!   └───────────────────────────────────────────────────┘
//!           (member leave/timeout/join)
//! ```

use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::time::{Duration, SystemTime};

/// Unique consumer group identifier
pub type GroupId = String;

/// Unique consumer member identifier (generated by coordinator)
pub type MemberId = String;

/// Static group instance identifier (KIP-345)
///
/// When set, this provides a stable identity for the consumer across restarts.
/// The coordinator maps `group_instance_id` → `member_id` to recognize returning members.
pub type GroupInstanceId = String;

/// Consumer group state
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum GroupState {
    /// No members, awaiting first join
    Empty,

    /// Rebalance triggered, waiting for all members to rejoin
    PreparingRebalance,

    /// Leader computing assignments, followers waiting
    CompletingRebalance,

    /// All members have assignments, operating normally
    Stable,

    /// Group marked for deletion (all members left)
    Dead,
}

/// Partition assignment strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum AssignmentStrategy {
    /// Range: Assign contiguous partition ranges (default)
    #[default]
    Range,

    /// RoundRobin: Distribute partitions evenly in round-robin
    RoundRobin,

    /// Sticky: Minimize partition movement during rebalance
    Sticky,
}

/// Rebalance protocol (KIP-429)
///
/// Determines how partition reassignment is handled during rebalance.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum RebalanceProtocol {
    /// Eager: All partitions revoked during rebalance (stop-the-world)
    /// This is the traditional Kafka behavior before version 2.4
    #[default]
    Eager,

    /// Cooperative: Incremental rebalance, only moved partitions are revoked
    /// Requires all group members to support cooperative protocol
    Cooperative,
}

impl RebalanceProtocol {
    /// Select the highest common protocol supported by all members
    /// Returns Cooperative only if ALL members support it
    pub fn select_common(protocols: &[Self]) -> Self {
        if protocols.is_empty() {
            return RebalanceProtocol::Eager;
        }

        // Cooperative requires all members to support it
        if protocols
            .iter()
            .all(|p| *p == RebalanceProtocol::Cooperative)
        {
            RebalanceProtocol::Cooperative
        } else {
            RebalanceProtocol::Eager
        }
    }
}

/// Consumer group member
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct GroupMember {
    /// Unique member ID (generated by coordinator)
    pub member_id: MemberId,

    /// Static instance ID (KIP-345) - stable identity across restarts
    /// When set, enables static membership for this consumer
    pub group_instance_id: Option<GroupInstanceId>,

    /// Client-provided ID (for sticky assignment)
    pub client_id: String,

    /// Topics subscribed to
    pub subscriptions: Vec<String>,

    /// Current partition assignment
    pub assignment: Vec<PartitionAssignment>,

    /// Partitions pending revocation (cooperative protocol only)
    /// These partitions must be explicitly revoked before being reassigned
    pub pending_revocation: Vec<PartitionAssignment>,

    /// Last heartbeat timestamp (milliseconds since epoch for serialization)
    #[serde(
        serialize_with = "serialize_systemtime",
        deserialize_with = "deserialize_systemtime"
    )]
    pub last_heartbeat: SystemTime,

    /// Member metadata (client version, etc.)
    pub metadata: Vec<u8>,

    /// Is this a static member (has group_instance_id)
    pub is_static: bool,

    /// Supported rebalance protocols (KIP-429)
    /// The group selects the highest common protocol
    pub supported_protocols: Vec<RebalanceProtocol>,
}

// Helper functions for SystemTime serialization
fn serialize_systemtime<S>(time: &SystemTime, serializer: S) -> Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    let duration = time
        .duration_since(SystemTime::UNIX_EPOCH)
        .map_err(serde::ser::Error::custom)?;
    serializer.serialize_u128(duration.as_millis())
}

fn deserialize_systemtime<'de, D>(deserializer: D) -> Result<SystemTime, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let millis = u128::deserialize(deserializer)?;
    Ok(SystemTime::UNIX_EPOCH + std::time::Duration::from_millis(millis as u64))
}

/// Partition assignment for a member
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct PartitionAssignment {
    pub topic: String,
    pub partition: u32,
}

/// Result of a rebalance operation (KIP-429)
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RebalanceResult {
    /// Rebalance completed immediately (eager protocol or no revocations needed)
    Complete,

    /// Awaiting partition revocations (cooperative protocol)
    /// Contains the revocations requested and the pending final assignments
    AwaitingRevocations {
        /// Partitions each member needs to revoke
        revocations: HashMap<MemberId, Vec<PartitionAssignment>>,
        /// Final assignments to apply after revocations are acknowledged
        pending_assignments: HashMap<MemberId, Vec<PartitionAssignment>>,
    },
}

/// Consumer group metadata
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ConsumerGroup {
    /// Group identifier
    pub group_id: GroupId,

    /// Current state
    pub state: GroupState,

    /// Generation ID (increments on each rebalance)
    pub generation_id: u32,

    /// Leader member ID (computes assignments)
    pub leader_id: Option<MemberId>,

    /// Protocol name (for compatibility)
    pub protocol_name: String,

    /// Assignment strategy
    pub assignment_strategy: AssignmentStrategy,

    /// Selected rebalance protocol (KIP-429)
    /// Determined by the highest common protocol all members support
    pub rebalance_protocol: RebalanceProtocol,

    /// Active members
    pub members: HashMap<MemberId, GroupMember>,

    /// Static membership mapping: group_instance_id → member_id (KIP-345)
    /// Used to recognize returning static members without triggering rebalance
    pub static_members: HashMap<GroupInstanceId, MemberId>,

    /// Pending static members awaiting sync (instance_id → saved assignment)
    /// Stores assignments for static members that disconnected but haven't timed out
    pub pending_static_members: HashMap<GroupInstanceId, Vec<PartitionAssignment>>,

    /// Partitions awaiting revocation acknowledgment (cooperative protocol)
    /// Maps member_id → partitions that need to be revoked before reassignment
    pub awaiting_revocation: HashMap<MemberId, Vec<PartitionAssignment>>,

    /// Committed offsets (topic → partition → offset)
    pub offsets: HashMap<String, HashMap<u32, i64>>,

    /// Session timeout (member considered dead if no heartbeat)
    #[serde(
        serialize_with = "serialize_duration",
        deserialize_with = "deserialize_duration"
    )]
    pub session_timeout: Duration,

    /// Rebalance timeout (max time for rebalance completion)
    #[serde(
        serialize_with = "serialize_duration",
        deserialize_with = "deserialize_duration"
    )]
    pub rebalance_timeout: Duration,
}

// Helper functions for Duration serialization
fn serialize_duration<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    serializer.serialize_u64(duration.as_millis() as u64)
}

fn deserialize_duration<'de, D>(deserializer: D) -> Result<Duration, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let millis = u64::deserialize(deserializer)?;
    Ok(Duration::from_millis(millis))
}

impl ConsumerGroup {
    /// Create a new consumer group
    pub fn new(group_id: GroupId, session_timeout: Duration, rebalance_timeout: Duration) -> Self {
        Self {
            group_id,
            state: GroupState::Empty,
            generation_id: 0,
            leader_id: None,
            protocol_name: "consumer".to_string(),
            assignment_strategy: AssignmentStrategy::default(),
            rebalance_protocol: RebalanceProtocol::Eager,
            members: HashMap::new(),
            static_members: HashMap::new(),
            pending_static_members: HashMap::new(),
            awaiting_revocation: HashMap::new(),
            offsets: HashMap::new(),
            session_timeout,
            rebalance_timeout,
        }
    }

    /// Add a member to the group
    ///
    /// For static members (with `group_instance_id`):
    /// - If the instance ID is known and has a pending assignment, restore it without rebalance
    /// - If the instance ID is known but has an active member, fence the old member
    /// - Otherwise, treat as new member and trigger rebalance
    ///
    /// For dynamic members (without `group_instance_id`):
    /// - Always trigger rebalance
    pub fn add_member(
        &mut self,
        member_id: MemberId,
        client_id: String,
        subscriptions: Vec<String>,
        metadata: Vec<u8>,
    ) {
        self.add_member_full(
            member_id,
            None,
            client_id,
            subscriptions,
            metadata,
            vec![RebalanceProtocol::Eager],
        )
    }

    /// Add a member with optional static instance ID (KIP-345)
    pub fn add_member_with_instance_id(
        &mut self,
        member_id: MemberId,
        group_instance_id: Option<GroupInstanceId>,
        client_id: String,
        subscriptions: Vec<String>,
        metadata: Vec<u8>,
    ) {
        self.add_member_full(
            member_id,
            group_instance_id,
            client_id,
            subscriptions,
            metadata,
            vec![RebalanceProtocol::Eager],
        )
    }

    /// Add a member with full configuration including supported protocols (KIP-429)
    pub fn add_member_full(
        &mut self,
        member_id: MemberId,
        group_instance_id: Option<GroupInstanceId>,
        client_id: String,
        subscriptions: Vec<String>,
        metadata: Vec<u8>,
        supported_protocols: Vec<RebalanceProtocol>,
    ) {
        let is_static = group_instance_id.is_some();
        let supported_protocols = if supported_protocols.is_empty() {
            vec![RebalanceProtocol::Eager]
        } else {
            supported_protocols
        };

        // Check if this is a returning static member
        if let Some(ref instance_id) = group_instance_id {
            // Check for existing member with same instance ID (fencing)
            if let Some(old_member_id) = self.static_members.get(instance_id).cloned() {
                if old_member_id != member_id {
                    // Fence the old member - remove it without triggering rebalance
                    self.members.remove(&old_member_id);
                }
            }

            // Check for pending assignment (member returning after disconnect)
            if let Some(saved_assignment) = self.pending_static_members.remove(instance_id) {
                // Static member rejoining - restore assignment without rebalance
                let member = GroupMember {
                    member_id: member_id.clone(),
                    group_instance_id: Some(instance_id.clone()),
                    client_id,
                    subscriptions,
                    assignment: saved_assignment,
                    pending_revocation: Vec::new(),
                    last_heartbeat: SystemTime::now(),
                    metadata,
                    is_static: true,
                    supported_protocols: supported_protocols.clone(),
                };

                self.members.insert(member_id.clone(), member);
                self.static_members
                    .insert(instance_id.clone(), member_id.clone());

                // Update group's selected protocol
                self.update_rebalance_protocol();

                // First member becomes leader
                if self.leader_id.is_none() {
                    self.leader_id = Some(member_id);
                }

                // No rebalance needed - member restored with previous assignment
                if self.state == GroupState::Empty {
                    self.state = GroupState::Stable;
                }
                return;
            }

            // Register static member mapping
            self.static_members
                .insert(instance_id.clone(), member_id.clone());
        }

        // Create new member
        let member = GroupMember {
            member_id: member_id.clone(),
            group_instance_id,
            client_id,
            subscriptions,
            assignment: Vec::new(),
            pending_revocation: Vec::new(),
            last_heartbeat: SystemTime::now(),
            metadata,
            is_static,
            supported_protocols,
        };

        self.members.insert(member_id.clone(), member);

        // Update group's selected protocol
        self.update_rebalance_protocol();

        // First member becomes leader
        if self.leader_id.is_none() {
            self.leader_id = Some(member_id);
        }

        // Trigger rebalance (but only if not Empty state)
        if self.state != GroupState::Empty {
            self.transition_to_preparing_rebalance();
        } else if self.members.len() == 1 {
            // First member added - move to PreparingRebalance
            self.state = GroupState::PreparingRebalance;
        }
    }

    /// Check if a static member with the given instance ID exists
    pub fn has_static_member(&self, instance_id: &GroupInstanceId) -> bool {
        self.static_members.contains_key(instance_id)
    }

    /// Get member ID for a static instance ID
    pub fn get_member_for_instance(&self, instance_id: &GroupInstanceId) -> Option<&MemberId> {
        self.static_members.get(instance_id)
    }

    /// Fence a static member (remove and replace with new member)
    ///
    /// Called when a new member joins with the same group.instance.id as an existing member.
    /// This prevents split-brain scenarios.
    pub fn fence_static_member(&mut self, instance_id: &GroupInstanceId) -> Option<MemberId> {
        if let Some(old_member_id) = self.static_members.get(instance_id).cloned() {
            // Save the old member's assignment for potential restoration
            if let Some(old_member) = self.members.get(&old_member_id) {
                if !old_member.assignment.is_empty() {
                    self.pending_static_members
                        .insert(instance_id.clone(), old_member.assignment.clone());
                }
            }

            // Remove the old member
            self.members.remove(&old_member_id);

            // If leader left, elect new leader
            if self.leader_id.as_ref() == Some(&old_member_id) {
                self.leader_id = self.members.keys().next().cloned();
            }

            Some(old_member_id)
        } else {
            None
        }
    }

    /// Remove a member from the group
    ///
    /// For static members: Saves assignment for potential restoration (no rebalance until timeout)
    /// For dynamic members: Triggers immediate rebalance
    pub fn remove_member(&mut self, member_id: &MemberId) -> bool {
        if let Some(member) = self.members.remove(member_id) {
            // Handle static member removal differently
            if member.is_static {
                if let Some(ref instance_id) = member.group_instance_id {
                    // Save assignment for potential restoration
                    if !member.assignment.is_empty() {
                        self.pending_static_members
                            .insert(instance_id.clone(), member.assignment);
                    }
                    // Keep the static_members mapping until timeout
                    // This allows the member to rejoin without rebalance
                }
            } else {
                // Dynamic member - trigger rebalance immediately
                if !self.members.is_empty() {
                    self.transition_to_preparing_rebalance();
                }
            }

            // If leader left, elect new leader
            if self.leader_id.as_ref() == Some(member_id) {
                self.leader_id = self.members.keys().next().cloned();
            }

            // If no members left, transition to Empty
            if self.members.is_empty() {
                self.state = GroupState::Empty;
                self.generation_id = 0;
                self.leader_id = None;
                // Clear all static member mappings
                self.static_members.clear();
                self.pending_static_members.clear();
            }

            true
        } else {
            false
        }
    }

    /// Remove a static member permanently (triggers rebalance)
    ///
    /// Called when a static member's session times out or explicitly leaves.
    pub fn remove_static_member(&mut self, instance_id: &GroupInstanceId) -> bool {
        if let Some(member_id) = self.static_members.remove(instance_id) {
            self.pending_static_members.remove(instance_id);

            if self.members.remove(&member_id).is_some() {
                // If leader left, elect new leader
                if self.leader_id.as_ref() == Some(&member_id) {
                    self.leader_id = self.members.keys().next().cloned();
                }

                // If no members left, transition to Empty
                if self.members.is_empty() {
                    self.state = GroupState::Empty;
                    self.generation_id = 0;
                    self.leader_id = None;
                    self.static_members.clear();
                    self.pending_static_members.clear();
                } else {
                    // Trigger rebalance
                    self.transition_to_preparing_rebalance();
                }

                return true;
            }
        }

        // Also check pending members
        if self.pending_static_members.remove(instance_id).is_some() {
            self.static_members.remove(instance_id);
            // Pending member timeout doesn't trigger rebalance since they weren't active
            return true;
        }

        false
    }

    /// Update member heartbeat
    pub fn heartbeat(&mut self, member_id: &MemberId) -> Result<(), String> {
        if let Some(member) = self.members.get_mut(member_id) {
            member.last_heartbeat = SystemTime::now();
            Ok(())
        } else {
            Err(format!("Unknown member: {}", member_id))
        }
    }

    /// Check for timed-out members
    ///
    /// For static members: Only removes if instance ID hasn't rejoined within session timeout
    /// For dynamic members: Removes immediately on timeout
    pub fn check_timeouts(&mut self) -> Vec<MemberId> {
        let now = SystemTime::now();
        let mut timed_out = Vec::new();
        let mut static_timeouts: Vec<GroupInstanceId> = Vec::new();

        // Check active members
        for (member_id, member) in &self.members {
            if let Ok(elapsed) = now.duration_since(member.last_heartbeat) {
                if elapsed > self.session_timeout {
                    timed_out.push(member_id.clone());

                    // Track static member timeouts separately
                    if let Some(ref instance_id) = member.group_instance_id {
                        static_timeouts.push(instance_id.clone());
                    }
                }
            }
        }

        // Remove timed-out dynamic members immediately
        for member_id in &timed_out {
            let member = self.members.get(member_id);
            if let Some(m) = member {
                if !m.is_static {
                    self.remove_member(member_id);
                }
            }
        }

        // For static members, save assignment but remove from active
        for instance_id in &static_timeouts {
            // The remove_static_member call will trigger rebalance
            self.remove_static_member(instance_id);
        }

        timed_out
    }

    /// Check for timed-out pending static members
    ///
    /// Pending static members are those that disconnected but haven't timed out yet.
    /// Their assignments are preserved for potential restoration.
    pub fn check_pending_static_timeouts(
        &mut self,
        _pending_timeout: Duration,
    ) -> Vec<GroupInstanceId> {
        // This would typically be called periodically to clean up
        // pending static members that never rejoined.
        // For now, we just return an empty vec as we don't track
        // pending timestamps separately (could be enhanced).
        Vec::new()
    }

    // =========================================================================
    // Cooperative Rebalancing (KIP-429)
    // =========================================================================

    /// Update the group's selected rebalance protocol based on member support
    ///
    /// The group uses cooperative protocol only if ALL members support it.
    /// Otherwise, falls back to eager protocol.
    fn update_rebalance_protocol(&mut self) {
        let protocols: Vec<RebalanceProtocol> = self
            .members
            .values()
            .flat_map(|m| {
                // Find the highest protocol each member supports
                if m.supported_protocols
                    .contains(&RebalanceProtocol::Cooperative)
                {
                    Some(RebalanceProtocol::Cooperative)
                } else {
                    Some(RebalanceProtocol::Eager)
                }
            })
            .collect();

        self.rebalance_protocol = RebalanceProtocol::select_common(&protocols);
    }

    /// Check if the group uses cooperative rebalancing
    pub fn is_cooperative(&self) -> bool {
        self.rebalance_protocol == RebalanceProtocol::Cooperative
    }

    /// Compute partitions that need to be revoked for cooperative rebalance
    ///
    /// Returns a map of member_id → partitions to revoke.
    /// Only partitions that are moving to a different consumer are revoked.
    pub fn compute_revocations(
        &self,
        new_assignments: &HashMap<MemberId, Vec<PartitionAssignment>>,
    ) -> HashMap<MemberId, Vec<PartitionAssignment>> {
        let mut revocations: HashMap<MemberId, Vec<PartitionAssignment>> = HashMap::new();

        for (member_id, member) in &self.members {
            let new_assignment = new_assignments.get(member_id);

            // Find partitions that are being taken away
            let mut to_revoke = Vec::new();
            for partition in &member.assignment {
                let still_assigned = new_assignment
                    .map(|a| a.contains(partition))
                    .unwrap_or(false);

                if !still_assigned {
                    to_revoke.push(partition.clone());
                }
            }

            if !to_revoke.is_empty() {
                revocations.insert(member_id.clone(), to_revoke);
            }
        }

        revocations
    }

    /// Start cooperative rebalance phase 1: Request revocations
    ///
    /// In cooperative rebalancing, we don't immediately reassign all partitions.
    /// Instead, we first ask members to revoke partitions that need moving.
    pub fn request_revocations(
        &mut self,
        revocations: HashMap<MemberId, Vec<PartitionAssignment>>,
    ) {
        // Store pending revocations
        self.awaiting_revocation = revocations.clone();

        // Mark partitions as pending revocation on each member
        for (member_id, partitions) in revocations {
            if let Some(member) = self.members.get_mut(&member_id) {
                member.pending_revocation = partitions;
            }
        }

        // Transition to waiting for revocation acknowledgments
        // (reuse CompletingRebalance state for this phase)
        self.state = GroupState::CompletingRebalance;
    }

    /// Acknowledge revocation from a member (cooperative protocol)
    ///
    /// Called when a member confirms it has stopped processing revoked partitions.
    /// Returns true if all pending revocations have been acknowledged.
    pub fn acknowledge_revocation(&mut self, member_id: &MemberId) -> bool {
        // Remove from awaiting list
        self.awaiting_revocation.remove(member_id);

        // Clear pending revocation on member
        if let Some(member) = self.members.get_mut(member_id) {
            // Remove revoked partitions from assignment
            let revoked: HashSet<_> = member.pending_revocation.drain(..).collect();
            member.assignment.retain(|p| !revoked.contains(p));
        }

        // Check if all revocations are complete
        self.awaiting_revocation.is_empty()
    }

    /// Check if there are pending revocations
    pub fn has_pending_revocations(&self) -> bool {
        !self.awaiting_revocation.is_empty()
    }

    /// Get pending revocations for a member
    pub fn get_pending_revocations(&self, member_id: &MemberId) -> Vec<PartitionAssignment> {
        self.members
            .get(member_id)
            .map(|m| m.pending_revocation.clone())
            .unwrap_or_default()
    }

    /// Complete cooperative rebalance with final assignments
    ///
    /// Called after all revocations are acknowledged. Assigns partitions
    /// that were revoked to their new owners.
    pub fn complete_cooperative_rebalance(
        &mut self,
        final_assignments: HashMap<MemberId, Vec<PartitionAssignment>>,
    ) {
        // Merge new assignments with existing (non-revoked) assignments
        for (member_id, new_partitions) in final_assignments {
            if let Some(member) = self.members.get_mut(&member_id) {
                // Add new partitions to existing assignment
                for partition in new_partitions {
                    if !member.assignment.contains(&partition) {
                        member.assignment.push(partition);
                    }
                }
            }
        }

        // Clear revocation state
        self.awaiting_revocation.clear();

        // Increment generation
        self.generation_id += 1;
        self.state = GroupState::Stable;
    }

    /// Perform a full rebalance (handles both eager and cooperative)
    ///
    /// For eager: Immediately assigns all partitions
    /// For cooperative: Initiates two-phase revocation process
    pub fn rebalance_with_strategy(
        &mut self,
        new_assignments: HashMap<MemberId, Vec<PartitionAssignment>>,
    ) -> RebalanceResult {
        match self.rebalance_protocol {
            RebalanceProtocol::Eager => {
                // Eager: Immediately replace all assignments
                self.complete_rebalance(new_assignments);
                RebalanceResult::Complete
            }
            RebalanceProtocol::Cooperative => {
                // Cooperative: Check if revocations are needed
                let revocations = self.compute_revocations(&new_assignments);

                if revocations.is_empty() {
                    // No revocations needed - can complete immediately
                    self.complete_cooperative_rebalance(new_assignments);
                    RebalanceResult::Complete
                } else {
                    // Need to revoke first
                    self.request_revocations(revocations.clone());
                    RebalanceResult::AwaitingRevocations {
                        revocations,
                        pending_assignments: new_assignments,
                    }
                }
            }
        }
    }

    /// Transition to PreparingRebalance state
    fn transition_to_preparing_rebalance(&mut self) {
        if self.state != GroupState::Empty {
            self.state = GroupState::PreparingRebalance;
        }
    }

    /// Complete rebalance with assignments
    pub fn complete_rebalance(&mut self, assignments: HashMap<MemberId, Vec<PartitionAssignment>>) {
        // Update member assignments
        for (member_id, partitions) in assignments {
            if let Some(member) = self.members.get_mut(&member_id) {
                member.assignment = partitions;
            }
        }

        // Increment generation
        self.generation_id += 1;
        self.state = GroupState::Stable;
    }

    /// Commit offset for a partition
    pub fn commit_offset(&mut self, topic: &str, partition: u32, offset: i64) {
        self.offsets
            .entry(topic.to_string())
            .or_default()
            .insert(partition, offset);
    }

    /// Fetch committed offset for a partition
    pub fn fetch_offset(&self, topic: &str, partition: u32) -> Option<i64> {
        self.offsets.get(topic)?.get(&partition).copied()
    }

    /// Get all partition assignments for the group
    pub fn all_assignments(&self) -> HashMap<PartitionAssignment, MemberId> {
        let mut assignments = HashMap::new();
        for (member_id, member) in &self.members {
            for partition in &member.assignment {
                assignments.insert(partition.clone(), member_id.clone());
            }
        }
        assignments
    }
}

/// Partition assignment strategies
pub mod assignment {
    use super::*;

    /// Range assignment: Assign contiguous partition ranges
    ///
    /// Example: 3 consumers, 10 partitions
    /// - Consumer 0: partitions 0-3
    /// - Consumer 1: partitions 4-6
    /// - Consumer 2: partitions 7-9
    pub fn range_assignment(
        members: &[MemberId],
        topic_partitions: &HashMap<String, u32>,
    ) -> HashMap<MemberId, Vec<PartitionAssignment>> {
        let mut assignments: HashMap<MemberId, Vec<PartitionAssignment>> = HashMap::new();

        if members.is_empty() {
            return assignments;
        }

        for (topic, partition_count) in topic_partitions {
            let partitions_per_member = partition_count / members.len() as u32;
            let extra_partitions = partition_count % members.len() as u32;

            let mut current_partition = 0;

            for (idx, member_id) in members.iter().enumerate() {
                let mut member_partitions = partitions_per_member;
                if (idx as u32) < extra_partitions {
                    member_partitions += 1;
                }

                for _ in 0..member_partitions {
                    assignments
                        .entry(member_id.clone())
                        .or_default()
                        .push(PartitionAssignment {
                            topic: topic.clone(),
                            partition: current_partition,
                        });
                    current_partition += 1;
                }
            }
        }

        assignments
    }

    /// Round-robin assignment: Distribute partitions evenly
    ///
    /// Example: 3 consumers, 10 partitions
    /// - Consumer 0: partitions 0, 3, 6, 9
    /// - Consumer 1: partitions 1, 4, 7
    /// - Consumer 2: partitions 2, 5, 8
    pub fn round_robin_assignment(
        members: &[MemberId],
        topic_partitions: &HashMap<String, u32>,
    ) -> HashMap<MemberId, Vec<PartitionAssignment>> {
        let mut assignments: HashMap<MemberId, Vec<PartitionAssignment>> = HashMap::new();

        if members.is_empty() {
            return assignments;
        }

        let mut member_idx = 0;

        for (topic, partition_count) in topic_partitions {
            for partition in 0..*partition_count {
                assignments
                    .entry(members[member_idx].clone())
                    .or_default()
                    .push(PartitionAssignment {
                        topic: topic.clone(),
                        partition,
                    });

                member_idx = (member_idx + 1) % members.len();
            }
        }

        assignments
    }

    /// Sticky assignment: Minimize partition movement during rebalance
    ///
    /// Preserves previous assignments as much as possible.
    pub fn sticky_assignment(
        members: &[MemberId],
        topic_partitions: &HashMap<String, u32>,
        previous_assignments: &HashMap<MemberId, Vec<PartitionAssignment>>,
    ) -> HashMap<MemberId, Vec<PartitionAssignment>> {
        let mut assignments: HashMap<MemberId, Vec<PartitionAssignment>> = HashMap::new();

        if members.is_empty() {
            return assignments;
        }

        // Collect all partitions
        let mut all_partitions = Vec::new();
        for (topic, partition_count) in topic_partitions {
            for partition in 0..*partition_count {
                all_partitions.push(PartitionAssignment {
                    topic: topic.clone(),
                    partition,
                });
            }
        }

        // Track which partitions are assigned
        let mut assigned: HashSet<PartitionAssignment> = HashSet::new();

        // Step 1: Preserve previous assignments for existing members
        for member_id in members {
            if let Some(prev_partitions) = previous_assignments.get(member_id) {
                let valid_partitions: Vec<_> = prev_partitions
                    .iter()
                    .filter(|p| all_partitions.contains(p) && !assigned.contains(p))
                    .cloned()
                    .collect();

                for partition in &valid_partitions {
                    assigned.insert(partition.clone());
                }

                assignments.insert(member_id.clone(), valid_partitions);
            }
        }

        // Step 2: Distribute unassigned partitions using round-robin
        let unassigned: Vec<_> = all_partitions
            .into_iter()
            .filter(|p| !assigned.contains(p))
            .collect();

        let mut member_idx = 0;
        for partition in unassigned {
            assignments
                .entry(members[member_idx].clone())
                .or_default()
                .push(partition);

            member_idx = (member_idx + 1) % members.len();
        }

        assignments
    }
}

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

    #[test]
    fn test_consumer_group_creation() {
        let group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        assert_eq!(group.group_id, "test-group");
        assert_eq!(group.state, GroupState::Empty);
        assert_eq!(group.generation_id, 0);
        assert!(group.leader_id.is_none());
        assert!(group.members.is_empty());
    }

    #[test]
    fn test_add_first_member_becomes_leader() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        group.add_member(
            "member-1".to_string(),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        assert_eq!(group.members.len(), 1);
        assert_eq!(group.leader_id, Some("member-1".to_string()));
        assert_eq!(group.state, GroupState::PreparingRebalance);
    }

    #[test]
    fn test_remove_member_triggers_rebalance() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        group.add_member(
            "member-1".to_string(),
            "client-1".to_string(),
            vec![],
            vec![],
        );
        group.add_member(
            "member-2".to_string(),
            "client-2".to_string(),
            vec![],
            vec![],
        );

        group.state = GroupState::Stable;

        group.remove_member(&"member-2".to_string());

        assert_eq!(group.members.len(), 1);
        assert_eq!(group.state, GroupState::PreparingRebalance);
    }

    #[test]
    fn test_remove_last_member_transitions_to_empty() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        group.add_member(
            "member-1".to_string(),
            "client-1".to_string(),
            vec![],
            vec![],
        );
        group.remove_member(&"member-1".to_string());

        assert_eq!(group.state, GroupState::Empty);
        assert_eq!(group.generation_id, 0);
        assert!(group.leader_id.is_none());
    }

    #[test]
    fn test_offset_commit_and_fetch() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        group.commit_offset("topic-1", 0, 100);
        group.commit_offset("topic-1", 1, 200);
        group.commit_offset("topic-2", 0, 300);

        assert_eq!(group.fetch_offset("topic-1", 0), Some(100));
        assert_eq!(group.fetch_offset("topic-1", 1), Some(200));
        assert_eq!(group.fetch_offset("topic-2", 0), Some(300));
        assert_eq!(group.fetch_offset("topic-1", 2), None);
    }

    #[test]
    fn test_range_assignment() {
        let members = vec!["m1".to_string(), "m2".to_string(), "m3".to_string()];
        let mut topic_partitions = HashMap::new();
        topic_partitions.insert("topic-1".to_string(), 10);

        let assignments = assignment::range_assignment(&members, &topic_partitions);

        // m1: 0-3 (4 partitions)
        // m2: 4-6 (3 partitions)
        // m3: 7-9 (3 partitions)
        assert_eq!(assignments.get("m1").unwrap().len(), 4);
        assert_eq!(assignments.get("m2").unwrap().len(), 3);
        assert_eq!(assignments.get("m3").unwrap().len(), 3);

        // Verify m1 has partitions 0-3
        let m1_partitions: Vec<u32> = assignments
            .get("m1")
            .unwrap()
            .iter()
            .map(|p| p.partition)
            .collect();
        assert_eq!(m1_partitions, vec![0, 1, 2, 3]);
    }

    #[test]
    fn test_round_robin_assignment() {
        let members = vec!["m1".to_string(), "m2".to_string(), "m3".to_string()];
        let mut topic_partitions = HashMap::new();
        topic_partitions.insert("topic-1".to_string(), 10);

        let assignments = assignment::round_robin_assignment(&members, &topic_partitions);

        // m1: 0, 3, 6, 9 (4 partitions)
        // m2: 1, 4, 7 (3 partitions)
        // m3: 2, 5, 8 (3 partitions)
        assert_eq!(assignments.get("m1").unwrap().len(), 4);
        assert_eq!(assignments.get("m2").unwrap().len(), 3);
        assert_eq!(assignments.get("m3").unwrap().len(), 3);

        let m1_partitions: Vec<u32> = assignments
            .get("m1")
            .unwrap()
            .iter()
            .map(|p| p.partition)
            .collect();
        assert_eq!(m1_partitions, vec![0, 3, 6, 9]);
    }

    #[test]
    fn test_sticky_assignment_preserves_assignments() {
        let members = vec!["m1".to_string(), "m2".to_string()];
        let mut topic_partitions = HashMap::new();
        topic_partitions.insert("topic-1".to_string(), 4);

        // Previous assignment: m1=[0,1], m2=[2,3]
        let mut previous = HashMap::new();
        previous.insert(
            "m1".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 0,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 1,
                },
            ],
        );
        previous.insert(
            "m2".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 2,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 3,
                },
            ],
        );

        let assignments = assignment::sticky_assignment(&members, &topic_partitions, &previous);

        // Should preserve previous assignments
        assert_eq!(assignments.get("m1").unwrap().len(), 2);
        assert_eq!(assignments.get("m2").unwrap().len(), 2);

        let m1_partitions: HashSet<u32> = assignments
            .get("m1")
            .unwrap()
            .iter()
            .map(|p| p.partition)
            .collect();
        assert!(m1_partitions.contains(&0));
        assert!(m1_partitions.contains(&1));
    }

    #[test]
    fn test_sticky_assignment_redistributes_on_new_member() {
        let members = vec!["m1".to_string(), "m2".to_string(), "m3".to_string()];
        let mut topic_partitions = HashMap::new();
        topic_partitions.insert("topic-1".to_string(), 6);

        // Previous: m1=[0,1,2], m2=[3,4,5], m3 is new
        let mut previous = HashMap::new();
        previous.insert(
            "m1".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 0,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 1,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 2,
                },
            ],
        );
        previous.insert(
            "m2".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 3,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 4,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 5,
                },
            ],
        );

        let assignments = assignment::sticky_assignment(&members, &topic_partitions, &previous);

        // Sticky assignment preserves previous, so m1 keeps 0,1,2 and m2 keeps 3,4,5
        // m3 gets none from previous. Then redistribution happens.
        // Total: 6 partitions across 3 members
        let total_assigned: usize = assignments.values().map(|v| v.len()).sum();
        assert_eq!(total_assigned, 6, "All 6 partitions should be assigned");

        // m1 and m2 should keep some of their previous assignments
        let m1_partitions: HashSet<u32> = assignments
            .get("m1")
            .unwrap()
            .iter()
            .map(|p| p.partition)
            .collect();
        let m2_partitions: HashSet<u32> = assignments
            .get("m2")
            .unwrap()
            .iter()
            .map(|p| p.partition)
            .collect();

        // Should have some overlap with previous
        let m1_kept = m1_partitions.iter().filter(|p| **p <= 2).count();
        let m2_kept = m2_partitions.iter().filter(|p| **p >= 3).count();

        assert!(
            m1_kept > 0 || m2_kept > 0,
            "Sticky assignment should preserve some assignments"
        );
    }

    // =========================================================================
    // Static Membership (KIP-345) Tests
    // =========================================================================

    #[test]
    fn test_static_member_add() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add static member with instance ID
        group.add_member_with_instance_id(
            "member-1".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        assert_eq!(group.members.len(), 1);
        assert!(group.has_static_member(&"instance-1".to_string()));
        assert_eq!(
            group.get_member_for_instance(&"instance-1".to_string()),
            Some(&"member-1".to_string())
        );

        let member = group.members.get("member-1").unwrap();
        assert!(member.is_static);
        assert_eq!(member.group_instance_id, Some("instance-1".to_string()));
    }

    #[test]
    fn test_static_member_rejoin_no_rebalance() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add two members
        group.add_member_with_instance_id(
            "member-1".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );
        group.add_member(
            "member-2".to_string(),
            "client-2".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        // Complete rebalance with assignments
        let mut assignments = HashMap::new();
        assignments.insert(
            "member-1".to_string(),
            vec![PartitionAssignment {
                topic: "topic-1".to_string(),
                partition: 0,
            }],
        );
        assignments.insert(
            "member-2".to_string(),
            vec![PartitionAssignment {
                topic: "topic-1".to_string(),
                partition: 1,
            }],
        );
        group.complete_rebalance(assignments);
        assert_eq!(group.state, GroupState::Stable);
        let gen_before = group.generation_id;

        // Remove static member (simulates disconnect)
        group.remove_member(&"member-1".to_string());

        // Assignment should be saved for restoration
        assert!(group.pending_static_members.contains_key("instance-1"));

        // Static member rejoins with same instance ID but new member ID
        group.add_member_with_instance_id(
            "member-1-new".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        // Should restore assignment without rebalance
        let member = group.members.get("member-1-new").unwrap();
        assert_eq!(member.assignment.len(), 1);
        assert_eq!(member.assignment[0].partition, 0);

        // Generation should not have changed (no rebalance)
        assert_eq!(group.generation_id, gen_before);

        // Pending assignment should be cleared
        assert!(!group.pending_static_members.contains_key("instance-1"));
    }

    #[test]
    fn test_static_member_fencing() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add static member
        group.add_member_with_instance_id(
            "member-1".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        assert!(group.members.contains_key("member-1"));

        // New member joins with same instance ID (fencing)
        group.add_member_with_instance_id(
            "member-1-new".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        // Old member should be fenced (removed)
        assert!(!group.members.contains_key("member-1"));
        // New member should be active
        assert!(group.members.contains_key("member-1-new"));
        // Instance mapping should point to new member
        assert_eq!(
            group.get_member_for_instance(&"instance-1".to_string()),
            Some(&"member-1-new".to_string())
        );
    }

    #[test]
    fn test_dynamic_member_removal_triggers_rebalance() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add one static and one dynamic member
        group.add_member_with_instance_id(
            "static-member".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );
        group.add_member(
            "dynamic-member".to_string(),
            "client-2".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        group.state = GroupState::Stable;

        // Remove dynamic member
        group.remove_member(&"dynamic-member".to_string());

        // Should trigger rebalance for dynamic member removal
        assert_eq!(group.state, GroupState::PreparingRebalance);
    }

    #[test]
    fn test_static_member_timeout_triggers_rebalance() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_millis(10), // Short timeout for testing
            Duration::from_secs(60),
        );

        // Add static member
        group.add_member_with_instance_id(
            "member-1".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );
        group.add_member(
            "member-2".to_string(),
            "client-2".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        group.state = GroupState::Stable;

        // Simulate timeout by removing static member permanently
        group.remove_static_member(&"instance-1".to_string());

        // Should trigger rebalance
        assert_eq!(group.state, GroupState::PreparingRebalance);
        // Static member mapping should be cleared
        assert!(!group.has_static_member(&"instance-1".to_string()));
    }

    #[test]
    fn test_mixed_static_and_dynamic_members() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add mix of static and dynamic members
        group.add_member_with_instance_id(
            "static-1".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );
        group.add_member_with_instance_id(
            "static-2".to_string(),
            Some("instance-2".to_string()),
            "client-2".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );
        group.add_member(
            "dynamic-1".to_string(),
            "client-3".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        assert_eq!(group.members.len(), 3);
        assert_eq!(group.static_members.len(), 2);

        // Static members should be identified correctly
        let static1 = group.members.get("static-1").unwrap();
        let static2 = group.members.get("static-2").unwrap();
        let dynamic1 = group.members.get("dynamic-1").unwrap();

        assert!(static1.is_static);
        assert!(static2.is_static);
        assert!(!dynamic1.is_static);
    }

    #[test]
    fn test_all_members_leave_clears_static_mappings() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        group.add_member_with_instance_id(
            "member-1".to_string(),
            Some("instance-1".to_string()),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        group.remove_member(&"member-1".to_string());

        // Even though member left, pending assignment is preserved
        // But let's also remove the static mapping to test full cleanup
        group.remove_static_member(&"instance-1".to_string());

        // Group should be empty and all mappings cleared
        assert_eq!(group.state, GroupState::Empty);
        assert!(group.static_members.is_empty());
        assert!(group.pending_static_members.is_empty());
    }

    // =========================================================================
    // Cooperative Rebalancing (KIP-429) Tests
    // =========================================================================

    #[test]
    fn test_rebalance_protocol_selection_all_eager() {
        let protocols = vec![RebalanceProtocol::Eager, RebalanceProtocol::Eager];
        assert_eq!(
            RebalanceProtocol::select_common(&protocols),
            RebalanceProtocol::Eager
        );
    }

    #[test]
    fn test_rebalance_protocol_selection_all_cooperative() {
        let protocols = vec![
            RebalanceProtocol::Cooperative,
            RebalanceProtocol::Cooperative,
        ];
        assert_eq!(
            RebalanceProtocol::select_common(&protocols),
            RebalanceProtocol::Cooperative
        );
    }

    #[test]
    fn test_rebalance_protocol_selection_mixed() {
        // Mixed protocols should fall back to eager
        let protocols = vec![RebalanceProtocol::Cooperative, RebalanceProtocol::Eager];
        assert_eq!(
            RebalanceProtocol::select_common(&protocols),
            RebalanceProtocol::Eager
        );
    }

    #[test]
    fn test_cooperative_member_add_updates_protocol() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add member with cooperative support
        group.add_member_full(
            "member-1".to_string(),
            None,
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
            vec![RebalanceProtocol::Cooperative],
        );

        // Single cooperative member means cooperative protocol
        assert!(group.is_cooperative());

        // Add member with only eager support
        group.add_member_full(
            "member-2".to_string(),
            None,
            "client-2".to_string(),
            vec!["topic-1".to_string()],
            vec![],
            vec![RebalanceProtocol::Eager],
        );

        // Mixed members means eager protocol
        assert!(!group.is_cooperative());
        assert_eq!(group.rebalance_protocol, RebalanceProtocol::Eager);
    }

    #[test]
    fn test_compute_revocations() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add two cooperative members
        group.add_member_full(
            "member-1".to_string(),
            None,
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
            vec![RebalanceProtocol::Cooperative],
        );
        group.add_member_full(
            "member-2".to_string(),
            None,
            "client-2".to_string(),
            vec!["topic-1".to_string()],
            vec![],
            vec![RebalanceProtocol::Cooperative],
        );

        // Assign partitions
        let mut initial = HashMap::new();
        initial.insert(
            "member-1".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 0,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 1,
                },
            ],
        );
        initial.insert(
            "member-2".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 2,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 3,
                },
            ],
        );
        group.complete_rebalance(initial);

        // New assignment moves partition 1 from member-1 to member-2
        let mut new_assignment = HashMap::new();
        new_assignment.insert(
            "member-1".to_string(),
            vec![PartitionAssignment {
                topic: "topic-1".to_string(),
                partition: 0,
            }],
        );
        new_assignment.insert(
            "member-2".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 1,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 2,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 3,
                },
            ],
        );

        let revocations = group.compute_revocations(&new_assignment);

        // Only member-1 should have revocations (partition 1)
        assert!(revocations.contains_key("member-1"));
        assert!(!revocations.contains_key("member-2"));

        let m1_revoked = revocations.get("member-1").unwrap();
        assert_eq!(m1_revoked.len(), 1);
        assert_eq!(m1_revoked[0].partition, 1);
    }

    #[test]
    fn test_cooperative_rebalance_two_phase() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add two cooperative members
        group.add_member_full(
            "member-1".to_string(),
            None,
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
            vec![RebalanceProtocol::Cooperative],
        );
        group.add_member_full(
            "member-2".to_string(),
            None,
            "client-2".to_string(),
            vec!["topic-1".to_string()],
            vec![],
            vec![RebalanceProtocol::Cooperative],
        );

        assert!(group.is_cooperative());

        // Initial assignment
        let mut initial = HashMap::new();
        initial.insert(
            "member-1".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 0,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 1,
                },
            ],
        );
        initial.insert("member-2".to_string(), vec![]);
        group.complete_rebalance(initial);

        let gen_before = group.generation_id;

        // New member-2 should get partition 1
        let mut new_assignment = HashMap::new();
        new_assignment.insert(
            "member-1".to_string(),
            vec![PartitionAssignment {
                topic: "topic-1".to_string(),
                partition: 0,
            }],
        );
        new_assignment.insert(
            "member-2".to_string(),
            vec![PartitionAssignment {
                topic: "topic-1".to_string(),
                partition: 1,
            }],
        );

        // Phase 1: Request revocations
        let result = group.rebalance_with_strategy(new_assignment.clone());

        match result {
            RebalanceResult::AwaitingRevocations {
                revocations,
                pending_assignments: _,
            } => {
                // Should have revocation for member-1
                assert!(revocations.contains_key("member-1"));
                assert!(group.has_pending_revocations());
                assert_eq!(group.state, GroupState::CompletingRebalance);
            }
            RebalanceResult::Complete => panic!("Expected AwaitingRevocations"),
        }

        // Member-1 still has both partitions (hasn't acked revocation yet)
        let m1 = group.members.get("member-1").unwrap();
        assert_eq!(m1.assignment.len(), 2);
        assert_eq!(m1.pending_revocation.len(), 1);

        // Phase 2: Acknowledge revocation
        let all_acked = group.acknowledge_revocation(&"member-1".to_string());
        assert!(all_acked);

        // After ack, member-1 should only have partition 0
        let m1 = group.members.get("member-1").unwrap();
        assert_eq!(m1.assignment.len(), 1);
        assert_eq!(m1.assignment[0].partition, 0);

        // Complete with final assignments
        group.complete_cooperative_rebalance(new_assignment);

        // Member-2 should now have partition 1
        let m2 = group.members.get("member-2").unwrap();
        assert_eq!(m2.assignment.len(), 1);
        assert_eq!(m2.assignment[0].partition, 1);

        // Generation should have incremented
        assert_eq!(group.generation_id, gen_before + 1);
        assert_eq!(group.state, GroupState::Stable);
    }

    #[test]
    fn test_eager_rebalance_immediate() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add two eager members
        group.add_member(
            "member-1".to_string(),
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );
        group.add_member(
            "member-2".to_string(),
            "client-2".to_string(),
            vec!["topic-1".to_string()],
            vec![],
        );

        assert!(!group.is_cooperative());

        // Eager rebalance should complete immediately
        let mut new_assignment = HashMap::new();
        new_assignment.insert(
            "member-1".to_string(),
            vec![PartitionAssignment {
                topic: "topic-1".to_string(),
                partition: 0,
            }],
        );
        new_assignment.insert(
            "member-2".to_string(),
            vec![PartitionAssignment {
                topic: "topic-1".to_string(),
                partition: 1,
            }],
        );

        let result = group.rebalance_with_strategy(new_assignment);

        assert_eq!(result, RebalanceResult::Complete);
        assert_eq!(group.state, GroupState::Stable);
        assert!(!group.has_pending_revocations());
    }

    #[test]
    fn test_cooperative_no_revocations_needed() {
        let mut group = ConsumerGroup::new(
            "test-group".to_string(),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );

        // Add cooperative member
        group.add_member_full(
            "member-1".to_string(),
            None,
            "client-1".to_string(),
            vec!["topic-1".to_string()],
            vec![],
            vec![RebalanceProtocol::Cooperative],
        );

        // Assign partition 0
        let mut initial = HashMap::new();
        initial.insert(
            "member-1".to_string(),
            vec![PartitionAssignment {
                topic: "topic-1".to_string(),
                partition: 0,
            }],
        );
        group.complete_rebalance(initial);

        // Add partition 1 (no revocation needed, just addition)
        let mut new_assignment = HashMap::new();
        new_assignment.insert(
            "member-1".to_string(),
            vec![
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 0,
                },
                PartitionAssignment {
                    topic: "topic-1".to_string(),
                    partition: 1,
                },
            ],
        );

        // Should complete immediately since no partitions need revocation
        let result = group.rebalance_with_strategy(new_assignment);

        assert_eq!(result, RebalanceResult::Complete);
        assert_eq!(group.state, GroupState::Stable);

        let m1 = group.members.get("member-1").unwrap();
        assert_eq!(m1.assignment.len(), 2);
    }
}