// src/sharded_client.rs
//
// SPDX-License-Identifier: Apache-2.0 OR MIT
// SPDX-FileCopyrightText: 2025 Russ Fellows <russ.fellows@gmail.com>

use anyhow::Result;
use std::sync::Arc;

/// A collection of HTTP clients sharded to reduce contention
/// 
/// Instead of a single shared HTTP client that can become a bottleneck,
/// this maintains multiple independent client instances and routes requests
/// across them to maximize concurrency.
#[derive(Debug, Clone)]
pub struct ShardedS3Clients {
    shards: Vec<ClientShard>,
    shard_count: usize,
}

/// Individual client shard
#[derive(Debug, Clone)]
pub struct ClientShard {
    #[cfg(feature = "native-backends")]
    aws_client: Option<Arc<aws_sdk_s3::Client>>,
    
    #[cfg(feature = "arrow-backend")]
    arrow_store: Option<Arc<object_store::aws::AmazonS3>>,
    
    shard_id: usize,
}

/// Configuration for sharded clients
#[derive(Debug, Clone)]
pub struct ShardedClientConfig {
    /// Number of client shards (default: 4)
    pub shard_count: usize,
    /// Maximum connections per shard (default: 128)
    pub max_connections_per_shard: usize,
    /// HTTP/2 configuration
    pub http2_enabled: bool,
    /// Connection pool idle timeout
    pub idle_timeout: std::time::Duration,
    /// Request timeout per shard
    pub request_timeout: std::time::Duration,
}

impl Default for ShardedClientConfig {
    fn default() -> Self {
        Self {
            shard_count: 4,
            max_connections_per_shard: 128,
            http2_enabled: true,
            idle_timeout: std::time::Duration::from_secs(30),
            request_timeout: std::time::Duration::from_secs(30),
        }
    }
}

impl ShardedS3Clients {
    /// Create a placeholder client for testing (empty shards)
    pub fn placeholder(shard_count: usize) -> Self {
        Self {
            shards: Vec::new(),
            shard_count,
        }
    }

    /// Create a new set of sharded S3 clients
    pub async fn new(config: ShardedClientConfig) -> Result<Self> {
        let mut shards = Vec::with_capacity(config.shard_count);
        
        for i in 0..config.shard_count {
            let shard = ClientShard::new(i, &config).await?;
            shards.push(shard);
        }

        Ok(Self {
            shards,
            shard_count: config.shard_count,
        })
    }

    /// Get a client shard by index
    /// 
    /// Use this to route requests to specific shards based on object key,
    /// part index, or other sharding criteria to balance load.
    pub fn get_shard(&self, shard_index: usize) -> &ClientShard {
        &self.shards[shard_index % self.shard_count]
    }

    /// Get a client shard using a hash-based selection
    /// 
    /// This provides automatic load balancing by routing requests
    /// based on a hash of the input (e.g., object key, part index).
    pub fn get_shard_for_key(&self, key: &str) -> &ClientShard {
        let hash = self.hash_key(key);
        self.get_shard(hash)
    }

    /// Get a client shard for a specific part index
    /// 
    /// This ensures that parts are distributed across shards in a
    /// round-robin fashion to maximize parallelism.
    pub fn get_shard_for_part(&self, part_index: usize) -> &ClientShard {
        self.get_shard(part_index)
    }

    /// Get the number of shards
    pub fn shard_count(&self) -> usize {
        self.shard_count
    }

    /// Simple hash function for key-based sharding
    fn hash_key(&self, key: &str) -> usize {
        use std::hash::{Hash, Hasher};
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish() as usize
    }
}

impl ClientShard {
    /// Create a new client shard with optimized configuration
    async fn new(shard_id: usize, config: &ShardedClientConfig) -> Result<Self> {
        let mut shard = Self {
            #[cfg(feature = "native-backends")]
            aws_client: None,
            #[cfg(feature = "arrow-backend")]
            arrow_store: None,
            shard_id,
        };

        // Initialize AWS SDK client if feature is enabled
        #[cfg(feature = "native-backends")]
        {
            shard.aws_client = Some(Self::create_aws_client(shard_id, config).await?);
        }

        // Initialize Arrow object store if feature is enabled
        #[cfg(feature = "arrow-backend")]
        {
            shard.arrow_store = Some(Self::create_arrow_store(shard_id, config).await?);
        }

        Ok(shard)
    }

    /// Get the AWS SDK client for this shard
    #[cfg(feature = "native-backends")]
    pub fn aws_client(&self) -> Option<&Arc<aws_sdk_s3::Client>> {
        self.aws_client.as_ref()
    }

    /// Get the Arrow object store for this shard
    #[cfg(feature = "arrow-backend")]
    pub fn arrow_store(&self) -> Option<&Arc<object_store::aws::AmazonS3>> {
        self.arrow_store.as_ref()
    }

    /// Get the shard ID
    pub fn shard_id(&self) -> usize {
        self.shard_id
    }

    #[cfg(feature = "native-backends")]
    async fn create_aws_client(_shard_id: usize, _config: &ShardedClientConfig) -> Result<Arc<aws_sdk_s3::Client>> {
        // Create optimized AWS client configuration
        let aws_config = aws_config::defaults(aws_config::BehaviorVersion::latest())
            .load()
            .await;

        // TODO: Apply HTTP client optimizations per shard
        // - Set max_connections based on config.max_connections_per_shard
        // - Configure HTTP/2 based on config.http2_enabled
        // - Set timeouts based on config values
        // - Use optimized hyper client if enhanced-http feature is enabled

        let client = aws_sdk_s3::Client::new(&aws_config);
        Ok(Arc::new(client))
    }

