use anyhow::Result;
use log::*;
use rusoto_credential::AutoRefreshingProvider;
use rusoto_s3::*;
use rusoto_sts::WebIdentityProvider;
use serde::{Deserialize, Serialize};
use std::{env, str::FromStr};
use tokio::{
    fs::File,
    io::{copy, AsyncReadExt},
};

use crate::*;

#[derive(Debug, Serialize, Deserialize)]
pub struct Config {
    pub endpoint: Option<String>,
    pub access_key: Option<String>,
    pub secret_access_key: Option<String>,
    pub region: Option<String>,
}

impl Config {
    pub fn new(
        endpoint: Option<String>,
        access_key: Option<String>,
        secret_access_key: Option<String>,
        region: Option<String>,
    ) -> Self {
        Config {
            endpoint,
            access_key,
            secret_access_key,
            region,
        }
    }

    pub fn new_from_env() -> Self {
        let (endpoint, access_key, secret_access_key, region) = (
            std::env::var("S3_ENDPOINT").ok(),
            std::env::var("S3_ACCESS_KEY").ok(),
            std::env::var("S3_SECRET_KEY").ok(),
            std::env::var("S3_REGION").ok(),
        );

        Self::new(endpoint, access_key, secret_access_key, region)
    }
}

impl Default for Config {
    fn default() -> Self {
        Self::new_from_env()
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Client<Config> {
    pub config: Config,
}

impl Default for Client<Config> {
    fn default() -> Self {
        Client {
            config: Config::default(),
        }
    }
}

pub struct S3Service {
    pub bucket: String,
    pub client: S3Client,
}

impl S3Service {
    pub fn new(bucket: String, client: Client<Config>) -> Self {
        debug!(
            "Creating S3 client: {:?}, bucket: {}",
            &client.config, bucket
        );

        let region = match (client.config.endpoint, client.config.region) {
            (Some(endpoint), Some(region)) => rusoto_core::Region::Custom {
                name: region,
                endpoint: endpoint + "/" + &bucket,
            },
            (Some(endpoint), None) => rusoto_core::Region::Custom {
                name: "us-east-1".to_string(),
                endpoint: endpoint + "/" + &bucket,
            },
            (None, Some(region)) => rusoto_core::Region::from_str(&region).unwrap(),
            _ => rusoto_core::Region::default(),
        };

        let request_dispatcher = rusoto_core::request::HttpClient::new().unwrap();
        let credentials_provider = match (client.config.access_key, client.config.secret_access_key)
        {
            (Some(access_key), Some(secret_access_key)) => Some(
                rusoto_credential::StaticProvider::new_minimal(access_key, secret_access_key),
            ),
            _ => None,
        };

        let s3_client = if let Some(static_provider) = credentials_provider {
            S3Client::new_with(request_dispatcher, static_provider, region)
        } else {
            let provider =
                AutoRefreshingProvider::new(WebIdentityProvider::from_k8s_env()).unwrap();
            // let provider = rusoto_credential::DefaultCredentialsProvider::new().unwrap();
            S3Client::new_with(request_dispatcher, provider, region)
            // S3Client::new(region)
        };

        S3Service {
            bucket,
            client: s3_client,
        }
    }

    async fn get_object_retries(&self, key: &str, retries: i32) -> Result<GetObjectOutput> {
        let get_request = GetObjectRequest {
            bucket: self.bucket.clone(),
            key: key.to_string(),
            ..Default::default()
        };

        let mut retries = retries;
        loop {
            match self.client.get_object(get_request.clone()).await {
                Ok(result) => return Ok(result),
                Err(e) => {
                    if retries == 0 {
                        return Err(e.into());
                    }
                    retries -= 1;
                    warn!("Failed to get object: {}, retries left: {}", e, retries);
                }
            }
        }
    }

    async fn list_objects_retries(
        &self,
        prefix: Option<String>,
        retries: i32,
    ) -> Result<ListObjectsV2Output> {
        let list_request = ListObjectsV2Request {
            prefix,
            bucket: self.bucket.clone(),
            ..Default::default()
        };

        let mut retries = retries;
        loop {
            match self.client.list_objects_v2(list_request.clone()).await {
                Ok(result) => return Ok(result),
                Err(e) => {
                    if retries == 0 {
                        return Err(e.into());
                    }
                    retries -= 1;
                    warn!("Failed to list objects: {}, retries left: {}", e, retries);
                }
            }
        }
    }
}

impl Default for S3Service {
    fn default() -> Self {
        let bucket: String = env::var("S3_FTS_BUCKET").unwrap();
        S3Service::new(bucket, Client::default())
    }
}

#[async_trait::async_trait]
impl Storage for S3Service {
    async fn upload_file(&self, key: &str, file_path: &str) -> Result<()> {
        let mut file = File::open(file_path).await?;
        let mut buffer = Vec::new();
        file.read_to_end(&mut buffer).await?;

        let put_request = PutObjectRequest {
            bucket: self.bucket.clone(),
            key: key.to_string(),
            body: Some(buffer.into()),
            ..Default::default()
        };

        self.client.put_object(put_request).await?;
        Ok(())
    }

    async fn upload_large_file(&self, key: &str, file_path: &str) -> Result<()> {
        let mut file = File::open(file_path).await?;

        let create_multipart_upload_request = CreateMultipartUploadRequest {
            bucket: self.bucket.clone(),
            key: key.to_string(),
            ..Default::default()
        };

        let multipart_upload = self
            .client
            .create_multipart_upload(create_multipart_upload_request)
            .await?;
        let upload_id = multipart_upload.upload_id.unwrap();

        let mut completed_parts = Vec::new();
        let mut part_number = 1;
        let mut buffer = vec![0; MIN_CHUNK_SIZE];

        loop {
            let bytes_read = file.read(&mut buffer).await?;
            if bytes_read == 0 {
                break;
            }

            let upload_part_request = UploadPartRequest {
                bucket: self.bucket.clone(),
                key: key.to_string(),
                part_number,
                upload_id: upload_id.clone(),
                body: Some(buffer[..bytes_read].to_vec().into()),
                ..Default::default()
            };

            let part = self.client.upload_part(upload_part_request).await?;
            completed_parts.push(CompletedPart {
                e_tag: part.e_tag,
                part_number: Some(part_number),
            });

            part_number += 1;
        }

        let complete_multipart_upload_request = CompleteMultipartUploadRequest {
            bucket: self.bucket.clone(),
            key: key.to_string(),
            upload_id,
            multipart_upload: Some(CompletedMultipartUpload {
                parts: Some(completed_parts),
            }),
            ..Default::default()
        };

        self.client
            .complete_multipart_upload(complete_multipart_upload_request)
            .await?;
        Ok(())
    }

    async fn upload_string(&self, key: &str, content: &str) -> Result<()> {
        let put_request = PutObjectRequest {
            bucket: self.bucket.clone(),
            key: key.to_string(),
            body: Some(content.as_bytes().to_vec().into()),
            ..Default::default()
        };

        self.client.put_object(put_request).await?;
        Ok(())
    }

    async fn read_file(&self, key: &str) -> Result<String> {
        let result = self.get_object_retries(key, 5).await?;
        let mut body = result.body.unwrap().into_async_read();
        let mut content = String::new();
        body.read_to_string(&mut content).await?;
        Ok(content)
    }

    async fn download_file(&self, key: &str, file_path: &str) -> Result<()> {
        let result = self.get_object_retries(key, 5).await?;
        let mut body = result.body.unwrap().into_async_read();
        let mut file = File::create(file_path).await?;
        copy(&mut body, &mut file).await?;
        Ok(())
    }

    async fn delete_file(&self, key: &str) -> Result<()> {
        let delete_request = DeleteObjectRequest {
            bucket: self.bucket.clone(),
            key: key.to_string(),
            ..Default::default()
        };

        self.client.delete_object(delete_request).await?;
        Ok(())
    }

    async fn expire_file(&self, key: &str, keep_days: i64) -> Result<()> {
        let lifecycle_rule = LifecycleRule {
            id: Some("expire_rule".to_string()),
            status: "Enabled".to_string(),
            filter: Some(LifecycleRuleFilter {
                prefix: Some(key.to_string()),
                ..Default::default()
            }),
            expiration: Some(LifecycleExpiration {
                days: Some(keep_days),
                ..Default::default()
            }),
            ..Default::default()
        };

        let lifecycle_configuration = BucketLifecycleConfiguration {
            rules: vec![lifecycle_rule],
        };

        let put_lifecycle_request = PutBucketLifecycleConfigurationRequest {
            bucket: self.bucket.clone(),
            lifecycle_configuration: Some(lifecycle_configuration),
            ..Default::default()
        };

        self.client
            .put_bucket_lifecycle_configuration(put_lifecycle_request)
            .await?;
        Ok(())
    }

    async fn list_files(&self, prefix: &str) -> Result<Vec<String>> {
        let result = self.list_objects_retries(Some(prefix.into()), 5).await?;
        let keys = result
            .contents
            .unwrap_or_default()
            .into_iter()
            .filter_map(|object| object.key)
            .collect();

        Ok(keys)
    }

    async fn move_file(&self, from_key: &str, to_key: &str) -> Result<()> {
        debug!("Moving file from {} to {}", from_key, to_key);
        let copy_request = CopyObjectRequest {
            bucket: self.bucket.clone(),
            key: to_key.to_string(),
            copy_source: format!("{}/{}", self.bucket, from_key),
            ..Default::default()
        };

        self.client.copy_object(copy_request).await?;

        let delete_request = DeleteObjectRequest {
            bucket: self.bucket.clone(),
            key: from_key.to_string(),
            ..Default::default()
        };

        self.client.delete_object(delete_request).await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::fs::File;
    use tokio::io::AsyncWriteExt;
    use tokio::runtime::Runtime;

    fn s3_service() -> S3Service {
        S3Service::default()
    }

    #[test]
    fn test_s3_service_upload_file() {
        let s3_service = s3_service();

        // Create a temp file
        let file_path = "/tmp/test-file-path-for-upload";
        let rt = Runtime::new().unwrap();
        rt.block_on(async {
            let mut temp_file = File::create(file_path).await.unwrap();
            temp_file.write_all(b"test content").await.unwrap();
        });

        // Upload file
        rt.block_on(async {
            let result = s3_service.upload_file("test-key", file_path).await;
            assert!(result.is_ok());
        });

        // Clean up
        rt.block_on(async {
            let result = tokio::fs::remove_file(file_path).await;
            assert!(result.is_ok());
        });
    }

    #[tokio::test]
    async fn test_s3_service_upload_large_file() {
        let s3_service = s3_service();

        // Create a temp file
        let file_path = "/tmp/test-file-path";

        let mut temp_file = File::create(file_path).await.unwrap();
        temp_file.write_all(b"test content").await.unwrap();

        // Upload file
        let result = s3_service.upload_large_file("test-key", file_path).await;
        assert!(result.is_ok());

        // Clean up
        let result = tokio::fs::remove_file(file_path).await;
        assert!(result.is_ok());
    }

    #[test]
    fn test_s3_service_read_file() {
        let rt = Runtime::new().unwrap();
        rt.block_on(async {
            let result = s3_service().read_file("test-key").await;
            assert!(result.is_ok());
            assert_eq!(result.unwrap(), "test content");
        });
    }

    #[test]
    fn test_s3_service_delete_file() {
        let s3_service = s3_service();

        // Upload file
        let file_path = "/tmp/test-file-path-for-delete";
        let key = "test-file-key-for-delete";
        let rt = Runtime::new().unwrap();
        rt.block_on(async {
            let mut temp_file = File::create(file_path).await.unwrap();
            temp_file
                .write_all(b"test content for delete")
                .await
                .unwrap();
        });
        rt.block_on(async {
            let result = s3_service.upload_file(key, file_path).await;
            assert!(result.is_ok());
        });

        // Delete file
        rt.block_on(async {
            let result = s3_service.delete_file(key).await;
            assert!(result.is_ok());
        });

        // Check if file is deleted
        rt.block_on(async {
            let result = s3_service.read_file(key).await;
            assert!(result.is_err());
        });

        // Clean up
        rt.block_on(async {
            let result = tokio::fs::remove_file(file_path).await;
            assert!(result.is_ok());
        });
    }

    #[test]
    fn test_s3_service_list_files() {
        let rt = Runtime::new().unwrap();
        rt.block_on(async {
            let result = s3_service().list_files("").await;
            assert!(result.is_ok());
        });
    }

    #[test]
    fn test_s3_service_move_file() {
        let s3_service = s3_service();

        // Upload file
        let file_path = "/tmp/test-file-path-for-move";
        let key = "test-file-key-for-move";
        let new_key = "test-file-new-key-for-move";
        let rt = Runtime::new().unwrap();
        rt.block_on(async {
            let mut temp_file = File::create(file_path).await.unwrap();
            temp_file.write_all(b"test content for move").await.unwrap();
        });
        rt.block_on(async {
            let result = s3_service.upload_file(key, file_path).await;
            assert!(result.is_ok());
        });

        // Move file
        rt.block_on(async {
            let result = s3_service.move_file(key, new_key).await;
            assert!(result.is_ok());
        });

        // Check if file is moved
        rt.block_on(async {
            let result = s3_service.read_file(new_key).await;
            assert!(result.is_ok());
            assert_eq!(result.unwrap(), "test content for move");
        });

        // Clean up
        rt.block_on(async {
            let result = s3_service.delete_file(new_key).await;
            assert!(result.is_ok());
        });
        rt.block_on(async {
            let result = tokio::fs::remove_file(file_path).await;
            assert!(result.is_ok());
        });
    }
}