use async_trait::async_trait;
use aws_sdk_s3::primitives::{ByteStream, ByteStreamError};
use aws_sdk_s3::types::{Delete, ObjectIdentifier};
use aws_sdk_s3::Client;
use vfs::engine::{BlockKey, BlockStore, VfsError, VfsResult};

#[derive(Debug, Clone, Default)]
pub struct S3BlockStoreOptions {
    pub prefix: String,
}

#[derive(Debug, Clone)]
pub struct S3BlockStore {
    client: Client,
    bucket: String,
    options: S3BlockStoreOptions,
}

impl S3BlockStore {
    pub fn new(client: Client, bucket: impl Into<String>) -> Self {
        Self::with_options(client, bucket, S3BlockStoreOptions::default())
    }

    pub fn with_options(
        client: Client,
        bucket: impl Into<String>,
        options: S3BlockStoreOptions,
    ) -> Self {
        Self {
            client,
            bucket: bucket.into(),
            options,
        }
    }

    fn key_for(&self, key: &BlockKey) -> String {
        format!("{}{}", self.options.prefix, key.0)
    }
}

#[async_trait]
impl BlockStore for S3BlockStore {
    async fn get(&self, key: &BlockKey) -> VfsResult<Vec<u8>> {
        let object_key = self.key_for(key);
        let response = self
            .client
            .get_object()
            .bucket(&self.bucket)
            .key(&object_key)
            .send()
            .await
            .map_err(|err| s3_error(format!("get s3 block '{object_key}': {err}")))?;
        collect_body(response.body, &object_key).await
    }

    async fn get_range(&self, key: &BlockKey, off: u64, len: u64) -> VfsResult<Vec<u8>> {
        if len == 0 {
            return Ok(Vec::new());
        }
        let object_key = self.key_for(key);
        let end = off
            .checked_add(len)
            .and_then(|value| value.checked_sub(1))
            .ok_or_else(|| VfsError::einval("invalid S3 byte range"))?;
        let response = self
            .client
            .get_object()
            .bucket(&self.bucket)
            .key(&object_key)
            .range(format!("bytes={off}-{end}"))
            .send()
            .await
            .map_err(|err| s3_error(format!("get s3 block range '{object_key}': {err}")))?;
        collect_body(response.body, &object_key).await
    }

    async fn put(&self, key: &BlockKey, data: &[u8]) -> VfsResult<()> {
        let object_key = self.key_for(key);
        self.client
            .put_object()
            .bucket(&self.bucket)
            .key(&object_key)
            .body(ByteStream::from(data.to_vec()))
            .send()
            .await
            .map_err(|err| s3_error(format!("put s3 block '{object_key}': {err}")))?;
        Ok(())
    }

    async fn exists(&self, key: &BlockKey) -> VfsResult<bool> {
        let object_key = self.key_for(key);
        match self
            .client
            .head_object()
            .bucket(&self.bucket)
            .key(&object_key)
            .send()
            .await
        {
            Ok(_) => Ok(true),
            Err(err)
                if err.as_service_error().is_some_and(|err| {
                    err.is_not_found()
                        || matches!(err.meta().code(), Some("NotFound" | "NoSuchKey"))
                }) =>
            {
                Ok(false)
            }
            Err(err) => Err(s3_error(format!("head s3 block '{object_key}': {err}"))),
        }
    }

    async fn delete_many(&self, keys: &[BlockKey]) -> VfsResult<()> {
        let mut errors = Vec::new();
        for chunk in keys.chunks(1000) {
            let mut object_keys = Vec::with_capacity(chunk.len());
            let mut objects = Vec::with_capacity(chunk.len());
            for key in chunk {
                let object_key = self.key_for(key);
                let object = ObjectIdentifier::builder()
                    .key(&object_key)
                    .build()
                    .map_err(|err| {
                        s3_error(format!("build s3 delete key '{object_key}': {err}"))
                    })?;
                object_keys.push(object_key);
                objects.push(object);
            }
            let delete = Delete::builder()
                .set_objects(Some(objects))
                .quiet(true)
                .build()
                .map_err(|err| s3_error(format!("build s3 delete batch: {err}")))?;
            match self
                .client
                .delete_objects()
                .bucket(&self.bucket)
                .delete(delete)
                .send()
                .await
            {
                Ok(response) => {
                    for error in response.errors() {
                        let key = error.key().unwrap_or("<unknown>");
                        let code = error.code().unwrap_or("unknown");
                        let message = error.message().unwrap_or("unknown error");
                        errors.push(format!("delete s3 block '{key}': {code}: {message}"));
                    }
                }
                Err(err) => errors.push(format!(
                    "delete s3 block batch [{}]: {err}",
                    object_keys.join(", ")
                )),
            }
        }
        if errors.is_empty() {
            Ok(())
        } else {
            Err(s3_error(format!(
                "delete {} s3 blocks failed: {}",
                errors.len(),
                errors.join("; ")
            )))
        }
    }

    async fn copy(&self, src: &BlockKey, dst: &BlockKey) -> VfsResult<()> {
        let src_key = self.key_for(src);
        let dst_key = self.key_for(dst);
        self.client
            .copy_object()
            .bucket(&self.bucket)
            .copy_source(format!("{}/{}", self.bucket, src_key))
            .key(&dst_key)
            .send()
            .await
            .map_err(|err| s3_error(format!("copy s3 block '{src_key}' to '{dst_key}': {err}")))?;
        Ok(())
    }
}

pub(crate) async fn collect_body(body: ByteStream, key: &str) -> VfsResult<Vec<u8>> {
    body.collect()
        .await
        .map_err(|err| body_error(key, err))
        .map(|bytes| bytes.into_bytes().to_vec())
}

pub(crate) fn s3_error(message: impl Into<String>) -> VfsError {
    VfsError::eio(message)
}

fn body_error(key: &str, err: ByteStreamError) -> VfsError {
    VfsError::eio(format!("read s3 object '{key}': {err}"))
}
