use crate::config::S3Config as LocalS3Config;
use crate::{BlobStorage, BlobStorageError, ByteStream, File, FileReader, PART_SIZE};
use async_trait::async_trait;
use aws_sdk_s3::Client as S3Client;
use aws_sdk_s3::config::Config as S3Config;
use tokio::sync::mpsc;

#[derive(Clone)]
pub struct S3Storage {
    client: S3Client,
    bucket_name: String,
    cdn_url: Option<String>,
    path_prefix: String,
}

impl S3Storage {
    pub fn new(config: &LocalS3Config) -> Self {
        let s3_config = S3Config::builder()
            .behavior_version_latest()
            .region(aws_sdk_s3::config::Region::new(config.s3_region.clone()))
            .endpoint_url(config.s3_endpoint_url.clone())
            .credentials_provider(aws_sdk_s3::config::Credentials::new(
                &config.s3_access_key_id,
                &config.s3_secret_access_key,
                None,
                None,
                "static",
            ))
            .force_path_style(true)
            .build();

        let client = S3Client::from_conf(s3_config);

        S3Storage {
            client,
            bucket_name: config.s3_bucket_name.clone(),
            cdn_url: config.cdn_url.clone(),
            path_prefix: config.s3_path_prefix.clone(),
        }
    }
}

#[async_trait]
impl BlobStorage for S3Storage {
    async fn download_files(
        &self,
        keys: &[&str],
    ) -> Result<mpsc::Receiver<FileReader>, BlobStorageError> {
        let (tx, rx) = mpsc::channel(100);
        for key in keys {
            let file = self.download_file(key).await?;
            tx.send(file).await?;
        }
        Ok(rx)
    }

    async fn download_file(&self, key: &str) -> Result<FileReader, BlobStorageError> {
        let completed_key = format!("{}/{}", self.path_prefix, key);
        let data = self
            .client
            .get_object()
            .bucket(&self.bucket_name)
            .key(completed_key)
            .send()
            .await?;
        Ok(FileReader {
            key: key.to_string(),
            data: Box::pin(data.body.into_async_read()),
        })
    }

    async fn upload_file(&self, data: File) -> Result<Option<String>, BlobStorageError> {
        let completed_key = if data.key.starts_with('/') {
            &data.key[1..]
        } else {
            &data.key
        };
        let completed_key = format!("{}/{}", self.path_prefix, completed_key);
        let body = aws_sdk_s3::primitives::ByteStream::from(data.data.to_vec());
        let _ = self
            .client
            .put_object()
            .bucket(&self.bucket_name)
            .key(&completed_key)
            .body(body)
            .send()
            .await.expect(format!("Failed to upload file to S3: {}", completed_key).as_str());

        let cdn_url = if let Some(cdn_url) = &self.cdn_url {
            Some(format!("{}/{}", cdn_url, completed_key))
        } else {
            None
        };
        Ok(cdn_url)
    }

    async fn upload_files(
        &self,
        mut files: mpsc::Receiver<File>,
    ) -> Result<Vec<String>, BlobStorageError> {
        let mut keys = Vec::new();
        while let Some(file_data) = files.recv().await {
            let key = self.upload_file(file_data).await?;
            if let Some(key) = key {
                keys.push(key);
            }
        }
        Ok(keys)
    }

    async fn delete_file(&self, key: &str) -> Result<(), BlobStorageError> {
        let completed_key = format!("{}/{}", self.path_prefix, key);
        let _ = self
            .client
            .delete_object()
            .bucket(&self.bucket_name)
            .key(&completed_key)
            .send()
            .await?;
        Ok(())
    }

    async fn stream_upload(
        &self,
        key: &str,
        mut part: ByteStream,
    ) -> Result<Option<String>, BlobStorageError> {
        // Upload part using S3 multipart upload API
        use tokio::io::AsyncReadExt;
        let completed_key = format!("{}/{}", self.path_prefix, key);

        // 1. Initiate multipart upload
        let resp = self
            .client
            .create_multipart_upload()
            .bucket(&self.bucket_name)
            .key(&completed_key)
            .send()
            .await?;

        let upload_id = resp
            .upload_id()
            .ok_or_else(|| BlobStorageError::S3NoUploadIdReturned)?;

        let mut part_number = 1;
        let mut buffer = vec![0u8; PART_SIZE];
        let mut completed_part = Vec::new();

        loop {
            let n = part.read(&mut buffer).await?;
            if n == 0 {
                break; // EOF
            }

            let chunk = &buffer[..n];

            let upload_part_resp = self
                .client
                .upload_part()
                .bucket(&self.bucket_name)
                .key(&completed_key)
                .upload_id(upload_id)
                .part_number(part_number)
                .body(aws_sdk_s3::primitives::ByteStream::from(chunk.to_vec())) // only this chunk in memory
                .send()
                .await?;

            let etag = upload_part_resp
                .e_tag()
                .ok_or_else(|| BlobStorageError::S3NoETagReturned)?;

            completed_part.push(
                aws_sdk_s3::types::CompletedPart::builder()
                    .part_number(part_number)
                    .e_tag(etag)
                    .build(),
            );

            part_number += 1;
        }

        let completed_upload = aws_sdk_s3::types::CompletedMultipartUpload::builder()
            .set_parts(Some(completed_part))
            .build();

        let _ = self
            .client
            .complete_multipart_upload()
            .bucket(&self.bucket_name)
            .key(&completed_key)
            .upload_id(upload_id)
            .multipart_upload(completed_upload)
            .send()
            .await?;

        let cdn_url = if let Some(cdn_url) = &self.cdn_url {
            Some(format!("{}/{}", cdn_url, completed_key))
        } else {
            None
        };
        Ok(cdn_url)
    }
}
