use std::pin::Pin;

use config::s3::S3Config;
use futures::StreamExt as _;
use http::HeaderMap;
use s3::Bucket;
use s3::bucket::CHUNK_SIZE;
use s3::creds::Credentials;
use s3::region::Region;
use s3::serde_types::{DateTime as S3DateTime, Object as S3Object};
use tokio::io::{AsyncRead, AsyncReadExt};

use crate::error::StorageError;

pub const MAX_PRESIGN_EXPIRY_SECS: u32 = 604_800;

#[derive(Debug, Clone)]
pub struct ObjectInfo {
    pub key: String,
    pub size: u64,
    pub last_modified: Option<S3DateTime>,
    pub etag: Option<String>,
}

impl From<S3Object> for ObjectInfo {
    fn from(obj: S3Object) -> Self {
        Self {
            key: obj.key,
            size: obj.size,
            last_modified: Some(obj.last_modified),
            etag: obj.e_tag,
        }
    }
}

#[derive(Debug, Clone)]
pub struct ObjectPage {
    pub objects: Vec<ObjectInfo>,
    pub next_continuation_token: Option<String>,
}

#[derive(Clone)]
pub struct S3Client {
    bucket: Bucket,
    /// Same bucket, but pointed at `s3.public_endpoint`. Presigned URLs carry
    /// the endpoint host inside their signature, so they must be produced by a
    /// client whose region points at the address the *caller* can reach (eg.
    /// https://s3.gitdata.ai) rather than the in-cluster service DNS name.
    presign_bucket: Bucket,
    bucket_name: String,
}

impl S3Client {
    pub fn new(config: &S3Config) -> Result<Self, StorageError> {
        let (bucket_name, region) = config.validate().map_err(StorageError::InvalidConfig)?;

        let region = {
            match config.endpoint.as_deref().filter(|s| !s.is_empty()) {
                Some(endpoint) => Ok(Region::Custom {
                    region: region.to_string(),
                    endpoint: endpoint.to_string(),
                }),
                None => region.parse::<Region>().map_err(|err| {
                    StorageError::InvalidConfig(format!("region {region:?}: {err}"))
                }),
            }
        }?;
        let credentials = {
            match (
                config.access_key_id.as_deref().filter(|s| !s.is_empty()),
                config
                    .secret_access_key
                    .as_deref()
                    .filter(|s| !s.is_empty()),
            ) {
                (Some(access), Some(secret)) => Ok(Credentials::new(
                    Some(access),
                    Some(secret),
                    None,
                    None,
                    None,
                )?),
                (Some(_), None) => Err(StorageError::InvalidConfig(
                    "access_key_id provided without secret_access_key".into(),
                )),
                (None, Some(_)) => Err(StorageError::InvalidConfig(
                    "secret_access_key provided without access_key_id".into(),
                )),
                (None, None) => Ok(Credentials::default()?),
            }
        }?;
        let mut bucket = Bucket::new(bucket_name, region, credentials)?;

        if config.use_path_style() {
            bucket = bucket.with_path_style();
        }
        if let Some(timeout) = config.request_timeout() {
            bucket = bucket.with_request_timeout(timeout);
        }

        let mut presign_bucket = bucket.clone();
        if let Some(public_endpoint) = config.public_endpoint.as_deref().filter(|s| !s.is_empty()) {
            presign_bucket.region = Region::Custom {
                region: config.region().to_string(),
                endpoint: public_endpoint.to_string(),
            };
        }

        Ok(Self {
            bucket,
            presign_bucket,
            bucket_name: bucket_name.to_string(),
        })
    }

    pub fn bucket_name(&self) -> &str {
        &self.bucket_name
    }

    pub async fn put(&self, key: &str, body: &[u8]) -> Result<(), StorageError> {
        self.put_with_content_type(key, body, "application/octet-stream")
            .await
    }

    pub async fn put_with_content_type(
        &self,
        key: &str,
        body: &[u8],
        content_type: &str,
    ) -> Result<(), StorageError> {
        let response = self
            .bucket
            .put_object_with_content_type(key, body, content_type)
            .await?;
        let status = response.status_code();
        if (200..300).contains(&status) {
            Ok(())
        } else {
            Err(StorageError::UnexpectedStatus {
                operation: "put_object",
                status,
            })
        }
    }

    /// True streaming upload with bounded memory: consumes the reader in
    /// 8 MiB chunks and uploads them serially (single PUT below one chunk,
    /// multipart above), hashing nothing but counting bytes through a tee.
    ///
    /// Returns the number of bytes actually read from the reader.
    pub async fn put_stream<R: AsyncRead + Unpin>(
        &self,
        key: &str,
        reader: &mut R,
        content_type: &str,
    ) -> Result<u64, StorageError> {
        let first = read_chunk(reader).await?;
        if first.len() < CHUNK_SIZE {
            self.put_with_content_type(key, &first, content_type)
                .await?;
            return Ok(first.len() as u64);
        }
        let msg = self
            .bucket
            .initiate_multipart_upload(key, content_type)
            .await?;
        let path = msg.key;
        let upload_id = msg.upload_id;

        let uploaded = upload_parts(reader, first, async |chunk, part_number| {
            self.bucket
                .put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)
                .await
                .map_err(StorageError::from)
        })
        .await;

