use std::{collections::VecDeque, future::Future};

use futures::{
    FutureExt as _, Stream, StreamExt as _,
    future::{BoxFuture, pending},
};
use k8s_openapi::apimachinery::pkg::apis::meta::v1::Status as RemoteStatus;
use kube::api::AttachedProcess;
use tokio::{
    io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _},
    sync::mpsc,
};
use tonic::Status;

use crate::plugin::{CopyInChunk, CopyOutChunk, DataChunk, ExecOutput, data_chunk, exec_output};

const STDERR_LIMIT: usize = 16 * 1024;
type Reader = Box<dyn AsyncRead + Unpin + Send>;
type Writer = Box<dyn AsyncWrite + Unpin + Send>;
type StatusFuture = BoxFuture<'static, Option<RemoteStatus>>;

/// Owns the WebSocket driver as well as its extracted streams. Dropping an
/// operation requests local driver abortion, not termination of the remote command.
pub struct Attachment {
    _process: Option<AttachedProcess>,
    stdin: Option<Writer>,
    stdout: Option<Reader>,
    stderr: Option<Reader>,
    status: Option<StatusFuture>,
    #[cfg(test)]
    _drop_guard: Option<tests::DropGuard>,
}

impl From<AttachedProcess> for Attachment {
    fn from(mut process: AttachedProcess) -> Self {
        Self {
            stdin: process.stdin().map(|writer| -> Writer { Box::new(writer) }),
            stdout: process
                .stdout()
                .map(|reader| -> Reader { Box::new(reader) }),
            stderr: process
                .stderr()
                .map(|reader| -> Reader { Box::new(reader) }),
            status: process
                .take_status()
                .map(|status| -> StatusFuture { Box::pin(status) }),
            _process: Some(process),
            #[cfg(test)]
            _drop_guard: None,
        }
    }
}

#[derive(Debug)]
pub enum Failure {
    Remote(String),
    Stream(Status),
}

/// Forwards data only; environment policy and terminal messages belong to the caller.
pub async fn command(
    process: Attachment,
    output: &mpsc::Sender<Result<ExecOutput, Status>>,
) -> Result<i32, Failure> {
    supervise(process, async { Ok(0) }, Output::Command(output))
        .await
        .map(|completion| completion.exit_code)
}

pub async fn upload_archive<Chunks>(
    mut process: Attachment,
    first_data: Vec<u8>,
    remaining: Chunks,
) -> Result<usize, Failure>
where
    Chunks: Stream<Item = Result<CopyInChunk, Status>> + Unpin,
{
    let stdin = process
        .stdin
        .take()
        .ok_or_else(|| internal("Kubernetes exec did not provide stdin"))?;
    supervise(
        process,
        write_archive(stdin, first_data, remaining),
        Output::Upload,
    )
    .await
    .map(|completion| completion.bytes)
}

pub async fn download_archive(
    process: Attachment,
    output: &mpsc::Sender<Result<CopyOutChunk, Status>>,
) -> Result<(), Failure> {
    supervise(process, async { Ok(0) }, Output::Download(output))
        .await
        .map(|_| ())
}

#[derive(Clone, Copy)]
enum Output<'sender> {
    Command(&'sender mpsc::Sender<Result<ExecOutput, Status>>),
    Upload,
    Download(&'sender mpsc::Sender<Result<CopyOutChunk, Status>>),
}

impl Output<'_> {
    async fn closed(self) {
        match self {
            Self::Command(sender) => sender.closed().await,
            Self::Download(sender) => sender.closed().await,
            Self::Upload => pending().await,
        }
    }

    fn cancelled(self) -> Failure {
        Failure::Stream(Status::cancelled(match self {
            Self::Command(_) => "Exec stream closed",
            Self::Upload => "CopyIn stream closed",
            Self::Download(_) => "CopyOut stream closed",
        }))
    }
}

struct Completion {
    exit_code: i32,
    bytes: usize,
}

enum InputFailure {
    Source(Status),
    Stdin(Status),
}

async fn supervise(
    process: Attachment,
    input: impl Future<Output = Result<usize, InputFailure>>,
    output: Output<'_>,
) -> Result<Completion, Failure> {
    let mut stderr = StderrTail::default();
    let result = supervise_inner(process, input, output, &mut stderr).await;
    result.map_err(|failure| stderr.annotate(failure))
}