    #[cfg(feature = "arrow-backend")]
    async fn create_arrow_store(_shard_id: usize, _config: &ShardedClientConfig) -> Result<Arc<object_store::aws::AmazonS3>> {
        use object_store::aws::AmazonS3Builder;

        // Create Arrow object store with optimized configuration
        let store = AmazonS3Builder::from_env()
            // TODO: Apply HTTP client optimizations per shard
            // - Configure connection limits
            // - Set timeouts
            // - Configure HTTP/2 if supported
            .build()?;

        Ok(Arc::new(store))
    }
}

/// Range request specification for sharded operations
#[derive(Debug, Clone)]
pub struct RangeRequest {
    /// Object key/URI
    pub key: String,
    /// Byte offset in the object
    pub offset: u64,
    /// Number of bytes to read
    pub length: usize,
    /// Part index (for tracking)
    pub part_index: usize,
}

/// Result of a range request
#[derive(Debug)]
pub struct RangeResponse {
    /// Original request
    pub request: RangeRequest,
    /// Response body bytes
    pub data: bytes::Bytes,
    /// Actual bytes read (may be less than requested at end of object)
    pub bytes_read: usize,
    /// Shard that handled the request
    pub shard_id: usize,
}

impl ShardedS3Clients {
    /// Execute a range request using the optimal shard
    /// 
    /// This automatically selects the appropriate shard and executes
    /// the range GET request with optimized settings.
    pub async fn get_range(&self, request: RangeRequest) -> Result<RangeResponse> {
        let shard = self.get_shard_for_part(request.part_index);
        
        // TODO: Implement actual range GET using the selected shard
        // This should:
        // 1. Use the shard's HTTP client (AWS SDK or Arrow)
        // 2. Make a range GET request with the specified offset/length
        // 3. Return the response data
        
        // Placeholder implementation
        let data = bytes::Bytes::new();
        
        Ok(RangeResponse {
            request,
            data,
            bytes_read: 0,
            shard_id: shard.shard_id(),
        })
    }

    /// Execute multiple range requests concurrently across shards
    /// 
    /// This distributes the requests across available shards to maximize
    /// parallelism and avoid bottlenecks on any single client.
    pub async fn get_ranges_concurrent(&self, requests: Vec<RangeRequest>) -> Result<Vec<RangeResponse>> {
        use tokio::task::JoinSet;
        
        let mut join_set = JoinSet::new();
        
        // Spawn concurrent tasks across shards
        for request in requests {
            let shards = self.clone(); // TODO: Make this more efficient
            join_set.spawn(async move {
                shards.get_range(request).await
            });
        }

        // Collect results
        let mut responses = Vec::new();
        while let Some(result) = join_set.join_next().await {
            match result {
                Ok(Ok(response)) => responses.push(response),
                Ok(Err(e)) => return Err(e),
                Err(e) => return Err(anyhow::anyhow!("Task join error: {}", e)),
            }
        }

        Ok(responses)
    }
}



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

    #[test]
    fn sharded_client_config() {
        let config = ShardedClientConfig {
            shard_count: 8,
            max_connections_per_shard: 256,
            http2_enabled: true,
            ..Default::default()
        };

        assert_eq!(config.shard_count, 8);
        assert_eq!(config.max_connections_per_shard, 256);
        assert!(config.http2_enabled);
    }

    #[test]
    fn hash_key_distribution() {
        let config = ShardedClientConfig {
            shard_count: 4,
            ..Default::default()
        };

        // Create a mock sharded client structure for testing
        let shards = ShardedS3Clients {
            shards: Vec::new(),
            shard_count: config.shard_count,
        };

        // Test that different keys hash to different shards
        let keys = ["key1", "key2", "key3", "key4", "key5"];
        let mut shard_counts = vec![0; config.shard_count];
        
        for key in &keys {
            let hash = shards.hash_key(key);
            let shard_index = hash % config.shard_count;
            shard_counts[shard_index] += 1;
        }

        // Verify distribution (not all keys should go to the same shard)
        let unique_shards = shard_counts.iter().filter(|&&count| count > 0).count();
        assert!(unique_shards > 1, "Keys should distribute across multiple shards");
    }

    #[test]
    fn part_index_distribution() {
        // Create test shards
        let mut shards = Vec::new();
        for i in 0..4 {
            shards.push(ClientShard {
                #[cfg(feature = "native-backends")]
                aws_client: None,
                #[cfg(feature = "arrow-backend")]
                arrow_store: None,
                shard_id: i,
            });
        }
        
        let sharded_clients = ShardedS3Clients {
            shards,
            shard_count: 4,
        };

        // Test round-robin distribution for part indices
        for i in 0..16 {
            let expected_shard = i % 4;
            let actual_shard = sharded_clients.get_shard(i).shard_id;
            assert_eq!(
                actual_shard, 
                expected_shard,
                "Part index {} should map to shard {}, got {}",
                i, expected_shard, actual_shard
            );
        }
    }

    #[tokio::test]
    async fn range_request_creation() {
        let request = RangeRequest {
            key: "test-object".to_string(),
            offset: 1024,
            length: 4096,
            part_index: 5,
        };

        assert_eq!(request.key, "test-object");
        assert_eq!(request.offset, 1024);
        assert_eq!(request.length, 4096);
        assert_eq!(request.part_index, 5);
    }
}