// Copyright 2026 PingCAP, Inc.

use std::{collections::VecDeque, mem, sync::Arc};

use bytes::{BufMut, Bytes, BytesMut};
use slog_global::warn;
use tokio::task::JoinHandle;

use super::metrics::KVENGINE_DFS_MULTIPART_RUNNING_UPLOADS;
use crate::dfs::{Dfs, MultipartUploadPart};

type PartUploadResult = std::result::Result<MultipartUploadPart, String>;

const MAX_MULTIPART_UPLOAD_PARTS: u64 = 10_000;

/// Buffers data and uploads it as a bounded-concurrency DFS multipart upload.
///
/// The uploader owns the DFS upload id and aborts the multipart upload on drop
/// until [`Self::complete`] succeeds. Callers may feed object bytes through
/// [`Self::upload_part`] repeatedly, then move the uploader to the component
/// that should make the object visible by calling [`Self::complete`].
pub struct MultiPartUploader {
    dfs: Arc<dyn Dfs>,
    key: String,
    upload_id: String,
    part_size: usize,
    max_uploading_parts: usize,
    next_part_number: u64,
    len: usize,
    current_part: BytesMut,
    parts: Vec<MultipartUploadPart>,
    upload_handles: VecDeque<JoinHandle<PartUploadResult>>,
    part_uploads_finished: bool,
    abort_on_drop: bool,
}

struct MultiPartRunningUploadGuard;

impl MultiPartRunningUploadGuard {
    fn new() -> Self {
        KVENGINE_DFS_MULTIPART_RUNNING_UPLOADS.inc();
        Self
    }
}

impl Drop for MultiPartRunningUploadGuard {
    fn drop(&mut self) {
        KVENGINE_DFS_MULTIPART_RUNNING_UPLOADS.dec();
    }
}

impl MultiPartUploader {
    /// Starts a DFS multipart upload.
    pub fn new(
        dfs: Arc<dyn Dfs>,
        key: String,
        part_size: usize,
        max_uploading_parts: usize,
    ) -> std::result::Result<Self, String> {
        if part_size == 0 {
            return Err("multipart upload part size must be non-zero".to_owned());
        }
        if max_uploading_parts == 0 {
            return Err("multipart upload concurrency must be non-zero".to_owned());
        }

        // Allocate before creating remote state, so an allocation failure cannot
        // leave an unfinished multipart upload behind.
        let current_part = BytesMut::with_capacity(part_size);
        let upload_id = dfs.get_runtime().block_on(dfs.begin_upload(&key))?;
        Ok(Self {
            dfs,
            key,
            upload_id,
            part_size,
            max_uploading_parts,
            next_part_number: 1,
            len: 0,
            current_part,
            parts: Vec::new(),
            upload_handles: VecDeque::new(),
            part_uploads_finished: false,
            abort_on_drop: true,
        })
    }

    /// Returns the destination key of this multipart upload.
    pub fn key(&self) -> &str {
        &self.key
    }

    /// Returns the number of object bytes accepted by this uploader.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Returns whether no object bytes have been accepted.
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Appends bytes to the multipart upload, uploading full parts in the
    /// background as they become available.
    pub fn upload_part(&mut self, mut data: &[u8]) -> std::result::Result<(), String> {
        if self.part_uploads_finished {
            return Err(format!(
                "multipart upload for {} already finished uploading parts",
                self.key
            ));
        }

        self.len = self.len.saturating_add(data.len());
        while !data.is_empty() {
            let remaining_part_capacity = self.part_size - self.current_part.len();
            let to_write = std::cmp::min(remaining_part_capacity, data.len());
            self.current_part.put_slice(&data[..to_write]);
            data = &data[to_write..];

            if self.current_part.len() == self.part_size {
                self.flush_current_part()?;
            }
        }
        Ok(())
    }

    /// Flushes all currently buffered data and waits for all part uploads.
    pub fn finish(mut self) -> std::result::Result<Self, String> {
        self.finish_part_uploads()?;
        Ok(self)
    }

    /// Completes the multipart upload and makes the object visible.
    pub fn complete(mut self) -> std::result::Result<(), String> {
        self.finish_part_uploads()?;
        let parts = mem::take(&mut self.parts);
        match self
            .dfs
            .get_runtime()
            .block_on(
                self.dfs
                    .complete_multipart_upload(&self.key, &self.upload_id, parts),
            ) {
            Ok(()) => {
                self.abort_on_drop = false;
                Ok(())
            }
            Err(err) => self.abort_after_error(err),
        }
    }