async fn supervise_inner(
    mut process: Attachment,
    input: impl Future<Output = Result<usize, InputFailure>>,
    output: Output<'_>,
    stderr_tail: &mut StderrTail,
) -> Result<Completion, Failure> {
    let stdout = process
        .stdout
        .take()
        .ok_or_else(|| internal("Kubernetes exec did not provide stdout"))?;
    let stderr = process
        .stderr
        .take()
        .ok_or_else(|| internal("Kubernetes exec did not provide stderr"))?;
    let mut status = process
        .status
        .take()
        .ok_or_else(|| internal("Kubernetes exec did not provide a status channel"))?;
    let capture = (!matches!(output, Output::Command(_))).then_some(stderr_tail);
    let drains = async {
        tokio::try_join!(
            pump(stdout, output, data_chunk::Stream::Stdout, None),
            pump(stderr, output, data_chunk::Stream::Stderr, capture),
        )?;
        Ok::<(), Failure>(())
    };
    tokio::pin!(drains);
    let mut drained = false;
    let mut bytes = None;
    let remote = {
        // Tar can consume its end markers and exit zero before the source's EOF.
        // Keep that upload alive; other statuses leave this scope and drop even
        // an upload blocked reading a chunk or writing to the bounded pipe.
        let mut succeeded = false;
        tokio::pin!(input);
        loop {
            if succeeded && bytes.is_some() {
                break Ok(0);
            }
            tokio::select! {
                biased;
                () = output.closed() => return Err(output.cancelled()),
                result = &mut input, if bytes.is_none() => match result {
                    Ok(count) => bytes = Some(count),
                    Err(error) => {
                        // A ready remote failure can explain a broken stdin pipe.
                        // Collect only ready diagnostics/status: an input error must
                        // not wait for remote completion or output EOF.
                        if !drained {
                            let _ = drains.as_mut().now_or_never();
                        }
                        if matches!(&error, InputFailure::Stdin(_)) && !succeeded
                            && let Some(Some(remote)) = status.as_mut().now_or_never()
                        {
                            let exit_code = parse_exit_code(Some(remote))?;
                            if exit_code != 0 {
                                return Err(internal(format!("tar exited with status {exit_code}")));
                            }
                        }
                        return Err(match error {
                            InputFailure::Source(error_status) | InputFailure::Stdin(error_status) => Failure::Stream(error_status),
                        });
                    }
                },
                remote = &mut status, if !succeeded => match parse_exit_code(remote) {
                    Ok(0) => succeeded = true,
                    result => break result,
                },
                result = &mut drains, if !drained => {
                    result?;
                    drained = true;
                }
            }
        }
    };
    if !drained {
        tokio::select! {
            biased;
            () = output.closed() => return Err(output.cancelled()),
            result = &mut drains => result?,
        }
    }
    // Do not join the driver: join consumes the owner and drops untaken
    // endpoints; cancelling that owned join would detach its driver task.
    let exit_code = remote?;
    if !matches!(output, Output::Command(_)) && exit_code != 0 {
        return Err(internal(format!("tar exited with status {exit_code}")));
    }
    let bytes = bytes.ok_or_else(|| {
        internal("CopyIn tar completed before all input was consumed and stdin reached EOF")
    })?;
    Ok(Completion { exit_code, bytes })
}

async fn write_archive(
    mut stdin: Writer,
    first_data: Vec<u8>,
    mut remaining: impl Stream<Item = Result<CopyInChunk, Status>> + Unpin,
) -> Result<usize, InputFailure> {
    stdin.write_all(&first_data).await.map_err(|error| {
        InputFailure::Stdin(Status::internal(format!(
            "write first CopyIn chunk: {error}"
        )))
    })?;
    let mut bytes = first_data.len();
    while let Some(chunk) = remaining.next().await {
        let chunk = chunk.map_err(InputFailure::Source)?;
        if chunk.environment_id.is_some() || chunk.dest_path.is_some() {
            return Err(InputFailure::Source(Status::invalid_argument(
                "only the first CopyIn chunk may contain environment_id or dest_path",
            )));
        }
        stdin.write_all(&chunk.data).await.map_err(|error| {
            InputFailure::Stdin(Status::internal(format!("write CopyIn chunk: {error}")))
        })?;
        bytes = bytes
            .checked_add(chunk.data.len())
            .ok_or_else(|| InputFailure::Source(Status::internal("CopyIn byte count overflow")))?;
    }
    stdin.shutdown().await.map_err(|error| {
        InputFailure::Stdin(Status::internal(format!("close CopyIn stdin: {error}")))
    })?;
    drop(stdin);
    Ok(bytes)
}

