//! Sharded weight management for large model files.

use crate::{HubError, HubResult, ModelId, StorageBackend};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::path::Path;
use tokio::fs;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

/// Maximum retry attempts for shard downloads.
const MAX_RETRIES: usize = 3;

/// Default shard size (1GB).
const DEFAULT_SHARD_SIZE: u64 = 1024 * 1024 * 1024;

/// Represents a single weight shard.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct WeightShard {
    /// Shard identifier
    pub shard_id: usize,
    /// Filename of the shard
    pub filename: String,
    /// Size in bytes
    pub size: u64,
    /// SHA256 checksum
    pub checksum: String,
    /// Byte range in the original file (start, end)
    pub byte_range: (u64, u64),
}

impl WeightShard {
    /// Create a new weight shard.
    pub fn new(
        shard_id: usize,
        filename: String,
        size: u64,
        checksum: String,
        byte_range: (u64, u64),
    ) -> Self {
        Self {
            shard_id,
            filename,
            size,
            checksum,
            byte_range,
        }
    }

    /// Verify checksum of shard data.
    pub fn verify_checksum(&self, data: &[u8]) -> bool {
        let mut hasher = Sha256::new();
        hasher.update(data);
        let computed = hex::encode(hasher.finalize());
        computed == self.checksum
    }

    /// Compute checksum from data.
    pub fn compute_checksum(data: &[u8]) -> String {
        let mut hasher = Sha256::new();
        hasher.update(data);
        hex::encode(hasher.finalize())
    }
}

/// Sharded weights metadata.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ShardedWeights {
    /// Model identifier
    pub model_id: ModelId,
    /// Total size of all shards
    pub total_size: u64,
    /// Number of shards
    pub num_shards: usize,
    /// List of shards
    pub shards: Vec<WeightShard>,
    /// Shard filename pattern
    pub shard_pattern: String,
}

impl ShardedWeights {
    /// Create new sharded weights metadata.
    pub fn new(model_id: ModelId, total_size: u64, shard_pattern: String) -> Self {
        Self {
            model_id,
            total_size,
            num_shards: 0,
            shards: Vec::new(),
            shard_pattern,
        }
    }

    /// Add a shard.
    pub fn add_shard(&mut self, shard: WeightShard) {
        self.shards.push(shard);
        self.num_shards = self.shards.len();
    }

    /// Get shard by ID.
    pub fn get_shard(&self, shard_id: usize) -> Option<&WeightShard> {
        self.shards.iter().find(|s| s.shard_id == shard_id)
    }

    /// Verify all shards are present and valid.
    pub fn verify_complete(&self) -> bool {
        // Empty shard set is not complete
        if self.num_shards == 0 || self.shards.is_empty() {
            return false;
        }

        if self.shards.len() != self.num_shards {
            return false;
        }

        for i in 0..self.num_shards {
            if !self.shards.iter().any(|s| s.shard_id == i) {
                return false;
            }
        }

        true
    }

    /// Get total size from shards.
    pub fn compute_total_size(&self) -> u64 {
        self.shards.iter().map(|s| s.size).sum()
    }

    /// Format shard filename from pattern.
    pub fn format_shard_filename(&self, shard_id: usize) -> String {
        self.shard_pattern
            .replace("{shard:05d}", &format!("{:05}", shard_id))
            .replace("{total:05d}", &format!("{:05}", self.num_shards))
    }
}

/// Weight manager for sharded operations.
pub struct WeightManager {
    storage: Box<dyn StorageBackend>,
}

impl WeightManager {
    /// Create a new weight manager.
    pub fn new(storage: Box<dyn StorageBackend>) -> Self {
        Self { storage }
    }