    fn finish_part_uploads(&mut self) -> std::result::Result<(), String> {
        if self.part_uploads_finished {
            return Ok(());
        }
        self.flush_current_part()?;
        self.join_all_uploads()?;
        self.part_uploads_finished = true;
        Ok(())
    }

    fn flush_current_part(&mut self) -> std::result::Result<(), String> {
        if self.current_part.is_empty() {
            return Ok(());
        }
        if self.next_part_number > MAX_MULTIPART_UPLOAD_PARTS {
            return self.abort_after_error(format!(
                "multipart upload for {} exceeds maximum part count {}",
                self.key, MAX_MULTIPART_UPLOAD_PARTS
            ));
        }

        let part = self.current_part.split().freeze();
        let part_number = self.next_part_number;
        self.next_part_number += 1;
        let handle = self.spawn_upload_part(part_number, part);
        self.upload_handles.push_back(handle);

        if self.upload_handles.len() >= self.max_uploading_parts {
            self.join_next_upload()?;
        }
        Ok(())
    }

    fn spawn_upload_part(&self, part_number: u64, part: Bytes) -> JoinHandle<PartUploadResult> {
        let dfs = Arc::clone(&self.dfs);
        let key = self.key.clone();
        let upload_id = self.upload_id.clone();
        self.dfs.get_runtime().spawn(async move {
            let _running_upload = MultiPartRunningUploadGuard::new();
            dfs.upload_part(&key, &upload_id, part_number, part).await
        })
    }

    fn join_next_upload(&mut self) -> std::result::Result<(), String> {
        let handle = self.upload_handles.pop_front().unwrap();
        let part = self.dfs.get_runtime().block_on(handle).map_err(|err| {
            format!(
                "multipart upload part task failed for file {}: {}",
                self.key, err
            )
        })??;
        self.parts.push(part);
        Ok(())
    }

    fn join_all_uploads(&mut self) -> std::result::Result<(), String> {
        while !self.upload_handles.is_empty() {
            self.join_next_upload()?;
        }
        Ok(())
    }

    fn abort_after_error(&mut self, upload_err: String) -> std::result::Result<(), String> {
        self.abort();
        Err(upload_err)
    }

    fn abort(&mut self) {
        if !self.abort_on_drop {
            return;
        }
        self.abort_on_drop = false;
        while let Some(handle) = self.upload_handles.pop_front() {
            if let Err(err) = self.dfs.get_runtime().block_on(handle) {
                warn!("multipart upload part task failed before abort";
                    "key" => &self.key,
                    "upload_id" => &self.upload_id,
                    "err" => ?err,
                );
            }
        }
        if let Err(abort_err) = self
            .dfs
            .get_runtime()
            .block_on(self.dfs.abort_multipart_upload(&self.key, &self.upload_id))
        {
            warn!("abort multipart upload failed";
                "key" => &self.key,
                "upload_id" => &self.upload_id,
                "err" => ?abort_err,
            );
        }
    }
}

impl Drop for MultiPartUploader {
    fn drop(&mut self) {
        self.abort();
    }
}

#[cfg(test)]
mod tests {
    use std::{
        sync::{
            Arc, Mutex,
            atomic::{AtomicUsize, Ordering},
        },
        time::Duration,
    };

    use async_trait::async_trait;
    use bytes::Bytes;

    use super::*;
    use crate::dfs::{Error, Options, Result};

    #[derive(Debug, Clone, PartialEq, Eq)]
    enum UploadCall {
        Begin {
            key: String,
        },
        Part {
            key: String,
            part_number: u64,
            len: usize,
        },
        Complete {
            key: String,
            parts: Vec<MultipartUploadPart>,
        },
        Abort {
            key: String,
            upload_id: String,
        },
    }

    struct TestDfs {
        runtime: tokio::runtime::Runtime,
        calls: Mutex<Vec<UploadCall>>,
        active_uploads: AtomicUsize,
        max_active_uploads: AtomicUsize,
        fail_upload_part: Option<u64>,
    }

