use std::sync::Arc;

use anyhow::{Result, anyhow};
use async_trait::async_trait;
use storage_backend::{
    s3_error::{UploadError, retry_with_backoff},
    s3_multi_file::{FileBundle, FileMeta, MultiFileStream, group_file_metas},
    storage::BackendStorage,
};
use tokio::sync::oneshot;

use crate::{
    config::FragBatchUploaderConfig,
    utils::{
        batch_queue::{BatchExecutor, BatchQueue, Weighted},
        frag_multi_path,
        metrics::{
            inc_frag_upload_merged_file_count, observe_frag_upload_group_bytes,
            observe_frag_upload_group_count, observe_frag_upload_group_file_count,
        },
        s3_frag::parse_frag_id,
    },
};

const BATCH_QUEUE_COMPONENT: &str = "frag_batch_uploader";

#[derive(Debug)]
pub struct FragUploadReq {
    pub bucket: String,
    pub key: String,
    pub frag_id: u64,
    pub local_frag_path: String,
    pub file_metas: Vec<FileMeta>,
}

impl Weighted for FragUploadReq {
    fn weight(&self) -> usize {
        self.file_metas.iter().map(|f| f.size as usize).sum()
    }
}

pub struct FragUploadResp {
    pub key: String,
    pub offset: u64,
    pub size: u64,
}

struct FragBatchUploadExecutor {
    backend_storage: Arc<dyn BackendStorage>,
}

#[async_trait]
impl BatchExecutor for FragBatchUploadExecutor {
    type Request = FragUploadReq;
    type Response = FragUploadResp;

    async fn execute_batch(&self, reqs: Vec<Self::Request>) -> Vec<Result<Self::Response>> {
        let mut responses = Vec::with_capacity(reqs.len());
        // TODO: support batch upload later
        for req in reqs {
            let result = retry_with_backoff(|| self.upload_request(&req)).await;
            responses.push(result.map_err(|e| anyhow!("{:?}", e)));
        }
        responses
    }
}

impl FragBatchUploadExecutor {
    async fn upload_request(&self, req: &FragUploadReq) -> Result<FragUploadResp, UploadError> {
        let bundle = FileBundle::new(&req.local_frag_path, req.frag_id, req.file_metas.clone());
        let stream = MultiFileStream::new(&bundle).map_err(UploadError::Io)?;
        let size = stream.size();
        self.backend_storage
            .put_object(&req.bucket, &req.key, stream.into_stream(), size)
            .await
            .map_err(UploadError::Backend)?;
        Ok(FragUploadResp {
            key: req.key.clone(),
            offset: 0,
            size,
        })
    }
}

pub struct FragBatchUploader {
    queue: BatchQueue<FragBatchUploadExecutor>,
    config: FragBatchUploaderConfig,
    backend_storage: Arc<dyn BackendStorage>,
}
impl FragBatchUploader {
    pub fn new(backend_storage: Arc<dyn BackendStorage>, config: FragBatchUploaderConfig) -> Self {
        let executor = FragBatchUploadExecutor {
            backend_storage: backend_storage.clone(),
        };
        let queue = BatchQueue::new(
            BATCH_QUEUE_COMPONENT,
            executor,
            1, // Do not batch upload for now. TODO: support batch upload later
            config.max_batch_delay.0,
            config.max_batch_concurrency,
        );
        Self {
            queue,
            config,
            backend_storage,
        }
    }

    pub fn backend_storage(&self) -> Arc<dyn BackendStorage> {
        self.backend_storage.clone()
    }

    pub async fn upload(
        &self,
        bucket: &str,
        remote_frag_path: &str,
        local_frag_path: &str,
        file_metas: Vec<FileMeta>,
    ) -> Result<Vec<oneshot::Receiver<Result<FragUploadResp>>>> {
        let mut receivers = Vec::new();
        let frag_id = parse_frag_id(remote_frag_path)?;
        let groups = group_file_metas(file_metas, self.config.max_batch_size.0);
        // One fragment upload request may be split into multiple uploaded `multi.N`
        // objects. This metric is request-level: it records how many upload
        // groups the fragment turns into after packing.
        observe_frag_upload_group_count(groups.len());
        for (i, group) in groups.into_iter().enumerate() {
            let group_bytes = group.iter().map(|file| file.size).sum::<u64>();
            // These metrics are group-level: they describe how many files and bytes are
            // packed into each produced upload group.
            observe_frag_upload_group_file_count(group.len());
            observe_frag_upload_group_bytes(group_bytes);
            // For a group of N files we add N - 1 here, which is exactly the number of
            // single-file uploads avoided by merging them into the same upload group.
            inc_frag_upload_merged_file_count(group.len().saturating_sub(1));
            let req = FragUploadReq {
                bucket: bucket.to_string(),
                key: frag_multi_path(remote_frag_path, i as u32),
                frag_id,
                local_frag_path: local_frag_path.to_string(),
                file_metas: group,
            };
            let rx = self
                .queue
                .submit(req)
                .await
                .map_err(|e| anyhow!("Failed to submit upload request: {:?}", e))?;
            receivers.push(rx);
        }
        Ok(receivers)
    }
}