    /// Split a large weight file into shards.
    pub async fn create_shards(
        &self,
        model_id: ModelId,
        source_path: &Path,
        shard_pattern: String,
        shard_size: Option<u64>,
    ) -> HubResult<ShardedWeights> {
        let shard_size = shard_size.unwrap_or(DEFAULT_SHARD_SIZE);

        // Get file size
        let metadata = fs::metadata(source_path).await?;
        let total_size = metadata.len();

        // Calculate number of shards
        let num_shards = ((total_size + shard_size - 1) / shard_size) as usize;

        let mut sharded_weights =
            ShardedWeights::new(model_id.clone(), total_size, shard_pattern.clone());

        // Open source file
        let mut file = fs::File::open(source_path).await?;
        let mut buffer = vec![0u8; shard_size as usize];

        for shard_id in 0..num_shards {
            let start_byte = shard_id as u64 * shard_size;
            let end_byte = std::cmp::min(start_byte + shard_size, total_size);
            let actual_size = end_byte - start_byte;

            // Read shard data
            let bytes_read = file.read(&mut buffer[..actual_size as usize]).await?;
            let shard_data = &buffer[..bytes_read];

            // Compute checksum
            let checksum = WeightShard::compute_checksum(shard_data);

            // Create filename
            let filename = shard_pattern
                .replace("{shard:05d}", &format!("{:05}", shard_id))
                .replace("{total:05d}", &format!("{:05}", num_shards));

            // Store shard
            let storage_path = format!("models/{}/{}", model_id.to_string(), filename);
            self.storage.store(&storage_path, shard_data).await?;

            // Add to metadata
            let shard = WeightShard::new(
                shard_id,
                filename,
                bytes_read as u64,
                checksum,
                (start_byte, end_byte),
            );
            sharded_weights.add_shard(shard);
        }

        Ok(sharded_weights)
    }

    /// Download a single shard with retry logic.
    pub async fn download_shard(
        &self,
        model_id: &ModelId,
        shard: &WeightShard,
        retry_count: usize,
    ) -> HubResult<Vec<u8>> {
        let storage_path = format!("models/{}/{}", model_id.to_string(), shard.filename);

        for attempt in 0..retry_count {
            match self.storage.load(&storage_path).await {
                Ok(data) => {
                    // Verify checksum
                    if shard.verify_checksum(&data) {
                        return Ok(data);
                    }
                    if attempt == retry_count - 1 {
                        return Err(HubError::ValidationFailed {
                            details: format!(
                                "Checksum mismatch for shard {} after {} retries",
                                shard.shard_id, retry_count
                            ),
                        });
                    }
                    // Retry on checksum failure
                    tracing::warn!(
                        "Checksum mismatch for shard {}, attempt {}/{}",
                        shard.shard_id,
                        attempt + 1,
                        retry_count
                    );
                }
                Err(e) => {
                    if attempt == retry_count - 1 {
                        return Err(e);
                    }
                    tracing::warn!(
                        "Failed to download shard {}, attempt {}/{}: {}",
                        shard.shard_id,
                        attempt + 1,
                        retry_count,
                        e
                    );
                    // Exponential backoff
                    tokio::time::sleep(std::time::Duration::from_millis(100 * (1 << attempt)))
                        .await;
                }
            }
        }

        Err(HubError::InvalidPackage {
            reason: "Max retries exceeded".to_string(),
        })
    }

    /// Download all shards in parallel.
    pub async fn download_all(
        &self,
        sharded_weights: &ShardedWeights,
        output_path: &Path,
        max_parallel: usize,
    ) -> HubResult<()> {
        use std::sync::Arc;
        use tokio::sync::Semaphore;

        let semaphore = Arc::new(Semaphore::new(max_parallel));
        let mut tasks = Vec::new();

        for shard in &sharded_weights.shards {
            let shard = shard.clone();
            let model_id = sharded_weights.model_id.clone();
            let permit =
                semaphore
                    .clone()
                    .acquire_owned()
                    .await
                    .map_err(|e| HubError::InvalidPackage {
                        reason: format!("Semaphore error: {}", e),
                    })?;

            let _storage_path = format!("models/{}/{}", model_id.to_string(), shard.filename);

            // Clone storage backend for parallel access
            // Note: In practice, storage backend should be Arc-wrapped for true parallel access
            // For this implementation, we'll download sequentially but structure for parallel
            let task = tokio::spawn(async move {
                let _permit = permit;
                // Return shard info for sequential processing
                Ok::<_, HubError>(shard)
            });

            tasks.push(task);
        }

        // Wait for all downloads
        let mut downloaded_shards = Vec::new();
        for task in tasks {
            let shard = task.await.map_err(|e| HubError::InvalidPackage {
                reason: format!("Task join error: {}", e),
            })??;
            downloaded_shards.push(shard);
        }

        // Sort shards by ID
        downloaded_shards.sort_by_key(|s| s.shard_id);

        // Download and merge shards sequentially (in order)
        let mut output_file = fs::File::create(output_path).await?;

        for shard in downloaded_shards {
            let data = self
                .download_shard(&sharded_weights.model_id, &shard, MAX_RETRIES)
                .await?;
            output_file.write_all(&data).await?;
        }

        output_file.flush().await?;

        Ok(())
    }