        let (total, parts) = match uploaded {
            Ok(result) => result,
            Err(error) => {
                abort_upload(&self.bucket, &path, &upload_id).await;
                return Err(error);
            }
        };

        match self
            .bucket
            .complete_multipart_upload(&path, &upload_id, parts)
            .await
        {
            Ok(r) if (200..300).contains(&r.status_code()) => Ok(total),
            Ok(r) => {
                abort_upload(&self.bucket, &path, &upload_id).await;
                Err(StorageError::UnexpectedStatus {
                    operation: "complete_multipart_upload",
                    status: r.status_code(),
                })
            }
            Err(e) => {
                abort_upload(&self.bucket, &path, &upload_id).await;
                Err(StorageError::from(e))
            }
        }
    }

    pub async fn get(&self, key: &str) -> Result<Vec<u8>, StorageError> {
        let response = self.bucket.get_object(key).await?;
        let status = response.status_code();
        if !(200..300).contains(&status) {
            return Err(StorageError::UnexpectedStatus {
                operation: "get_object",
                status,
            });
        }
        Ok(response.as_slice().to_vec())
    }

    /// Stream an object out as a byte stream (ready for an HTTP response
    /// body). The returned stream yields `Result<Bytes, StorageError>`.
    pub async fn get_stream(
        &self,
        key: &str,
    ) -> Result<
        Pin<Box<dyn futures::Stream<Item = Result<bytes::Bytes, StorageError>> + Send>>,
        StorageError,
    > {
        let stream = self.bucket.get_object_stream(key).await?;
        if !(200..300).contains(&stream.status_code) {
            return Err(StorageError::UnexpectedStatus {
                operation: "get_object",
                status: stream.status_code,
            });
        }
        Ok(Box::pin(
            stream
                .bytes
                .map(move |item| item.map_err(StorageError::from)),
        ))
    }

    pub async fn delete(&self, key: &str) -> Result<(), StorageError> {
        let response = self.bucket.delete_object(key).await?;
        let status = response.status_code();
        if (200..300).contains(&status) || status == 404 {
            Ok(())
        } else {
            Err(StorageError::UnexpectedStatus {
                operation: "delete_object",
                status,
            })
        }
    }

    /// Whether a key is present. A missing key is a `false`, not an error: the
    /// underlying `HEAD` answers 404, which rust-s3 surfaces as a failure, and
    /// callers use this to choose between "create" and "already there".
    pub async fn exists(&self, key: &str) -> Result<bool, StorageError> {
        match self.bucket.head_object(key).await {
            Ok((_, status)) => Ok((200..300).contains(&status)),
            Err(s3::error::S3Error::HttpFailWithBody(404, _))
            | Err(s3::error::S3Error::HttpFailWithBody(403, _)) => Ok(false),
            Err(err) => Err(StorageError::S3(err)),
        }
    }

    pub async fn head(&self, key: &str) -> Result<Option<(u64, String)>, StorageError> {
        let (meta, status) = self.bucket.head_object(key).await?;
        if !(200..300).contains(&status) {
            return Ok(None);
        }
        Ok(Some((
            meta.content_length.unwrap_or(0).max(0) as u64,
            meta.content_type.unwrap_or_default(),
        )))
    }

    pub async fn probe(&self) -> Result<(), StorageError> {
        self.list_page("", None, Some(1)).await.map(|_| ())
    }

    pub async fn list_page(
        &self,
        prefix: &str,
        continuation_token: Option<&str>,
        max_keys: Option<usize>,
    ) -> Result<ObjectPage, StorageError> {
        let (page, status) = self
            .bucket
            .list_page(
                prefix.to_string(),
                None,
                continuation_token.map(str::to_owned),
                None,
                max_keys,
            )
            .await?;

        if !(200..300).contains(&status) {
            return Err(StorageError::UnexpectedStatus {
                operation: "list",
                status,
            });
        }

        Ok(ObjectPage {
            objects: page.contents.into_iter().map(ObjectInfo::from).collect(),
            next_continuation_token: page.next_continuation_token,
        })
    }

    pub async fn list(&self, prefix: &str) -> Result<Vec<ObjectInfo>, StorageError> {
        let mut out = Vec::new();
        let mut token: Option<String> = None;
        loop {
            let page = self.list_page(prefix, token.as_deref(), None).await?;
            let more = page.next_continuation_token.is_some();
            out.extend(page.objects);
            match page.next_continuation_token {
                Some(next) => token = Some(next),
                None => {
                    if !more {
                        break;
                    }
                }
            }
            if token.is_none() {
                break;
            }
        }
        Ok(out)
    }

    pub fn bucket(&self) -> &Bucket {
        &self.bucket
    }

    pub fn presign_get(&self, key: &str, expiry_secs: u32) -> Result<String, StorageError> {
        check_expiry(expiry_secs)?;
        Ok(self.presign_bucket.presign_get(key, expiry_secs, None)?)
    }

    pub fn presign_put(&self, key: &str, expiry_secs: u32) -> Result<String, StorageError> {
        check_expiry(expiry_secs)?;
        Ok(self.presign_bucket.presign_put(key, expiry_secs, None)?)
    }

    pub fn presign_put_with_headers(
        &self,
        key: &str,
        expiry_secs: u32,
        custom_headers: HeaderMap,
    ) -> Result<String, StorageError> {
        check_expiry(expiry_secs)?;
        Ok(self
            .presign_bucket
            .presign_put(key, expiry_secs, Some(custom_headers))?)
    }

    pub fn presign_delete(&self, key: &str, expiry_secs: u32) -> Result<String, StorageError> {
        check_expiry(expiry_secs)?;
        Ok(self.presign_bucket.presign_delete(key, expiry_secs)?)
    }

    pub fn presign_post(
        &self,
        key: &str,
        expiry_secs: u32,
        post_policy: String,
    ) -> Result<String, StorageError> {
        check_expiry(expiry_secs)?;
        Ok(self
            .presign_bucket
            .presign_post(key, expiry_secs, post_policy)?)
    }
}