#[cfg(test)]
mod tests {
    use std::{io::Write, time::Duration};

    use anyhow::bail;
    use tempfile::TempDir;
    use tici_system::readable::{ReadableDuration, ReadableSize};

    use super::*;
    use crate::test_utils::{TEST_S3_BUCKET, test_create_backend_storage};

    fn create_test_frag_dir(dir: &TempDir, frag_id: u64) -> Result<String> {
        let frag_dir = dir.path().join(format!("f_{}", frag_id));
        std::fs::create_dir_all(&frag_dir)?;
        Ok(frag_dir.to_str().unwrap().to_string())
    }

    fn test_remote_frag_path(prefix: &str, frag_id: u64) -> String {
        format!("{}/f_{}", prefix, frag_id)
    }

    fn create_test_file(base_dir: &str, name: &str, size: usize) -> FileMeta {
        let file_path = std::path::Path::new(base_dir).join(name);
        let mut file = std::fs::File::create(&file_path).unwrap();
        let data = vec![b'a'; size];
        file.write_all(&data).unwrap();

        FileMeta {
            relative_path: name.to_string(),
            size: size as u64,
        }
    }

    #[tokio::test]
    async fn test_upload_splitting_via_config() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(100),
            max_batch_delay: ReadableDuration(Duration::from_millis(50)),
            max_batch_concurrency: 2,
        };

        let uploader = FragBatchUploader::new(backend_storage.clone(), config);

        let files = vec![
            create_test_file(&base_dir, "part_1.bin", 60),
            create_test_file(&base_dir, "part_2.bin", 60),
        ];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        assert_eq!(
            receivers.len(),
            2,
            "Should be split into 2 batches due to 100 byte limit"
        );

        for (i, rx) in receivers.into_iter().enumerate() {
            let res = rx.await??;
            assert_eq!(res.key, frag_multi_path(&remote_frag_path, i as u32));
        }
        Ok(())
    }

    #[tokio::test]
    async fn test_upload_aggregation_via_config() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(1024 * 1024),
            max_batch_delay: ReadableDuration(Duration::from_millis(50)),
            max_batch_concurrency: 2,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        let files = vec![
            create_test_file(&base_dir, "small_1.txt", 10),
            create_test_file(&base_dir, "small_2.txt", 10),
            create_test_file(&base_dir, "small_3.txt", 10),
        ];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        assert_eq!(
            receivers.len(),
            1,
            "Small files should be aggregated into 1 batch"
        );

        let Some(res) = receivers.into_iter().next() else {
            bail!("Expected one receiver, but got none");
        };
        let res = res.await??;
        assert_eq!(res.key, frag_multi_path(&remote_frag_path, 0));

        Ok(())
    }

    #[tokio::test]
    async fn test_concurrent_uploads_use_distinct_object_keys() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir_a = TempDir::new()?;
        let temp_dir_b = TempDir::new()?;
        let base_dir_a = create_test_frag_dir(&temp_dir_a, 1)?;
        let base_dir_b = create_test_frag_dir(&temp_dir_b, 2)?;
        let remote_frag_path_a = test_remote_frag_path(&prefix, 1);
        let remote_frag_path_b = test_remote_frag_path(&prefix, 2);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(100),
            max_batch_delay: ReadableDuration(Duration::from_millis(1000)),
            max_batch_concurrency: 2,
        };

        let uploader = Arc::new(FragBatchUploader::new(backend_storage, config));

        let files_a = vec![create_test_file(&base_dir_a, "file_a.txt", 50)];
        let files_b = vec![create_test_file(&base_dir_b, "file_b.txt", 100)];

        let uploader_a = uploader.clone();
        let bucket_a = bucket.clone();
        let files_a_clone = files_a.clone();

        let uploader_b = uploader.clone();
        let bucket_b = bucket.clone();
        let files_b_clone = files_b.clone();

        let (res_a, res_b) = tokio::join!(
            tokio::spawn(async move {
                uploader_a
                    .upload(&bucket_a, &remote_frag_path_a, &base_dir_a, files_a_clone)
                    .await
            }),
            tokio::spawn(async move {
                tokio::time::sleep(Duration::from_millis(100)).await;
                uploader_b
                    .upload(&bucket_b, &remote_frag_path_b, &base_dir_b, files_b_clone)
                    .await
            })
        );

        let receivers_a = res_a??;
        let receivers_b = res_b??;

        assert_eq!(receivers_a.len(), 1);
        assert_eq!(receivers_b.len(), 1);

        let result_a = receivers_a.into_iter().next().unwrap().await??;
        let result_b = receivers_b.into_iter().next().unwrap().await??;

        assert_eq!(result_a.offset, 0, "First request offset should be 0");
        assert_eq!(result_b.offset, 0, "Second request offset should be 0");
        assert_eq!(
            result_a.key,
            frag_multi_path(&test_remote_frag_path(&prefix, 1), 0)
        );
        assert_eq!(
            result_b.key,
            frag_multi_path(&test_remote_frag_path(&prefix, 2), 0)
        );

        Ok(())
    }

    #[test]
    fn test_weighted_trait() {
        let req = FragUploadReq {
            bucket: "test-bucket".to_string(),
            key: "test-id".to_string(),
            frag_id: 1,
            local_frag_path: "/tmp".to_string(),
            file_metas: vec![
                FileMeta {
                    relative_path: "file1.txt".to_string(),
                    size: 100,
                },
                FileMeta {
                    relative_path: "file2.txt".to_string(),
                    size: 200,
                },
            ],
        };
        assert_eq!(req.weight(), 300);
    }

    #[test]
    fn test_weighted_trait_empty() {
        let req = FragUploadReq {
            bucket: "test-bucket".to_string(),
            key: "test-id".to_string(),
            frag_id: 1,
            local_frag_path: "/tmp".to_string(),
            file_metas: vec![],
        };
        assert_eq!(req.weight(), 0);
    }

    #[tokio::test]
    async fn test_single_file_upload() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(1024),
            max_batch_delay: ReadableDuration(Duration::from_millis(100)),
            max_batch_concurrency: 1,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        let files = vec![create_test_file(&base_dir, "single.txt", 50)];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        assert_eq!(receivers.len(), 1);
        let res = receivers.into_iter().next().unwrap().await??;
        assert_eq!(res.key, frag_multi_path(&remote_frag_path, 0));
        assert_eq!(res.offset, 0);

        Ok(())
    }

    #[tokio::test]
    async fn test_large_file_exceeds_batch_size() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(100),
            max_batch_delay: ReadableDuration(Duration::from_millis(50)),
            max_batch_concurrency: 2,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        // File size exceeds max_batch_size
        let files = vec![create_test_file(&base_dir, "large.bin", 200)];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        // Should still process the file even though it exceeds batch size
        assert_eq!(receivers.len(), 1);
        let res = receivers.into_iter().next().unwrap().await??;
        assert_eq!(res.key, frag_multi_path(&remote_frag_path, 0));

        Ok(())
    }

    #[tokio::test]
    async fn test_multiple_groups_splitting() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(100),
            max_batch_delay: ReadableDuration(Duration::from_millis(50)),
            max_batch_concurrency: 3,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        // Create files that will be split into multiple groups
        // Each file is 60 bytes, max_batch_size is 100 bytes
        let files = vec![
            create_test_file(&base_dir, "file1.bin", 60),
            create_test_file(&base_dir, "file2.bin", 60),
            create_test_file(&base_dir, "file3.bin", 60),
        ];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        // Files should be split based on the grouping algorithm
        // Since no two 60-byte files can fit together in 100 bytes,
        // we expect at least 2 groups (but algorithm determines exact number)
        assert!(
            receivers.len() >= 2,
            "Expected at least 2 groups, got {}",
            receivers.len()
        );

        for rx in receivers.into_iter() {
            let res = rx.await??;
            assert!(res.key.starts_with(&format!("{}/multi.", remote_frag_path)));
        }

        Ok(())
    }

    #[tokio::test]
    async fn test_file_grouping_within_batch_size() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(200),
            max_batch_delay: ReadableDuration(Duration::from_millis(50)),
            max_batch_concurrency: 2,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        // Create files that can potentially be grouped together
        // Files: 60, 50, 40, 30 bytes (total 180 bytes < 200)
        let files = vec![
            create_test_file(&base_dir, "file1.bin", 60),
            create_test_file(&base_dir, "file2.bin", 50),
            create_test_file(&base_dir, "file3.bin", 40),
            create_test_file(&base_dir, "file4.bin", 30),
        ];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        // The grouping algorithm should create fewer groups than the number of files
        // since the total size fits within the batch limit
        assert!(
            receivers.len() < 4,
            "Expected efficient grouping (fewer than 4 groups), got {}",
            receivers.len()
        );

        // Verify all uploads succeed
        for rx in receivers {
            rx.await??;
        }

        Ok(())
    }

    #[tokio::test]
    async fn test_empty_file_list() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let base_dir = create_test_frag_dir(&temp_dir, 1)?;

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(1024),
            max_batch_delay: ReadableDuration(Duration::from_millis(50)),
            max_batch_concurrency: 2,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        let files = vec![];

        let receivers = uploader
            .upload(
                &bucket,
                &test_remote_frag_path(&prefix, 1),
                &base_dir,
                files,
            )
            .await?;

        // Should return empty receivers list
        assert_eq!(receivers.len(), 0);

        Ok(())
    }

    #[tokio::test]
    async fn test_offset_calculation_multiple_files() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(1024 * 1024),
            max_batch_delay: ReadableDuration(Duration::from_millis(100)),
            max_batch_concurrency: 1,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        let file1 = create_test_file(&base_dir, "file1.txt", 100);
        let file2 = create_test_file(&base_dir, "file2.txt", 150);

        let files = vec![file1, file2];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        assert_eq!(receivers.len(), 1);
        let res = receivers.into_iter().next().unwrap().await??;
        assert_eq!(res.key, frag_multi_path(&remote_frag_path, 0));
        assert_eq!(res.offset, 0);

        Ok(())
    }

    #[tokio::test]
    async fn test_zero_size_file() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(1024),
            max_batch_delay: ReadableDuration(Duration::from_millis(50)),
            max_batch_concurrency: 1,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        let files = vec![create_test_file(&base_dir, "empty.txt", 0)];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        assert_eq!(receivers.len(), 1);
        let res = receivers.into_iter().next().unwrap().await??;
        assert_eq!(res.key, frag_multi_path(&remote_frag_path, 0));

        Ok(())
    }

    #[tokio::test]
    async fn test_mixed_size_files() -> Result<()> {
        let backend_storage = test_create_backend_storage().await?;
        let bucket = TEST_S3_BUCKET.to_string();
        let prefix = uuid::Uuid::new_v4().to_string();
        let temp_dir = TempDir::new()?;
        let frag_id = 1;
        let base_dir = create_test_frag_dir(&temp_dir, frag_id)?;
        let remote_frag_path = test_remote_frag_path(&prefix, frag_id);

        let config = FragBatchUploaderConfig {
            max_batch_size: ReadableSize(150),
            max_batch_delay: ReadableDuration(Duration::from_millis(50)),
            max_batch_concurrency: 2,
        };

        let uploader = FragBatchUploader::new(backend_storage, config);

        // Mix of small and medium files
        let files = vec![
            create_test_file(&base_dir, "tiny.txt", 10),
            create_test_file(&base_dir, "small.txt", 40),
            create_test_file(&base_dir, "medium.txt", 80),
            create_test_file(&base_dir, "large.txt", 120),
        ];

        let receivers = uploader
            .upload(&bucket, &remote_frag_path, &base_dir, files)
            .await?;

        // Verify all uploads complete successfully
        assert!(!receivers.is_empty());
        for rx in receivers {
            rx.await??;
        }

        Ok(())
    }

    #[test]
    fn test_upload_resp_structure() {
        let resp = FragUploadResp {
            key: "test_key".to_string(),
            offset: 1024,
            size: 2048,
        };
        assert_eq!(resp.key, "test_key");
        assert_eq!(resp.offset, 1024);
        assert_eq!(resp.size, 2048);
    }

    #[test]
    fn test_upload_req_debug() {
        let req = FragUploadReq {
            bucket: "test-bucket".to_string(),
            key: "test-id".to_string(),
            frag_id: 1,
            local_frag_path: "/tmp".to_string(),
            file_metas: vec![],
        };
        let debug_str = format!("{:?}", req);
        assert!(debug_str.contains("test-bucket"));
        assert!(debug_str.contains("test-id"));
    }
}