article_scraper 3.0.0-alpha1

use futures::channel::mpsc::Sender;
use futures::{SinkExt, StreamExt};
use reqwest::{Client, Response};

use super::image_data::ImageData;
use crate::constants;
use crate::error::ImageDownloadError;
use crate::util::Util;

/// Progress of one download. The size of an image is only known once its response arrives
/// (or, without `Content-Length`, once its body is read), so the expected total grows too.
#[derive(Debug)]
pub enum Progress {
    Expected(usize),
    Downloaded(usize),
}

#[derive(Debug)]
pub struct ImageRequest {
    url: String,
    response: Response,
    content_length: Option<usize>,
    content_type: String,
}

impl ImageRequest {
    pub async fn new(url: String, client: &Client) -> Result<Self, ImageDownloadError> {
        let response = client.get(&url).send().await?;

        let content_type = Util::get_content_type(&response)?;
        // Chunked responses have no length; `download` enforces the size limit either way.
        let content_length = Util::get_content_length(&response).ok();

        if !content_type.contains("image") {
            return Err(ImageDownloadError::ContentType);
        }

        if let Some(content_length) = content_length
            && content_length > constants::MAX_IMAGE_SIZE
        {
            tracing::warn!(%url, content_length, "Image is too large");
            return Err(ImageDownloadError::TooLarge);
        }

        Ok(Self {
            url,
            response,
            content_length,
            content_type,
        })
    }

    pub async fn download(
        self,
        tx: &mut Sender<Progress>,
    ) -> Result<ImageData, ImageDownloadError> {
        if let Some(content_length) = self.content_length {
            _ = tx.send(Progress::Expected(content_length)).await;
        }

        let mut stream = self.response.bytes_stream();

        let mut result = Vec::with_capacity(self.content_length.unwrap_or(0));
        while let Some(item) = stream.next().await {
            let chunk = item?;
            if result.len() + chunk.len() > constants::MAX_IMAGE_SIZE {
                tracing::warn!(url = %self.url, "Image is too large");
                return Err(ImageDownloadError::TooLarge);
            }
            if self.content_length.is_none() {
                _ = tx.send(Progress::Expected(chunk.len())).await;
            }
            _ = tx.send(Progress::Downloaded(chunk.len())).await;
            result.extend_from_slice(&chunk);
        }

        Ok(ImageData {
            url: self.url,
            data: result,
            content_type: self.content_type,
        })
    }
}