fn check_expiry(expiry_secs: u32) -> Result<(), StorageError> {
    if expiry_secs > MAX_PRESIGN_EXPIRY_SECS {
        return Err(StorageError::InvalidConfig(format!(
            "presign expiry must be <= {MAX_PRESIGN_EXPIRY_SECS} seconds (got {expiry_secs})"
        )));
    }
    Ok(())
}

/// Read up to one `CHUNK_SIZE` block from the reader (EOF-tolerant).
async fn read_chunk<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec<u8>, StorageError> {
    let mut chunk = Vec::with_capacity(CHUNK_SIZE);
    let mut take = reader.take(CHUNK_SIZE as u64);
    let read = AsyncReadExt::read_to_end(&mut take, &mut chunk).await?;
    chunk.truncate(read);
    Ok(chunk)
}

/// Upload the first chunk plus every following chunk as a multipart part.
/// A zero-byte read is EOF and simply ends the upload, so a body whose size is
/// an exact multiple of `CHUNK_SIZE` completes instead of being aborted.
async fn upload_parts<R, F, T>(
    reader: &mut R,
    first: Vec<u8>,
    mut put_part: F,
) -> Result<(u64, Vec<T>), StorageError>
where
    R: AsyncRead + Unpin,
    F: AsyncFnMut(Vec<u8>, u32) -> Result<T, StorageError>,
{
    let mut chunk = first;
    let mut part_number: u32 = 0;
    let mut parts = Vec::new();
    let mut total: u64 = 0;
    loop {
        total += chunk.len() as u64;
        part_number += 1;
        parts.push(put_part(chunk, part_number).await?);
        chunk = read_chunk(reader).await?;
        if chunk.is_empty() {
            return Ok((total, parts));
        }
    }
}

/// Fire-and-forget multipart abort (errors are intentionally swallowed: the
/// original failure is what the caller cares about).
async fn abort_upload(bucket: &Bucket, path: &str, upload_id: &str) {
    if let Err(error) = bucket.abort_upload(path, upload_id).await {
        tracing::warn!(%error, "failed to abort multipart upload");
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io;

    #[tokio::test]
    async fn read_chunk_respects_chunk_size() {
        let mut data = (0..CHUNK_SIZE as u32)
            .map(|i| (i % 251) as u8)
            .collect::<Vec<_>>();
        let mut cursor = io::Cursor::new(&data);
        let first = read_chunk(&mut cursor).await.expect("read");
        assert_eq!(first.len(), CHUNK_SIZE);
        let second = read_chunk(&mut cursor).await.expect("read");
        assert!(second.is_empty());

        data.truncate(CHUNK_SIZE - 1);
        let mut cursor = io::Cursor::new(&data);
        let first = read_chunk(&mut cursor).await.expect("read");
        assert_eq!(first.len(), CHUNK_SIZE - 1);
    }

    #[tokio::test]
    async fn multipart_upload_handles_exact_chunk_multiple() {
        let data = vec![7u8; CHUNK_SIZE * 2];
        let mut cursor = io::Cursor::new(&data);
        let first = read_chunk(&mut cursor).await.expect("read");
        assert_eq!(first.len(), CHUNK_SIZE);

        let mut uploaded: Vec<usize> = Vec::new();
        let (total, part_sizes) = upload_parts(&mut cursor, first, async |chunk, _part_number| {
            uploaded.push(chunk.len());
            Ok::<usize, StorageError>(chunk.len())
        })
        .await
        .expect("an exact chunk multiple must complete");

        assert_eq!(uploaded, vec![CHUNK_SIZE, CHUNK_SIZE]);
        assert_eq!(part_sizes, vec![CHUNK_SIZE, CHUNK_SIZE]);
        assert_eq!(total, (CHUNK_SIZE * 2) as u64);
    }
}