//! Range management for dynamic shard split and merge operations.
//!
//! This module handles the automatic rebalancing of shards based on data size:
//!
//! - Shards exceeding 64MB are split into two
//! - Adjacent shards under 20MB combined are merged
//! - Meta-range changes are broadcast to all nodes via gossip

use std::sync::Arc;
use tracing::{info, instrument, warn};

use crate::config::ShardingConfig;
use crate::membership::{NodeId, SwimMembership};
use crate::sharding::router::RangeRouter;
use crate::sharding::types::{MetaRange, RangeDescriptor, ShardId};
use crate::Result;
use stemedb_core::types::HlcTimestamp;

/// Manages shard split and merge operations.
///
/// The manager periodically checks shard sizes and triggers split/merge
/// when thresholds are exceeded. Changes to the meta-range are propagated
/// to all nodes via the membership gossip layer.
pub struct RangeManager {
    /// Router for shard lookups and meta-range updates.
    router: Arc<RangeRouter>,

    /// Membership for discovering nodes and broadcasting updates.
    membership: Arc<SwimMembership>,

    /// Configuration thresholds.
    config: ShardingConfig,

    /// Local node ID.
    local_node_id: NodeId,

    /// HLC clock for timestamps.
    clock: uhlc::HLC,
}

impl RangeManager {
    /// Creates a new range manager.
    pub fn new(
        router: Arc<RangeRouter>,
        membership: Arc<SwimMembership>,
        config: ShardingConfig,
        local_node_id: NodeId,
    ) -> Self {
        Self { router, membership, config, local_node_id, clock: uhlc::HLCBuilder::new().build() }
    }

    /// Checks all shards for split conditions.
    ///
    /// Returns a list of shard IDs that should be split.
    #[instrument(skip(self))]
    pub fn check_splits(&self) -> Vec<ShardId> {
        let meta = self.router.get_meta_range();
        let threshold = self.config.split_threshold_bytes;

        meta.descriptors
            .iter()
            .filter_map(|(&shard_id, desc)| {
                if desc.should_split(threshold) {
                    // Only leader should initiate split
                    if desc.leader() == Some(self.local_node_id) {
                        Some(shard_id)
                    } else {
                        None
                    }
                } else {
                    None
                }
            })
            .collect()
    }

    /// Checks for merge candidates.
    ///
    /// Returns pairs of adjacent shard IDs that can be merged.
    #[instrument(skip(self))]
    pub fn check_merges(&self) -> Vec<(ShardId, ShardId)> {
        let meta = self.router.get_meta_range();
        let threshold = self.config.merge_threshold_bytes;

        let mut merge_candidates = Vec::new();
        let shard_ids: Vec<_> = meta.descriptors.keys().copied().collect();

        // Check adjacent pairs
        for i in 0..shard_ids.len().saturating_sub(1) {
            let shard1 = shard_ids[i];
            let shard2 = shard_ids[i + 1];

            if let (Some(desc1), Some(desc2)) = (meta.get(shard1), meta.get(shard2)) {
                if desc1.can_merge_with(desc2, threshold) {
                    // Only leader of first shard should initiate merge
                    if desc1.leader() == Some(self.local_node_id) {
                        merge_candidates.push((shard1, shard2));
                    }
                }
            }
        }

        merge_candidates
    }