async fn pump(
    mut reader: Reader,
    output: Output<'_>,
    stream: data_chunk::Stream,
    mut tail: Option<&mut StderrTail>,
) -> Result<(), Failure> {
    let mut buffer = vec![0; 8192];
    loop {
        let count = reader
            .read(&mut buffer)
            .await
            .map_err(|error| internal(format!("read Kubernetes exec {stream:?}: {error}")))?;
        if count == 0 {
            return Ok(());
        }
        let data = &buffer[..count];
        if let Some(tail) = tail.as_mut() {
            tail.push(data);
        }
        match output {
            Output::Command(sender) => sender
                .send(Ok(ExecOutput {
                    output: Some(exec_output::Output::Data(DataChunk {
                        stream: stream as i32,
                        data: data.to_vec(),
                    })),
                }))
                .await
                .map_err(|_| output.cancelled())?,
            Output::Download(sender) if stream == data_chunk::Stream::Stdout => sender
                .send(Ok(CopyOutChunk {
                    data: data.to_vec(),
                }))
                .await
                .map_err(|_| output.cancelled())?,
            Output::Upload | Output::Download(_) => {}
        }
    }
}

#[derive(Default)]
struct StderrTail {
    bytes: VecDeque<u8>,
    truncated: bool,
}

impl StderrTail {
    fn push(&mut self, data: &[u8]) {
        for byte in data {
            if self.bytes.len() == STDERR_LIMIT {
                self.bytes.pop_front();
                self.truncated = true;
            }
            self.bytes.push_back(*byte);
        }
    }

    fn annotate(self, failure: Failure) -> Failure {
        if self.bytes.is_empty() {
            return failure;
        }
        let bytes: Vec<_> = self.bytes.into_iter().collect();
        let marker = if self.truncated {
            "[stderr truncated; showing last 16384 bytes]\n"
        } else {
            ""
        };
        let diagnostic = format!("\ntar stderr:\n{marker}{}", String::from_utf8_lossy(&bytes));
        match failure {
            Failure::Remote(message) => Failure::Remote(format!("{message}{diagnostic}")),
            Failure::Stream(status) => Failure::Stream(Status::with_details_and_metadata(
                status.code(),
                format!("{}{diagnostic}", status.message()),
                status.details().to_vec().into(),
                status.metadata().clone(),
            )),
        }
    }
}

fn internal(message: impl Into<String>) -> Failure {
    Failure::Stream(Status::internal(message))
}