    impl TestDfs {
        fn new(fail_upload_part: Option<u64>) -> Self {
            Self {
                runtime: tokio::runtime::Builder::new_multi_thread()
                    .worker_threads(2)
                    .enable_all()
                    .build()
                    .unwrap(),
                calls: Mutex::new(vec![]),
                active_uploads: AtomicUsize::new(0),
                max_active_uploads: AtomicUsize::new(0),
                fail_upload_part,
            }
        }

        fn calls(&self) -> Vec<UploadCall> {
            self.calls.lock().unwrap().clone()
        }

        fn record(&self, call: UploadCall) {
            self.calls.lock().unwrap().push(call);
        }

        fn observe_active_upload(&self) {
            let active = self.active_uploads.fetch_add(1, Ordering::SeqCst) + 1;
            let mut max_active = self.max_active_uploads.load(Ordering::SeqCst);
            while active > max_active {
                match self.max_active_uploads.compare_exchange(
                    max_active,
                    active,
                    Ordering::SeqCst,
                    Ordering::SeqCst,
                ) {
                    Ok(_) => break,
                    Err(current) => max_active = current,
                }
            }
        }
    }

    #[async_trait]
    impl Dfs for TestDfs {
        async fn exists(&self, _file_id: u64, _opts: Options) -> Result<bool> {
            Err(Error::Other("exists is unsupported in TestDfs".to_owned()))
        }

        async fn size(&self, _file_id: u64, _opts: Options) -> Result<u64> {
            Err(Error::Other("size is unsupported in TestDfs".to_owned()))
        }

        async fn read_file(&self, _file_id: u64, _opts: Options) -> Result<Bytes> {
            Err(Error::Other(
                "read_file is unsupported in TestDfs".to_owned(),
            ))
        }

        async fn create(&self, _file_id: u64, _data: Bytes, _opts: Options) -> Result<()> {
            Err(Error::Other("create is unsupported in TestDfs".to_owned()))
        }

        async fn remove(&self, _file_id: u64, _file_len: Option<u64>, _opts: Options) {}

        async fn permanently_remove(&self, _file_id: u64, _opts: Options) -> Result<()> {
            Err(Error::Other(
                "permanently_remove is unsupported in TestDfs".to_owned(),
            ))
        }

        fn get_runtime(&self) -> &tokio::runtime::Runtime {
            &self.runtime
        }

        fn get_prefix(&self) -> String {
            String::new()
        }

        async fn begin_upload(&self, key: &str) -> std::result::Result<String, String> {
            self.record(UploadCall::Begin {
                key: key.to_owned(),
            });
            Ok("upload-id".to_owned())
        }

        async fn upload_part(
            &self,
            key: &str,
            _upload_id: &str,
            part_number: u64,
            data: Bytes,
        ) -> std::result::Result<MultipartUploadPart, String> {
            self.observe_active_upload();
            tokio::time::sleep(Duration::from_millis(10)).await;
            self.record(UploadCall::Part {
                key: key.to_owned(),
                part_number,
                len: data.len(),
            });
            self.active_uploads.fetch_sub(1, Ordering::SeqCst);
            if self.fail_upload_part == Some(part_number) {
                return Err(format!("part {} failed", part_number));
            }
            Ok(MultipartUploadPart {
                part_number,
                e_tag: format!("etag-{}", part_number),
                checksum_crc32: None,
            })
        }

        async fn complete_multipart_upload(
            &self,
            key: &str,
            _upload_id: &str,
            parts: Vec<MultipartUploadPart>,
        ) -> std::result::Result<(), String> {
            self.record(UploadCall::Complete {
                key: key.to_owned(),
                parts,
            });
            Ok(())
        }

        async fn abort_multipart_upload(
            &self,
            key: &str,
            upload_id: &str,
        ) -> std::result::Result<(), String> {
            self.record(UploadCall::Abort {
                key: key.to_owned(),
                upload_id: upload_id.to_owned(),
            });
            Ok(())
        }
    }