    /// Splits a shard into two at the midpoint.
    ///
    /// # Algorithm
    ///
    /// 1. Find the midpoint key in the shard's data
    /// 2. Create two new range descriptors
    /// 3. Assign replicas (maintain replication factor)
    /// 4. Update meta-range atomically
    /// 5. Broadcast to all nodes
    ///
    /// # Returns
    ///
    /// The IDs of the two new shards (left, right).
    #[instrument(skip(self))]
    pub async fn split_range(&self, shard_id: ShardId) -> Result<(ShardId, ShardId)> {
        let mut meta = self.router.get_meta_range();
        let timestamp = HlcTimestamp::now(&self.clock);

        let original =
            meta.get(shard_id).ok_or(crate::ClusterError::ShardNotFound(shard_id))?.clone();

        info!(shard_id = shard_id, size_bytes = original.size_bytes, "Splitting shard");

        // Generate midpoint key
        // In a real implementation, this would query the actual data distribution
        // For now, we create a synthetic midpoint based on the key range
        let midpoint = self.compute_midpoint(&original);

        // Generate new shard IDs
        let left_shard_id = self.next_shard_id(&meta);
        let right_shard_id = left_shard_id + 1;

        // Create left shard (start to midpoint)
        let left = RangeDescriptor {
            shard_id: left_shard_id,
            start_key: original.start_key.clone(),
            end_key: Some(midpoint.clone()),
            replicas: original.replicas.clone(),
            size_bytes: original.size_bytes / 2,
            assertion_count: original.assertion_count / 2,
            updated_at: timestamp,
            generation: 1,
        };

        // Create right shard (midpoint to end)
        let right = RangeDescriptor {
            shard_id: right_shard_id,
            start_key: Some(midpoint),
            end_key: original.end_key.clone(),
            replicas: original.replicas.clone(),
            size_bytes: original.size_bytes / 2,
            assertion_count: original.assertion_count / 2,
            updated_at: timestamp,
            generation: 1,
        };

        // Remove original, add new shards
        meta.remove(shard_id, timestamp);
        meta.upsert(left, timestamp);
        meta.upsert(right, timestamp);

        // Update router
        self.router.update_meta_range(meta.clone());

        // Broadcast to cluster
        self.broadcast_meta_range(&meta).await;

        info!(
            original_shard = shard_id,
            left_shard = left_shard_id,
            right_shard = right_shard_id,
            "Split complete"
        );

        Ok((left_shard_id, right_shard_id))
    }

    /// Merges two adjacent shards into one.
    ///
    /// # Algorithm
    ///
    /// 1. Verify ranges are adjacent
    /// 2. Create merged range descriptor
    /// 3. Update meta-range atomically
    /// 4. Broadcast to all nodes
    ///
    /// # Returns
    ///
    /// The ID of the merged shard.
    #[instrument(skip(self))]
    pub async fn merge_ranges(&self, left_id: ShardId, right_id: ShardId) -> Result<ShardId> {
        let mut meta = self.router.get_meta_range();
        let timestamp = HlcTimestamp::now(&self.clock);

        let left = meta.get(left_id).ok_or(crate::ClusterError::ShardNotFound(left_id))?.clone();

        let right = meta.get(right_id).ok_or(crate::ClusterError::ShardNotFound(right_id))?.clone();

        if !left.is_adjacent_to(&right) {
            return Err(crate::ClusterError::Sharding(format!(
                "Shards {left_id} and {right_id} are not adjacent"
            )));
        }

        info!(
            left_shard = left_id,
            right_shard = right_id,
            combined_size = left.size_bytes + right.size_bytes,
            "Merging shards"
        );

        // Create merged descriptor
        let merged_id = left_id; // Reuse left ID
        let merged = RangeDescriptor {
            shard_id: merged_id,
            start_key: left.start_key.clone(),
            end_key: right.end_key.clone(),
            replicas: left.replicas.clone(), // Keep left's replicas
            size_bytes: left.size_bytes.saturating_add(right.size_bytes),
            assertion_count: left.assertion_count.saturating_add(right.assertion_count),
            updated_at: timestamp,
            generation: left.generation.max(right.generation).saturating_add(1),
        };

        // Remove both, add merged
        meta.remove(left_id, timestamp);
        meta.remove(right_id, timestamp);
        meta.upsert(merged, timestamp);

        // Update router
        self.router.update_meta_range(meta.clone());

        // Broadcast to cluster
        self.broadcast_meta_range(&meta).await;

        info!(
            left_shard = left_id,
            right_shard = right_id,
            merged_shard = merged_id,
            "Merge complete"
        );

        Ok(merged_id)
    }

