use bytes::Bytes;
use futures_util::StreamExt;
use opendal::FuturesBytesStream;
use remora_stream::{SourceBody, StreamError, StreamSide};

pub struct ObjectSourceBody {
    stream: FuturesBytesStream,
    pending: Bytes,
    pending_offset: usize,
    remaining: u64,
    finished: bool,
}

impl ObjectSourceBody {
    pub(crate) fn new(stream: FuturesBytesStream, expected_length: u64) -> Self {
        Self {
            stream,
            pending: Bytes::new(),
            pending_offset: 0,
            remaining: expected_length,
            finished: false,
        }
    }

    fn copy_pending(&mut self, output: &mut [u8]) -> usize {
        let available = &self.pending[self.pending_offset..];
        let count = available.len().min(output.len());
        output[..count].copy_from_slice(&available[..count]);
        self.pending_offset += count;
        self.remaining -= count as u64;
        if self.pending_offset == self.pending.len() {
            self.pending = Bytes::new();
            self.pending_offset = 0;
        }
        count
    }
}

impl SourceBody for ObjectSourceBody {
    async fn read(&mut self, output: &mut [u8]) -> Result<usize, StreamError> {
        if output.is_empty() {
            return Ok(0);
        }
        if self.pending_offset < self.pending.len() {
            return Ok(self.copy_pending(output));
        }
        if self.finished {
            return Ok(0);
        }
        match self.stream.next().await {
            Some(Ok(chunk)) if chunk.len() as u64 <= self.remaining => {
                self.pending = chunk;
                Ok(self.copy_pending(output))
            }
            Some(Ok(_)) => Err(StreamError::SourceChanged),
            Some(Err(error)) => Err(map_stream_error(&error)),
            None if self.remaining != 0 => Err(StreamError::transient_source(
                "object source ended before its declared length",
            )),
            None => {
                self.finished = true;
                Ok(0)
            }
        }
    }
}

fn map_stream_error(error: &std::io::Error) -> StreamError {
    let kind = error
        .get_ref()
        .and_then(|source| source.downcast_ref::<opendal::Error>())
        .map(opendal::Error::kind);
    match kind {
        Some(opendal::ErrorKind::PermissionDenied) => StreamError::Authorization {
            side: StreamSide::Source,
        },
        Some(
            opendal::ErrorKind::ConditionNotMatch
            | opendal::ErrorKind::NotFound
            | opendal::ErrorKind::RangeNotSatisfied,
        ) => StreamError::SourceChanged,
        _ => StreamError::transient_source("object source read failed"),
    }
}