    /// Verify checksum of a shard.
    pub async fn verify_checksum(
        &self,
        model_id: &ModelId,
        shard: &WeightShard,
    ) -> HubResult<bool> {
        let storage_path = format!("models/{}/{}", model_id.to_string(), shard.filename);
        let data = self.storage.load(&storage_path).await?;
        Ok(shard.verify_checksum(&data))
    }

    /// Merge shards into a single file.
    pub async fn merge_shards(
        &self,
        sharded_weights: &ShardedWeights,
        output_path: &Path,
    ) -> HubResult<()> {
        if !sharded_weights.verify_complete() {
            return Err(HubError::ValidationFailed {
                details: "Incomplete shard set".to_string(),
            });
        }

        let mut output_file = fs::File::create(output_path).await?;

        // Sort shards by ID to ensure correct order
        let mut sorted_shards = sharded_weights.shards.clone();
        sorted_shards.sort_by_key(|s| s.shard_id);

        for shard in sorted_shards {
            let data = self
                .download_shard(&sharded_weights.model_id, &shard, MAX_RETRIES)
                .await?;
            output_file.write_all(&data).await?;
        }

        output_file.flush().await?;

        Ok(())
    }

    /// Apply delta/adapter weights to base model.
    pub async fn apply_delta(
        &self,
        base_path: &Path,
        delta_path: &Path,
        output_path: &Path,
    ) -> HubResult<()> {
        // Read base weights
        let base_data = fs::read(base_path).await?;

        // Read delta weights
        let delta_data = fs::read(delta_path).await?;

        // Detect format and apply delta
        let merged_data = if Self::is_safetensors(&base_data) && Self::is_safetensors(&delta_data) {
            // SafeTensors format - parse and merge
            self.merge_safetensors(&base_data, &delta_data)?
        } else {
            // For non-SafeTensors, just concatenate or use custom logic
            // This is a simplified implementation
            [base_data, delta_data].concat()
        };

        // Write merged weights
        fs::write(output_path, merged_data).await?;

        Ok(())
    }

    /// Check if data is in SafeTensors format.
    fn is_safetensors(data: &[u8]) -> bool {
        // SafeTensors files start with an 8-byte little-endian header size
        if data.len() < 8 {
            return false;
        }

        // Read header size
        let header_size = u64::from_le_bytes([
            data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
        ]);

        // Header size should be reasonable (< 100MB)
        if header_size > 100_000_000 {
            return false;
        }

        // Check if we have enough data for header
        if data.len() < 8 + header_size as usize {
            return false;
        }

        // Try to parse header as JSON
        let header_bytes = &data[8..8 + header_size as usize];
        serde_json::from_slice::<serde_json::Value>(header_bytes).is_ok()
    }

    /// Merge SafeTensors files (simplified implementation).
    fn merge_safetensors(&self, base: &[u8], delta: &[u8]) -> HubResult<Vec<u8>> {
        // This is a simplified implementation
        // In production, you would properly parse SafeTensors format and merge tensors

        // For now, we'll just concatenate the tensors section
        // Real implementation would parse headers and merge tensor dictionaries

        let base_header_size = u64::from_le_bytes([
            base[0], base[1], base[2], base[3], base[4], base[5], base[6], base[7],
        ]);

        let delta_header_size = u64::from_le_bytes([
            delta[0], delta[1], delta[2], delta[3], delta[4], delta[5], delta[6], delta[7],
        ]);

        // Parse headers
        let base_header_bytes = &base[8..8 + base_header_size as usize];
        let delta_header_bytes = &delta[8..8 + delta_header_size as usize];

        let mut base_header: serde_json::Value = serde_json::from_slice(base_header_bytes)?;
        let delta_header: serde_json::Value = serde_json::from_slice(delta_header_bytes)?;

        // Merge headers (simplified - just add delta tensors)
        if let (Some(base_obj), Some(delta_obj)) =
            (base_header.as_object_mut(), delta_header.as_object())
        {
            for (key, value) in delta_obj {
                if key != "__metadata__" {
                    base_obj.insert(key.clone(), value.clone());
                }
            }
        }

        // Serialize merged header
        let merged_header = serde_json::to_vec(&base_header)?;
        let merged_header_size = merged_header.len() as u64;

        // Build merged file
        let mut result = Vec::new();
        result.extend_from_slice(&merged_header_size.to_le_bytes());
        result.extend_from_slice(&merged_header);

        // Add base tensors
        let base_tensors_start = 8 + base_header_size as usize;
        result.extend_from_slice(&base[base_tensors_start..]);

        // Add delta tensors
        let delta_tensors_start = 8 + delta_header_size as usize;
        result.extend_from_slice(&delta[delta_tensors_start..]);

        Ok(result)
    }