    /// Broadcasts the meta-range to all cluster nodes.
    #[instrument(skip(self, meta), fields(version = meta.version))]
    pub async fn broadcast_meta_range(&self, meta: &MetaRange) {
        let members = self.membership.members();

        // RPC-based meta-range broadcast is not yet wired.
        // Once stemedb-rpc integration is complete, this will send
        // UpdateMetaRange RPCs to all peers.
        for node in members {
            if node.id != self.local_node_id {
                info!(
                    target_node = %node.id.short_hex(),
                    version = meta.version,
                    "Broadcasting meta-range update (RPC pending integration)"
                );
            }
        }
    }

    /// Updates a shard's statistics (size, count).
    #[instrument(skip(self))]
    pub fn update_shard_stats(
        &self,
        shard_id: ShardId,
        size_bytes: u64,
        assertion_count: u64,
    ) -> Result<()> {
        let mut meta = self.router.get_meta_range();
        let timestamp = HlcTimestamp::now(&self.clock);

        let desc = meta.get_mut(shard_id).ok_or(crate::ClusterError::ShardNotFound(shard_id))?;

        desc.update_stats(size_bytes, assertion_count, timestamp);

        self.router.update_meta_range(meta);
        Ok(())
    }

    /// Initializes the meta-range with the given number of shards.
    ///
    /// This should be called on cluster bootstrap.
    #[instrument(skip(self))]
    pub fn initialize_shards(&self) -> Result<()> {
        let members = self.membership.members();
        if members.is_empty() {
            warn!("No members available, creating single-node meta-range");
            let node_ids = vec![self.local_node_id];
            let meta = MetaRange::with_initial_shards(
                self.config.num_shards,
                &node_ids,
                self.config.replication_factor,
            );
            self.router.update_meta_range(meta);
        } else {
            let node_ids: Vec<_> = members.iter().map(|n| n.id).collect();
            let meta = MetaRange::with_initial_shards(
                self.config.num_shards,
                &node_ids,
                self.config.replication_factor,
            );
            self.router.update_meta_range(meta);
        }

        info!(
            num_shards = self.config.num_shards,
            replication_factor = self.config.replication_factor,
            "Initialized shard meta-range"
        );

        Ok(())
    }

    /// Computes the midpoint key for splitting a range.
    fn compute_midpoint(&self, desc: &RangeDescriptor) -> Vec<u8> {
        // If we have concrete bounds, compute actual midpoint
        match (&desc.start_key, &desc.end_key) {
            (Some(start), Some(end)) => {
                // Find midpoint byte-by-byte
                let mut mid = Vec::with_capacity(start.len().max(end.len()));
                let max_len = start.len().max(end.len());

                for i in 0..max_len {
                    let s = start.get(i).copied().unwrap_or(0);
                    let e = end.get(i).copied().unwrap_or(255);
                    mid.push(s.saturating_add(e.saturating_sub(s) / 2));
                }

                mid
            }
            (None, Some(end)) => {
                // Start is min, compute midpoint towards end
                let mut mid = Vec::with_capacity(end.len());
                for &b in end {
                    mid.push(b / 2);
                }
                mid
            }
            (Some(start), None) => {
                // End is max, compute midpoint from start
                let mut mid = Vec::with_capacity(start.len());
                for &b in start {
                    mid.push(b.saturating_add((255 - b) / 2));
                }
                mid
            }
            (None, None) => {
                // Full range, split at 0x80
                vec![0x80]
            }
        }
    }

    /// Gets the next available shard ID.
    fn next_shard_id(&self, meta: &MetaRange) -> ShardId {
        meta.descriptors.keys().max().map(|&max| max + 1).unwrap_or(0)
    }
}

#[cfg(test)]
#[path = "manager_tests.rs"]
mod tests;