fn parse_exit_code(status: Option<RemoteStatus>) -> Result<i32, Failure> {
    let status = status
        .ok_or_else(|| Failure::Remote("Kubernetes exec ended without a status".to_owned()))?;
    if status.status.as_deref() == Some("Success") {
        return Ok(0);
    }
    if status.reason.as_deref() == Some("NonZeroExitCode")
        && let Some(cause) = status
            .details
            .as_ref()
            .and_then(|details| details.causes.as_ref())
            .and_then(|causes| {
                causes
                    .iter()
                    .find(|cause| cause.reason.as_deref() == Some("ExitCode"))
            })
    {
        return cause
            .message
            .as_deref()
            .ok_or_else(|| Failure::Remote("Kubernetes exit-code cause had no message".to_owned()))?
            .parse::<i32>()
            .map_err(|error| Failure::Remote(format!("invalid Kubernetes exit code: {error}")));
    }
    Err(Failure::Remote(format!(
        "Kubernetes exec failed without a parseable exit code: reason={:?}, message={:?}",
        status.reason, status.message
    )))
}

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

    use futures::{
        future::{Either, select},
        stream,
    };
    use k8s_openapi::apimachinery::pkg::apis::meta::v1::{StatusCause, StatusDetails};
    use tokio::{
        io::{DuplexStream, duplex},
        sync::oneshot,
        time::timeout,
    };

    use tokio_stream::wrappers::ReceiverStream;

    use super::*;

    pub(super) struct DropGuard(Arc<AtomicBool>);

    impl Drop for DropGuard {
        fn drop(&mut self) {
            self.0.store(true, Ordering::SeqCst);
        }
    }

    struct Script {
        stdin: DuplexStream,
        stdout: DuplexStream,
        stderr: DuplexStream,
        status: oneshot::Sender<RemoteStatus>,
    }

    fn scripted() -> (Attachment, Script, Arc<AtomicBool>) {
        let (stdin, remote_stdin) = duplex(8);
        let (stdout, remote_stdout) = duplex(8);
        let (stderr, remote_stderr) = duplex(8);
        let (status_tx, status_rx) = oneshot::channel();
        let dropped = Arc::new(AtomicBool::new(false));
        (
            Attachment {
                _process: None,
                stdin: Some(Box::new(stdin)),
                stdout: Some(Box::new(stdout)),
                stderr: Some(Box::new(stderr)),
                status: Some(Box::pin(async { status_rx.await.ok() })),
                _drop_guard: Some(DropGuard(Arc::clone(&dropped))),
            },
            Script {
                stdin: remote_stdin,
                stdout: remote_stdout,
                stderr: remote_stderr,
                status: status_tx,
            },
            dropped,
        )
    }

    fn exit_status(code: i32) -> RemoteStatus {
        if code == 0 {
            return RemoteStatus {
                status: Some("Success".to_owned()),
                ..RemoteStatus::default()
            };
        }
        RemoteStatus {
            status: Some("Failure".to_owned()),
            reason: Some("NonZeroExitCode".to_owned()),
            details: Some(StatusDetails {
                causes: Some(vec![StatusCause {
                    reason: Some("ExitCode".to_owned()),
                    message: Some(code.to_string()),
                    ..StatusCause::default()
                }]),
                ..StatusDetails::default()
            }),
            ..RemoteStatus::default()
        }
    }

    async fn bounded<T>(future: impl Future<Output = T>) -> T {
        timeout(Duration::from_secs(2), future)
            .await
            .expect("scripted execution should not hang")
    }

    fn stream_failure(failure: Failure, code: tonic::Code) -> Status {
        let Failure::Stream(status) = failure else {
            panic!("expected a stream failure, got {failure:?}");
        };
        assert_eq!(status.code(), code);
        status
    }

    #[tokio::test]
    async fn command_forwards_both_streams_and_returns_zero_or_nonzero() {
        for code in [0, 23] {
            let (attachment, mut script, dropped) = scripted();
            let (sender, mut receiver) = mpsc::channel(1);
            let operation = async {
                let result = command(attachment, &sender).await;
                drop(sender);
                result
            };
            let driver = async {
                script.stdout.write_all(b"out").await.unwrap();
                script.stderr.write_all(b"err").await.unwrap();
                script.status.send(exit_status(code)).unwrap();
                drop(script.stdout);
                drop(script.stderr);
            };
            let consume = async {
                let mut stdout = Vec::new();
                let mut stderr = Vec::new();
                while let Some(message) = receiver.recv().await {
                    let Some(exec_output::Output::Data(chunk)) = message.unwrap().output else {
                        panic!("command must not send terminal messages");
                    };
                    if chunk.stream == data_chunk::Stream::Stdout as i32 {
                        stdout.extend(chunk.data);
                    } else {
                        stderr.extend(chunk.data);
                    }
                }
                (stdout, stderr)
            };
            let (result, (), (stdout, stderr)) =
                bounded(async { tokio::join!(operation, driver, consume) }).await;
            assert_eq!(result.unwrap(), code);
            assert_eq!(stdout, b"out");
            assert_eq!(stderr, b"err");
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn download_requires_zero_and_preserves_a_bounded_lossy_stderr_tail() {
        for code in [0, 2] {
            let (attachment, mut script, dropped) = scripted();
            let (sender, mut receiver) = mpsc::channel(1);
            let mut diagnostic = vec![b'x'; 100];
            diagnostic.extend_from_slice("\u{20ac}".as_bytes());
            diagnostic.extend(vec![b'z'; STDERR_LIMIT - 2]);
            let operation = async {
                let result = download_archive(attachment, &sender).await;
                drop(sender);
                result
            };
            let driver = async {
                script.stdout.write_all(b"archive bytes").await.unwrap();
                script.stderr.write_all(&diagnostic).await.unwrap();
                script.status.send(exit_status(code)).unwrap();
                drop(script.stdout);
                drop(script.stderr);
            };
            let consume = async {
                let mut archive = Vec::new();
                while let Some(chunk) = receiver.recv().await {
                    archive.extend(chunk.unwrap().data);
                }
                archive
            };
            let (result, (), archive) =
                bounded(async { tokio::join!(operation, driver, consume) }).await;
            assert_eq!(archive, b"archive bytes");
            if code == 0 {
                result.unwrap();
            } else {
                let status = stream_failure(result.unwrap_err(), tonic::Code::Internal);
                assert!(status.message().contains("tar exited with status 2"));
                assert!(
                    status
                        .message()
                        .contains("[stderr truncated; showing last 16384 bytes]")
                );
                assert!(status.message().contains('\u{fffd}'));
                assert!(!status.message().contains(&"x".repeat(100)));
                assert!(status.message().ends_with(&"z".repeat(STDERR_LIMIT - 2)));
            }
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn upload_drains_both_outputs_concurrently_and_requires_zero() {
        for code in [0, 3] {
            let (attachment, mut script, dropped) = scripted();
            let first = vec![b'a'; 4096];
            let remaining = stream::iter([Ok(CopyInChunk {
                data: vec![b'b'; 4096],
                ..CopyInChunk::default()
            })]);
            let operation = upload_archive(attachment, first, remaining);
            let driver = async {
                let mut uploaded = Vec::new();
                let stdout = vec![b'o'; 32 * 1024];
                let stderr = vec![b'e'; 32 * 1024];
                let (input, out, err) = tokio::join!(
                    script.stdin.read_to_end(&mut uploaded),
                    script.stdout.write_all(&stdout),
                    script.stderr.write_all(&stderr),
                );
                assert_eq!(input.unwrap(), 8192);
                out.unwrap();
                err.unwrap();
                script.status.send(exit_status(code)).unwrap();
                drop(script.stdout);
                drop(script.stderr);
                uploaded
            };
            let (result, uploaded) = bounded(async { tokio::join!(operation, driver) }).await;
            assert_eq!(&uploaded[..4096], vec![b'a'; 4096]);
            assert_eq!(&uploaded[4096..], vec![b'b'; 4096]);
            if code == 0 {
                assert_eq!(result.unwrap(), 8192);
            } else {
                let status = stream_failure(result.unwrap_err(), tonic::Code::Internal);
                assert!(status.message().contains("tar exited with status 3"));
                assert!(
                    status
                        .message()
                        .contains("[stderr truncated; showing last 16384 bytes]")
                );
                assert!(status.message().ends_with(&"e".repeat(STDERR_LIMIT)));
            }
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn missing_status_is_remote_for_all_operations() {
        let (attachment, script, dropped) = scripted();
        drop(script);
        let (sender, _receiver) = mpsc::channel(1);
        let failure = bounded(command(attachment, &sender)).await.unwrap_err();
        assert!(
            matches!(failure, Failure::Remote(message) if message.contains("without a status"))
        );
        assert!(dropped.load(Ordering::SeqCst));

        let (attachment, script, dropped) = scripted();
        drop(script);
        let (sender, _receiver) = mpsc::channel(1);
        let failure = bounded(download_archive(attachment, &sender))
            .await
            .unwrap_err();
        assert!(
            matches!(failure, Failure::Remote(message) if message.contains("without a status"))
        );
        assert!(dropped.load(Ordering::SeqCst));

        let (attachment, script, dropped) = scripted();
        drop(script.stdout);
        drop(script.stderr);
        drop(script.status);
        let failure = bounded(upload_archive(attachment, Vec::new(), stream::empty()))
            .await
            .unwrap_err();
        assert!(
            matches!(failure, Failure::Remote(message) if message.contains("without a status"))
        );
        assert!(dropped.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn malformed_remote_status_preserves_stopped_container_detail() {
        let malformed = RemoteStatus {
            reason: Some("NonZeroExitCode".to_owned()),
            details: Some(StatusDetails {
                causes: Some(vec![StatusCause {
                    reason: Some("ExitCode".to_owned()),
                    message: Some("not a number".to_owned()),
                    ..StatusCause::default()
                }]),
                ..StatusDetails::default()
            }),
            ..RemoteStatus::default()
        };
        let stopped = RemoteStatus {
            message: Some("cannot exec in a stopped state".to_owned()),
            ..RemoteStatus::default()
        };
        for (remote, expected) in [
            (malformed, "invalid Kubernetes exit code"),
            (stopped, "cannot exec in a stopped state"),
        ] {
            let (attachment, script, dropped) = scripted();
            script.status.send(remote).unwrap();
            drop(script.stdout);
            drop(script.stderr);
            let (sender, _receiver) = mpsc::channel(1);
            let Failure::Remote(message) = bounded(command(attachment, &sender)).await.unwrap_err()
            else {
                panic!("remote status failure must remain distinguishable");
            };
            assert!(message.contains(expected));
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn remote_failure_interrupts_pending_upload_source_or_write() {
        for blocked_write in [false, true] {
            for (remote, remote_failure) in [
                (Some(exit_status(9)), false),
                (None, true),
                (Some(RemoteStatus::default()), true),
            ] {
                let (attachment, script, dropped) = scripted();
                let first = if blocked_write {
                    vec![b'a'; 128]
                } else {
                    Vec::new()
                };
                let mut operation = Box::pin(upload_archive(attachment, first, stream::pending()));
                assert!(futures::poll!(operation.as_mut()).is_pending());
                assert!(!dropped.load(Ordering::SeqCst));
                if let Some(status) = remote {
                    script.status.send(status).unwrap();
                } else {
                    drop(script.status);
                }
                drop(script.stdout);
                drop(script.stderr);
                let failure = bounded(operation).await.unwrap_err();
                if remote_failure {
                    assert!(matches!(failure, Failure::Remote(_)));
                } else {
                    let status = stream_failure(failure, tonic::Code::Internal);
                    assert!(status.message().contains("tar exited with status 9"));
                }
                assert!(dropped.load(Ordering::SeqCst));
                let mut stdin = script.stdin;
                let mut bytes = Vec::new();
                bounded(stdin.read_to_end(&mut bytes)).await.unwrap();
                assert!(bytes.len() <= 8);
            }
        }
    }

    #[tokio::test]
    async fn upload_zero_exit_before_source_eof_waits_then_succeeds() {
        let (attachment, mut script, dropped) = scripted();
        let (chunks, receiver) = mpsc::channel(1);
        // Two zero-filled blocks are a complete empty tar archive.
        let archive = vec![0; 1024];
        let mut uploaded = vec![0; archive.len()];
        let mut operation = Box::pin(upload_archive(
            attachment,
            archive.clone(),
            ReceiverStream::new(receiver),
        ));
        let read = Box::pin(script.stdin.read_exact(&mut uploaded));
        assert!(matches!(
            bounded(select(operation.as_mut(), read)).await,
            Either::Right((Ok(_), _))
        ));
        assert_eq!(uploaded, archive);
        assert!(futures::poll!(operation.as_mut()).is_pending());
        script.status.send(exit_status(0)).unwrap();
        drop(script.stdin);
        drop(script.stdout);
        drop(script.stderr);
        assert!(futures::poll!(operation.as_mut()).is_pending());
        assert!(!dropped.load(Ordering::SeqCst));

        drop(chunks);
        assert_eq!(bounded(operation).await.unwrap(), archive.len());
        assert!(dropped.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn upload_zero_exit_keeps_pending_input_until_future_drop() {
        for blocked_write in [false, true] {
            let (attachment, script, dropped) = scripted();
            let first = if blocked_write {
                vec![b'a'; 128]
            } else {
                Vec::new()
            };
            let mut operation = Box::pin(upload_archive(attachment, first, stream::pending()));
            assert!(futures::poll!(operation.as_mut()).is_pending());
            script.status.send(exit_status(0)).unwrap();
            drop(script.stdout);
            drop(script.stderr);
            assert!(futures::poll!(operation.as_mut()).is_pending());
            assert!(!dropped.load(Ordering::SeqCst));
            drop(operation);
            assert!(dropped.load(Ordering::SeqCst));
            let mut stdin = script.stdin;
            let mut bytes = Vec::new();
            bounded(stdin.read_to_end(&mut bytes)).await.unwrap();
            assert!(bytes.len() <= 8);
        }
    }

    #[tokio::test]
    async fn upload_zero_exit_does_not_discard_later_payload_for_closed_stdin() {
        let (attachment, script, dropped) = scripted();
        let (chunks, receiver) = mpsc::channel(1);
        let mut operation = Box::pin(upload_archive(
            attachment,
            Vec::new(),
            ReceiverStream::new(receiver),
        ));
        assert!(futures::poll!(operation.as_mut()).is_pending());
        script.status.send(exit_status(0)).unwrap();
        drop(script.stdin);
        drop(script.stdout);
        drop(script.stderr);
        assert!(futures::poll!(operation.as_mut()).is_pending());
        chunks
            .send(Ok(CopyInChunk {
                data: vec![b'a'],
                ..CopyInChunk::default()
            }))
            .await
            .unwrap();
        drop(chunks);
        let status = stream_failure(bounded(operation).await.unwrap_err(), tonic::Code::Internal);
        assert!(status.message().contains("write CopyIn chunk"));
        assert!(dropped.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn upload_ready_nonzero_does_not_replace_forbidden_metadata_error() {
        for (environment_id, dest_path) in
            [(Some(String::new()), None), (None, Some(String::new()))]
        {
            let (attachment, script, dropped) = scripted();
            script.status.send(exit_status(2)).unwrap();
            let remaining = stream::iter([Ok(CopyInChunk {
                environment_id,
                dest_path,
                data: Vec::new(),
            })]);
            let status = stream_failure(
                bounded(upload_archive(attachment, Vec::new(), remaining))
                    .await
                    .unwrap_err(),
                tonic::Code::InvalidArgument,
            );
            assert!(status.message().contains("only the first CopyIn chunk"));
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn upload_ready_nonzero_does_not_replace_source_error() {
        for code in [tonic::Code::Cancelled, tonic::Code::DataLoss] {
            let (attachment, script, dropped) = scripted();
            script.status.send(exit_status(2)).unwrap();
            let remaining = stream::iter([Err(Status::new(code, "archive input failed"))]);
            let status = stream_failure(
                bounded(upload_archive(attachment, Vec::new(), remaining))
                    .await
                    .unwrap_err(),
                code,
            );
            assert!(status.message().contains("archive input failed"));
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn upload_broken_pipe_preserves_ready_tar_failure_and_stderr() {
        let (attachment, mut script, dropped) = scripted();
        script.stderr.write_all(b"bad tar").await.unwrap();
        script.status.send(exit_status(2)).unwrap();
        drop(script.stdin);
        drop(script.stdout);
        drop(script.stderr);

        let status = stream_failure(
            bounded(upload_archive(attachment, vec![b'a'], stream::pending()))
                .await
                .unwrap_err(),
            tonic::Code::Internal,
        );
        assert!(status.message().contains("tar exited with status 2"));
        assert!(status.message().contains("bad tar"));
        assert!(dropped.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn upload_errors_do_not_wait_for_status_or_output_eof() {
        for remote_code in [None, Some(0)] {
            for broken_pipe in [false, true] {
                let (attachment, mut script, dropped) = scripted();
                script.stderr.write_all(b"context").await.unwrap();
                if let Some(code) = remote_code {
                    script.status.send(exit_status(code)).unwrap();
                }
                let first = if broken_pipe {
                    drop(script.stdin);
                    vec![b'a']
                } else {
                    Vec::new()
                };
                let remaining = stream::iter([Err(Status::data_loss("archive input failed"))]);
                let expected = if broken_pipe {
                    tonic::Code::Internal
                } else {
                    tonic::Code::DataLoss
                };
                let status = stream_failure(
                    bounded(upload_archive(attachment, first, remaining))
                        .await
                        .unwrap_err(),
                    expected,
                );
                assert!(status.message().contains("context"));
                assert!(dropped.load(Ordering::SeqCst));
            }
        }
    }

    #[tokio::test]
    async fn upload_rejects_metadata_on_later_chunks() {
        for (environment_id, dest_path) in
            [(Some(String::new()), None), (None, Some(String::new()))]
        {
            let (attachment, _script, dropped) = scripted();
            let remaining = stream::iter([Ok(CopyInChunk {
                environment_id,
                dest_path,
                data: Vec::new(),
            })]);
            let status = stream_failure(
                bounded(upload_archive(attachment, Vec::new(), remaining))
                    .await
                    .unwrap_err(),
                tonic::Code::InvalidArgument,
            );
            assert!(status.message().contains("only the first CopyIn chunk"));
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn upload_preserves_input_error_status() {
        let (attachment, _script, dropped) = scripted();
        let remaining = stream::iter([Err(Status::data_loss("archive input failed"))]);
        let status = stream_failure(
            bounded(upload_archive(attachment, Vec::new(), remaining))
                .await
                .unwrap_err(),
            tonic::Code::DataLoss,
        );
        assert!(status.message().contains("archive input failed"));
        assert!(dropped.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn command_cancels_during_silence_and_when_output_is_backpressured() {
        for blocked in [false, true] {
            let (attachment, mut script, dropped) = scripted();
            let (sender, receiver) = mpsc::channel(1);
            if blocked {
                sender.send(Ok(ExecOutput::default())).await.unwrap();
                script.stdout.write_all(b"out").await.unwrap();
                script.status.send(exit_status(0)).unwrap();
                drop(script.stdout);
                drop(script.stderr);
            }
            let mut operation = Box::pin(command(attachment, &sender));
            assert!(futures::poll!(operation.as_mut()).is_pending());
            drop(receiver);
            stream_failure(
                bounded(operation).await.unwrap_err(),
                tonic::Code::Cancelled,
            );
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn download_cancels_during_silence_and_when_output_is_backpressured() {
        for blocked in [false, true] {
            let (attachment, mut script, dropped) = scripted();
            let (sender, receiver) = mpsc::channel(1);
            if blocked {
                sender.send(Ok(CopyOutChunk::default())).await.unwrap();
                script.stdout.write_all(b"tar").await.unwrap();
                script.status.send(exit_status(0)).unwrap();
                drop(script.stdout);
                drop(script.stderr);
            }
            let mut operation = Box::pin(download_archive(attachment, &sender));
            assert!(futures::poll!(operation.as_mut()).is_pending());
            drop(receiver);
            stream_failure(
                bounded(operation).await.unwrap_err(),
                tonic::Code::Cancelled,
            );
            assert!(dropped.load(Ordering::SeqCst));
        }
    }

    #[tokio::test]
    async fn abrupt_future_drop_releases_every_operation_owner_and_endpoint() {
        let (attachment, script, dropped) = scripted();
        let (sender, _receiver) = mpsc::channel(1);
        let mut operation = Box::pin(command(attachment, &sender));
        assert!(futures::poll!(operation.as_mut()).is_pending());
        drop(operation);
        assert!(dropped.load(Ordering::SeqCst));
        assert!(script.status.send(exit_status(0)).is_err());

        let (attachment, script, dropped) = scripted();
        let (sender, _receiver) = mpsc::channel(1);
        let mut operation = Box::pin(download_archive(attachment, &sender));
        assert!(futures::poll!(operation.as_mut()).is_pending());
        drop(operation);
        assert!(dropped.load(Ordering::SeqCst));
        assert!(script.status.send(exit_status(0)).is_err());

        let (attachment, mut script, dropped) = scripted();
        let mut operation = Box::pin(upload_archive(
            attachment,
            vec![b'a'; 128],
            stream::pending(),
        ));
        assert!(futures::poll!(operation.as_mut()).is_pending());
        drop(operation);
        assert!(dropped.load(Ordering::SeqCst));
        assert!(script.status.send(exit_status(0)).is_err());
        let mut bytes = Vec::new();
        bounded(script.stdin.read_to_end(&mut bytes)).await.unwrap();
        assert_eq!(bytes.len(), 8);
        assert!(script.stdout.write_all(b"out").await.is_err());
        assert!(script.stderr.write_all(b"err").await.is_err());
    }

    #[tokio::test]
    async fn completion_waits_for_output_drain_after_status() {
        let (attachment, mut script, dropped) = scripted();
        let (sender, mut receiver) = mpsc::channel(1);
        script.status.send(exit_status(0)).unwrap();
        let mut operation = Box::pin(command(attachment, &sender));
        assert!(futures::poll!(operation.as_mut()).is_pending());
        assert!(!dropped.load(Ordering::SeqCst));
        script.stdout.write_all(b"late").await.unwrap();
        drop(script.stdout);
        drop(script.stderr);
        assert_eq!(bounded(operation).await.unwrap(), 0);
        let Some(exec_output::Output::Data(chunk)) = receiver.recv().await.unwrap().unwrap().output
        else {
            panic!("expected the final stdout bytes");
        };
        assert_eq!(chunk.data, b"late");
        assert!(dropped.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn missing_required_endpoints_fail_without_waiting_for_status() {
        let (mut attachment, _script, dropped) = scripted();
        attachment.stdout = None;
        let (sender, _receiver) = mpsc::channel(1);
        let status = stream_failure(
            bounded(command(attachment, &sender)).await.unwrap_err(),
            tonic::Code::Internal,
        );
        assert!(status.message().contains("did not provide stdout"));
        assert!(dropped.load(Ordering::SeqCst));

        let (mut attachment, _script, dropped) = scripted();
        attachment.stdin = None;
        let status = stream_failure(
            bounded(upload_archive(attachment, Vec::new(), stream::empty()))
                .await
                .unwrap_err(),
            tonic::Code::Internal,
        );
        assert!(status.message().contains("did not provide stdin"));
        assert!(dropped.load(Ordering::SeqCst));
    }
}