use std::collections::HashMap;
use std::io;
use std::pin::Pin;
use std::task::{Context as TaskContext, Poll};

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

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,
    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 = build_region(config, region)?;
        let credentials = build_credentials(config)?;
        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);
        }

        Ok(Self {
            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 mut chunk = first;
        let mut part_number: u32 = 0;
        let mut parts: Vec<Part> = Vec::new();
        let mut total: u64 = 0;
        loop {
            total += chunk.len() as u64;
            part_number += 1;
            match self
                .bucket
                .put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)
                .await
            {
                Ok(part) => parts.push(part),
                Err(e) => {
                    abort_upload(&self.bucket, &path, &upload_id).await;
                    return Err(StorageError::from(e));
                }
            }

            chunk = match read_chunk(reader).await {
                Ok(c) => c,
                Err(e) => {
                    abort_upload(&self.bucket, &path, &upload_id).await;
                    return Err(e);
                }
            };
            if chunk.is_empty() {
                abort_upload(&self.bucket, &path, &upload_id).await;
                return Err(StorageError::UnexpectedStatus {
                    operation: "put_stream",
                    status: 400,
                });
            }
            if chunk.len() < CHUNK_SIZE {
                total += chunk.len() as u64;
                part_number += 1;
                match self
                    .bucket
                    .put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)
                    .await
                {
                    Ok(part) => parts.push(part),
                    Err(e) => {
                        abort_upload(&self.bucket, &path, &upload_id).await;
                        return Err(StorageError::from(e));
                    }
                }
                break;
            }
        }

        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 into `writer` without buffering it in memory; the
    /// writer is expected to be Send + Unpin so HTTP body writers can be fed
    /// directly. Returns the number of bytes written.
    pub async fn get_to_writer<W: AsyncWrite + Send + Unpin>(
        &self,
        key: &str,
        writer: &mut W,
    ) -> Result<u64, StorageError> {
        let mut counting = CountingWriter {
            inner: writer,
            count: 0,
        };
        let status = self.bucket.get_object_to_writer(key, &mut counting).await?;
        if !(200..300).contains(&status) {
            return Err(StorageError::UnexpectedStatus {
                operation: "get_object",
                status,
            });
        }
        Ok(counting.count)
    }

    /// 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,
            })
        }
    }

    pub async fn exists(&self, key: &str) -> Result<bool, StorageError> {
        let (_, status) = self.bucket.head_object(key).await?;
        Ok((200..300).contains(&status))
    }

    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 bucket_exists(&self) -> Result<bool, StorageError> {
        Ok(self.bucket.exists().await?)
    }

    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.bucket.presign_get(key, expiry_secs, None)?)
    }

    pub fn presign_get_with_queries(
        &self,
        key: &str,
        expiry_secs: u32,
        custom_queries: HashMap<String, String>,
    ) -> Result<String, StorageError> {
        check_expiry(expiry_secs)?;
        Ok(self
            .bucket
            .presign_get(key, expiry_secs, Some(custom_queries))?)
    }

    pub fn presign_put(&self, key: &str, expiry_secs: u32) -> Result<String, StorageError> {
        check_expiry(expiry_secs)?;
        Ok(self.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
            .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.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.bucket.presign_post(key, expiry_secs, post_policy)?)
    }
}

fn build_region(config: &S3Config, region: &str) -> Result<Region, StorageError> {
    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}"))),
    }
}

fn build_credentials(config: &S3Config) -> Result<Credentials, StorageError> {
    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()?),
    }
}

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)
}

/// 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");
    }
}

struct CountingWriter<'a, W> {
    inner: &'a mut W,
    count: u64,
}

impl<W: AsyncWrite + Unpin> AsyncWrite for CountingWriter<'_, W> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut TaskContext<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let this = &mut *self;
        match Pin::new(&mut *this.inner).poll_write(cx, buf) {
            Poll::Ready(Ok(n)) => {
                this.count += n as u64;
                Poll::Ready(Ok(n))
            }
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut *self.inner).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut *self.inner).poll_shutdown(cx)
    }

    fn is_write_vectored(&self) -> bool {
        self.inner.is_write_vectored()
    }
}

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

    #[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 counting_writer_tallies_bytes() {
        let mut buf = Vec::new();
        {
            let mut cw = CountingWriter {
                inner: &mut buf,
                count: 0,
            };
            cw.write_all(b"hello ").await.expect("write");
            cw.write_all(b"world").await.expect("write");
            cw.flush().await.expect("flush");
            assert_eq!(cw.count, 11);
        }
        assert_eq!(buf, b"hello world");
    }
}