use super::*;
use crate::config::SwimConfig;
use crate::membership::NodeInfo;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};

fn test_node_id(n: u8) -> NodeId {
    NodeId::from_bytes([n; 16])
}

fn test_addr(port: u16) -> SocketAddr {
    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), port)
}

fn create_test_membership(local_id: NodeId) -> Arc<SwimMembership> {
    let local_info = NodeInfo::new(local_id, test_addr(9090), test_addr(8080));
    let config = SwimConfig::default();
    Arc::new(SwimMembership::new(local_info, config))
}

#[test]
fn test_compute_midpoint_full_range() {
    let local_id = test_node_id(1);
    let router = Arc::new(RangeRouter::new(local_id));
    let membership = create_test_membership(local_id);
    let config = ShardingConfig::testing();

    let manager = RangeManager::new(router, membership, config, local_id);

    let desc = RangeDescriptor::new_full_range(0, vec![local_id]);
    let midpoint = manager.compute_midpoint(&desc);

    assert_eq!(midpoint, vec![0x80]);
}

#[test]
fn test_compute_midpoint_bounded() {
    let local_id = test_node_id(1);
    let router = Arc::new(RangeRouter::new(local_id));
    let membership = create_test_membership(local_id);
    let config = ShardingConfig::testing();

    let manager = RangeManager::new(router, membership, config, local_id);

    let desc = RangeDescriptor::new(0, Some(vec![0x00]), Some(vec![0x80]), vec![local_id]);
    let midpoint = manager.compute_midpoint(&desc);

    assert_eq!(midpoint, vec![0x40]);
}

#[test]
fn test_check_splits_empty() {
    let local_id = test_node_id(1);
    let router = Arc::new(RangeRouter::new(local_id));
    let membership = create_test_membership(local_id);
    let config = ShardingConfig::testing();

    let manager = RangeManager::new(router.clone(), membership, config, local_id);

    // Initialize with small shards
    let meta = MetaRange::with_initial_shards(4, &[local_id], 1);
    router.update_meta_range(meta);

    // No splits needed (shards are empty)
    let splits = manager.check_splits();
    assert!(splits.is_empty());
}

#[test]
fn test_check_splits_needed() {
    let local_id = test_node_id(1);
    let router = Arc::new(RangeRouter::new(local_id));
    let membership = create_test_membership(local_id);
    let config = ShardingConfig::testing(); // 1MB split threshold

    let manager = RangeManager::new(router.clone(), membership, config, local_id);

    // Create meta with one oversized shard
    let mut meta = MetaRange::with_initial_shards(2, &[local_id], 1);
    if let Some(desc) = meta.get_mut(0) {
        desc.size_bytes = 2 * 1024 * 1024; // 2MB > 1MB threshold
    }
    router.update_meta_range(meta);

    let splits = manager.check_splits();
    assert_eq!(splits, vec![0]);
}

#[test]
fn test_initialize_shards() {
    let local_id = test_node_id(1);
    let router = Arc::new(RangeRouter::new(local_id));
    let membership = create_test_membership(local_id);
    let config = ShardingConfig::testing();

    let manager = RangeManager::new(router.clone(), membership, config.clone(), local_id);

    manager.initialize_shards().unwrap();

    assert_eq!(router.num_shards(), config.num_shards);
}

#[tokio::test]
async fn test_split_range() {
    let local_id = test_node_id(1);
    let router = Arc::new(RangeRouter::new(local_id));
    let membership = create_test_membership(local_id);
    let config = ShardingConfig::testing();

    let manager = RangeManager::new(router.clone(), membership, config, local_id);

    // Initialize with one shard
    let meta = MetaRange::with_initial_shards(1, &[local_id], 1);
    router.update_meta_range(meta);

    // Split shard 0
    let (left, right) = manager.split_range(0).await.unwrap();

    // Should have 2 shards now (original removed, 2 new added)
    assert_eq!(router.num_shards(), 2);

    // Verify the new shards exist
    let left_desc = router.get_descriptor(left).unwrap();
    let right_desc = router.get_descriptor(right).unwrap();

    // Left ends where right begins
    assert_eq!(left_desc.end_key, right_desc.start_key);
}

#[tokio::test]
async fn test_merge_ranges() {
    let local_id = test_node_id(1);
    let router = Arc::new(RangeRouter::new(local_id));
    let membership = create_test_membership(local_id);
    let config = ShardingConfig::testing();

    let manager = RangeManager::new(router.clone(), membership, config, local_id);

    // Create two adjacent shards
    let mut meta = MetaRange::new();
    meta.upsert(
        RangeDescriptor::new(0, None, Some(vec![0x80]), vec![local_id]),
        HlcTimestamp::default(),
    );
    meta.upsert(
        RangeDescriptor::new(1, Some(vec![0x80]), None, vec![local_id]),
        HlcTimestamp::default(),
    );
    router.update_meta_range(meta);

    // Merge them
    let merged = manager.merge_ranges(0, 1).await.unwrap();

    // Should have 1 shard now
    assert_eq!(router.num_shards(), 1);

    // Merged shard should cover full range
    let desc = router.get_descriptor(merged).unwrap();
    assert!(desc.start_key.is_none());
    assert!(desc.end_key.is_none());
}