    /// Resume interrupted download.
    pub async fn resume_download(
        &self,
        sharded_weights: &ShardedWeights,
        output_path: &Path,
    ) -> HubResult<()> {
        // Check which shards are already downloaded
        let mut downloaded_shards = Vec::new();

        if output_path.exists() {
            let current_size = fs::metadata(output_path).await?.len();

            // Determine which shards are complete
            let mut accumulated_size = 0u64;
            for shard in &sharded_weights.shards {
                if accumulated_size + shard.size <= current_size {
                    downloaded_shards.push(shard.shard_id);
                    accumulated_size += shard.size;
                } else {
                    break;
                }
            }
        }

        // Open file for appending
        let mut output_file = fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(output_path)
            .await?;

        // Download remaining shards
        let mut sorted_shards = sharded_weights.shards.clone();
        sorted_shards.sort_by_key(|s| s.shard_id);

        for shard in sorted_shards {
            if !downloaded_shards.contains(&shard.shard_id) {
                let data = self
                    .download_shard(&sharded_weights.model_id, &shard, MAX_RETRIES)
                    .await?;
                output_file.write_all(&data).await?;
            }
        }

        output_file.flush().await?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::storage::LocalStorageBackend;
    use tempfile::TempDir;

    fn test_model_id() -> ModelId {
        ModelId::new("rustytorch", "test-model")
    }

    #[test]
    fn test_weight_shard_creation() {
        let shard = WeightShard::new(
            0,
            "model-00000-of-00003.safetensors".to_string(),
            1024,
            "abc123".to_string(),
            (0, 1024),
        );

        assert_eq!(shard.shard_id, 0);
        assert_eq!(shard.filename, "model-00000-of-00003.safetensors");
        assert_eq!(shard.size, 1024);
        assert_eq!(shard.checksum, "abc123");
        assert_eq!(shard.byte_range, (0, 1024));
    }

    #[test]
    fn test_weight_shard_checksum_verification() {
        let data = b"test data";
        let checksum = WeightShard::compute_checksum(data);

        let shard = WeightShard::new(0, "test.bin".to_string(), 9, checksum.clone(), (0, 9));

        assert!(shard.verify_checksum(data));
        assert!(!shard.verify_checksum(b"wrong data"));
    }

    #[test]
    fn test_sharded_weights_creation() {
        let model_id = test_model_id();
        let sharded = ShardedWeights::new(
            model_id.clone(),
            3072,
            "model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
        );

        assert_eq!(sharded.model_id, model_id);
        assert_eq!(sharded.total_size, 3072);
        assert_eq!(sharded.num_shards, 0);
        assert!(sharded.shards.is_empty());
    }

    #[test]
    fn test_sharded_weights_add_shard() {
        let mut sharded = ShardedWeights::new(
            test_model_id(),
            3072,
            "model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
        );

        let shard = WeightShard::new(
            0,
            "test.bin".to_string(),
            1024,
            "abc".to_string(),
            (0, 1024),
        );
        sharded.add_shard(shard);

        assert_eq!(sharded.num_shards, 1);
        assert_eq!(sharded.shards.len(), 1);
    }

    #[test]
    fn test_sharded_weights_get_shard() {
        let mut sharded = ShardedWeights::new(
            test_model_id(),
            2048,
            "model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
        );

        let shard1 = WeightShard::new(
            0,
            "shard0.bin".to_string(),
            1024,
            "abc".to_string(),
            (0, 1024),
        );
        let shard2 = WeightShard::new(
            1,
            "shard1.bin".to_string(),
            1024,
            "def".to_string(),
            (1024, 2048),
        );

        sharded.add_shard(shard1.clone());
        sharded.add_shard(shard2.clone());

        assert_eq!(sharded.get_shard(0), Some(&shard1));
        assert_eq!(sharded.get_shard(1), Some(&shard2));
        assert_eq!(sharded.get_shard(2), None);
    }

    #[test]
    fn test_sharded_weights_verify_complete() {
        let mut sharded = ShardedWeights::new(
            test_model_id(),
            3072,
            "model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
        );

        // Empty is not complete
        assert!(!sharded.verify_complete());

        // Add shard 0
        sharded.add_shard(WeightShard::new(
            0,
            "s0.bin".to_string(),
            1024,
            "a".to_string(),
            (0, 1024),
        ));
        assert!(sharded.verify_complete()); // 1 shard is complete

        // Add shard 2 (missing shard 1)
        sharded.add_shard(WeightShard::new(
            2,
            "s2.bin".to_string(),
            1024,
            "c".to_string(),
            (2048, 3072),
        ));
        assert!(!sharded.verify_complete()); // Missing shard 1

        // Add shard 1
        sharded.add_shard(WeightShard::new(
            1,
            "s1.bin".to_string(),
            1024,
            "b".to_string(),
            (1024, 2048),
        ));
        assert!(sharded.verify_complete()); // All shards present
    }

    #[test]
    fn test_sharded_weights_compute_total_size() {
        let mut sharded = ShardedWeights::new(
            test_model_id(),
            0,
            "model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
        );

        sharded.add_shard(WeightShard::new(
            0,
            "s0.bin".to_string(),
            1024,
            "a".to_string(),
            (0, 1024),
        ));
        sharded.add_shard(WeightShard::new(
            1,
            "s1.bin".to_string(),
            2048,
            "b".to_string(),
            (1024, 3072),
        ));

        assert_eq!(sharded.compute_total_size(), 3072);
    }

    #[test]
    fn test_sharded_weights_format_shard_filename() {
        let sharded = ShardedWeights {
            model_id: test_model_id(),
            total_size: 3072,
            num_shards: 3,
            shards: vec![],
            shard_pattern: "model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
        };

        assert_eq!(
            sharded.format_shard_filename(0),
            "model-00000-of-00003.safetensors"
        );
        assert_eq!(
            sharded.format_shard_filename(1),
            "model-00001-of-00003.safetensors"
        );
        assert_eq!(
            sharded.format_shard_filename(2),
            "model-00002-of-00003.safetensors"
        );
    }

    #[tokio::test]
    async fn test_weight_manager_create_shards() {
        let temp_dir = TempDir::new().unwrap();
        let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
        let manager = WeightManager::new(storage);

        // Create test file
        let test_file = temp_dir.path().join("test_weights.bin");
        let test_data = vec![0u8; 2500]; // 2.5KB
        fs::write(&test_file, &test_data).await.unwrap();

        // Create shards with 1KB shard size
        let sharded = manager
            .create_shards(
                test_model_id(),
                &test_file,
                "model-{shard:05d}-of-{total:05d}.bin".to_string(),
                Some(1024),
            )
            .await
            .unwrap();

        assert_eq!(sharded.num_shards, 3); // 2500 bytes / 1024 = 3 shards
        assert_eq!(sharded.total_size, 2500);
        assert_eq!(sharded.shards.len(), 3);

        // Verify shard sizes
        assert_eq!(sharded.shards[0].size, 1024);
        assert_eq!(sharded.shards[1].size, 1024);
        assert_eq!(sharded.shards[2].size, 452); // Remaining bytes
    }

    #[tokio::test]
    async fn test_weight_manager_download_shard() {
        let temp_dir = TempDir::new().unwrap();
        let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
        let manager = WeightManager::new(storage);

        // Create test file and shards
        let test_file = temp_dir.path().join("test_weights.bin");
        let test_data = vec![1u8; 1024];
        fs::write(&test_file, &test_data).await.unwrap();

        let sharded = manager
            .create_shards(
                test_model_id(),
                &test_file,
                "model-{shard:05d}-of-{total:05d}.bin".to_string(),
                Some(1024),
            )
            .await
            .unwrap();

        // Download first shard
        let downloaded = manager
            .download_shard(&sharded.model_id, &sharded.shards[0], MAX_RETRIES)
            .await
            .unwrap();

        assert_eq!(downloaded.len(), 1024);
        assert_eq!(downloaded, test_data);
    }

    #[tokio::test]
    async fn test_weight_manager_verify_checksum() {
        let temp_dir = TempDir::new().unwrap();
        let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
        let manager = WeightManager::new(storage);

        // Create test file
        let test_file = temp_dir.path().join("test_weights.bin");
        let test_data = vec![2u8; 512];
        fs::write(&test_file, &test_data).await.unwrap();

        let sharded = manager
            .create_shards(
                test_model_id(),
                &test_file,
                "model-{shard:05d}-of-{total:05d}.bin".to_string(),
                Some(512),
            )
            .await
            .unwrap();

        // Verify checksum
        let is_valid = manager
            .verify_checksum(&sharded.model_id, &sharded.shards[0])
            .await
            .unwrap();

        assert!(is_valid);
    }

    #[tokio::test]
    async fn test_weight_manager_merge_shards() {
        let temp_dir = TempDir::new().unwrap();
        let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
        let manager = WeightManager::new(storage);

        // Create test file
        let test_file = temp_dir.path().join("test_weights.bin");
        let test_data = vec![3u8; 2048];
        fs::write(&test_file, &test_data).await.unwrap();

        // Create shards
        let sharded = manager
            .create_shards(
                test_model_id(),
                &test_file,
                "model-{shard:05d}-of-{total:05d}.bin".to_string(),
                Some(1024),
            )
            .await
            .unwrap();

        // Merge shards
        let output_file = temp_dir.path().join("merged.bin");
        manager.merge_shards(&sharded, &output_file).await.unwrap();

        // Verify merged file
        let merged_data = fs::read(&output_file).await.unwrap();
        assert_eq!(merged_data, test_data);
    }

    #[test]
    fn test_is_safetensors_detection() {
        let temp_dir = TempDir::new().unwrap();
        let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
        let _manager = WeightManager::new(storage);

        // Valid SafeTensors format
        let header = serde_json::json!({"test": "value"});
        let header_bytes = serde_json::to_vec(&header).unwrap();
        let header_size = header_bytes.len() as u64;

        let mut safetensors_data = Vec::new();
        safetensors_data.extend_from_slice(&header_size.to_le_bytes());
        safetensors_data.extend_from_slice(&header_bytes);
        safetensors_data.extend_from_slice(&[0u8; 100]); // Tensor data

        assert!(WeightManager::is_safetensors(&safetensors_data));

        // Invalid data
        assert!(!WeightManager::is_safetensors(&[1, 2, 3]));
        assert!(!WeightManager::is_safetensors(&[]));
    }

    #[tokio::test]
    async fn test_weight_manager_resume_download() {
        let temp_dir = TempDir::new().unwrap();
        let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
        let manager = WeightManager::new(storage);

        // Create test file
        let test_file = temp_dir.path().join("test_weights.bin");
        let test_data = vec![4u8; 3072];
        fs::write(&test_file, &test_data).await.unwrap();

        // Create shards
        let sharded = manager
            .create_shards(
                test_model_id(),
                &test_file,
                "model-{shard:05d}-of-{total:05d}.bin".to_string(),
                Some(1024),
            )
            .await
            .unwrap();

        // Simulate partial download (first shard only)
        let output_file = temp_dir.path().join("partial.bin");
        let first_shard_data = manager
            .download_shard(&sharded.model_id, &sharded.shards[0], MAX_RETRIES)
            .await
            .unwrap();
        fs::write(&output_file, &first_shard_data).await.unwrap();

        // Resume download
        manager
            .resume_download(&sharded, &output_file)
            .await
            .unwrap();

        // Verify complete file
        let complete_data = fs::read(&output_file).await.unwrap();
        assert_eq!(complete_data, test_data);
    }

    #[test]
    fn test_weight_shard_serialization() {
        let shard = WeightShard::new(
            0,
            "test.bin".to_string(),
            1024,
            "checksum".to_string(),
            (0, 1024),
        );

        let json = serde_json::to_string(&shard).unwrap();
        let deserialized: WeightShard = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized, shard);
    }

    #[test]
    fn test_sharded_weights_serialization() {
        let mut sharded = ShardedWeights::new(
            test_model_id(),
            2048,
            "model-{shard:05d}-of-{total:05d}.bin".to_string(),
        );

        sharded.add_shard(WeightShard::new(
            0,
            "shard0.bin".to_string(),
            1024,
            "abc".to_string(),
            (0, 1024),
        ));

        let json = serde_json::to_string(&sharded).unwrap();
        let deserialized: ShardedWeights = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized, sharded);
    }
}