    #[test]
    fn test_multipart_uploader_limits_upload_concurrency() {
        let dfs = Arc::new(TestDfs::new(None));
        let dfs_for_upload: Arc<dyn Dfs> = dfs.clone();
        let mut uploader = MultiPartUploader::new(dfs_for_upload, "rlog".to_owned(), 4, 2).unwrap();

        uploader.upload_part(b"abcdefghijklmnopqrst").unwrap();
        assert_eq!(uploader.len(), 20);
        uploader.complete().unwrap();

        let calls = dfs.calls();
        assert_eq!(
            calls.first(),
            Some(&UploadCall::Begin {
                key: "rlog".to_owned()
            })
        );
        assert!(
            !calls
                .iter()
                .any(|call| matches!(call, UploadCall::Abort { .. }))
        );
        let mut parts = calls
            .iter()
            .filter_map(|call| match call {
                UploadCall::Part {
                    key,
                    part_number,
                    len,
                } => Some((key.clone(), *part_number, *len)),
                _ => None,
            })
            .collect::<Vec<_>>();
        parts.sort_by_key(|(_, part_number, _)| *part_number);
        assert_eq!(
            parts,
            vec![
                ("rlog".to_owned(), 1, 4),
                ("rlog".to_owned(), 2, 4),
                ("rlog".to_owned(), 3, 4),
                ("rlog".to_owned(), 4, 4),
                ("rlog".to_owned(), 5, 4),
            ]
        );
        assert_eq!(
            calls.last(),
            Some(&UploadCall::Complete {
                key: "rlog".to_owned(),
                parts: vec![
                    MultipartUploadPart {
                        part_number: 1,
                        e_tag: "etag-1".to_owned(),
                        checksum_crc32: None,
                    },
                    MultipartUploadPart {
                        part_number: 2,
                        e_tag: "etag-2".to_owned(),
                        checksum_crc32: None,
                    },
                    MultipartUploadPart {
                        part_number: 3,
                        e_tag: "etag-3".to_owned(),
                        checksum_crc32: None,
                    },
                    MultipartUploadPart {
                        part_number: 4,
                        e_tag: "etag-4".to_owned(),
                        checksum_crc32: None,
                    },
                    MultipartUploadPart {
                        part_number: 5,
                        e_tag: "etag-5".to_owned(),
                        checksum_crc32: None,
                    },
                ],
            })
        );
        assert!(dfs.max_active_uploads.load(Ordering::SeqCst) <= 2);
    }

    #[test]
    fn test_multipart_uploader_aborts_on_drop() {
        let dfs = Arc::new(TestDfs::new(None));
        let dfs_for_upload: Arc<dyn Dfs> = dfs.clone();
        let mut uploader = MultiPartUploader::new(dfs_for_upload, "rlog".to_owned(), 4, 1).unwrap();

        uploader.upload_part(b"abcd").unwrap();
        let uploader = uploader.finish().unwrap();
        drop(uploader);

        let calls = dfs.calls();
        assert!(
            !calls
                .iter()
                .any(|call| matches!(call, UploadCall::Complete { .. }))
        );
        assert!(calls.iter().any(|call| matches!(
            call,
            UploadCall::Abort { key, upload_id }
                if key == "rlog" && upload_id == "upload-id"
        )));
    }

    #[test]
    fn test_multipart_uploader_aborts_on_part_error() {
        let dfs = Arc::new(TestDfs::new(Some(2)));
        let dfs_for_upload: Arc<dyn Dfs> = dfs.clone();
        let mut uploader = MultiPartUploader::new(dfs_for_upload, "rlog".to_owned(), 4, 1).unwrap();

        let err = uploader.upload_part(b"abcdefgh").unwrap_err();

        assert_eq!(err, "part 2 failed");
        drop(uploader);
        assert!(dfs.calls().iter().any(|call| matches!(
            call,
            UploadCall::Abort { key, upload_id }
                if key == "rlog" && upload_id == "upload-id"
        )));
    }

    #[test]
    fn test_multipart_uploader_rejects_too_many_parts() {
        let dfs = Arc::new(TestDfs::new(None));
        let dfs_for_upload: Arc<dyn Dfs> = dfs.clone();
        let mut uploader = MultiPartUploader::new(dfs_for_upload, "rlog".to_owned(), 1, 1).unwrap();
        uploader.next_part_number = MAX_MULTIPART_UPLOAD_PARTS + 1;

        let err = uploader.upload_part(b"x").unwrap_err();

        assert!(err.contains("exceeds maximum part count"));
        let calls = dfs.calls();
        assert!(
            !calls
                .iter()
                .any(|call| matches!(call, UploadCall::Part { .. }))
        );
        assert!(calls.iter().any(|call| matches!(
            call,
            UploadCall::Abort { key, upload_id }
                if key == "rlog" && upload_id == "upload-id"
        )));
    }
}