pub mod svg;

use std::{io::Cursor, str::FromStr};

use anyhow::{Context, Result, bail};
use async_trait::async_trait;
use base64::{display::Base64Display, engine::general_purpose::STANDARD};
use bincode::{Decode, Encode};
use image::{
    DynamicImage, GenericImageView, ImageEncoder, ImageFormat,
    codecs::{
        bmp::BmpEncoder,
        ico::IcoEncoder,
        jpeg::JpegEncoder,
        png::{CompressionType, PngEncoder},
    },
    imageops::FilterType,
};
use mime::Mime;
use turbo_rcstr::rcstr;
use turbo_tasks::{
    NonLocalValue, PrettyPrintError, ResolvedVc, Vc, debug::ValueDebugFormat, trace::TraceRawVcs,
};
use turbo_tasks_fs::{File, FileContent, FileSystemPath};
use turbopack_core::{
    issue::{Issue, IssueExt, IssueSeverity, IssueSource, IssueStage, StyledString},
    source::Source,
};

use self::svg::calculate;

/// Small placeholder version of the image.
#[derive(PartialEq, Eq, TraceRawVcs, ValueDebugFormat, NonLocalValue, Encode, Decode)]
pub struct BlurPlaceholder {
    pub data_url: String,
    pub width: u32,
    pub height: u32,
}

impl BlurPlaceholder {
    pub fn fallback() -> Self {
        BlurPlaceholder {
            data_url: "data:image/gif;base64,R0lGODlhAQABAIAAAP///\
                       wAAACH5BAEAAAAALAAAAAABAAEAAAICRAEAOw=="
                .to_string(),
            width: 1,
            height: 1,
        }
    }
}

/// Gathered meta information about an image.
#[allow(clippy::manual_non_exhaustive)]
#[turbo_tasks::value]
#[derive(Default)]
#[non_exhaustive]
pub struct ImageMetaData {
    pub width: u32,
    pub height: u32,
    #[turbo_tasks(trace_ignore, debug_ignore)]
    #[bincode(with = "turbo_bincode::mime_option")]
    pub mime_type: Option<Mime>,
    pub blur_placeholder: Option<BlurPlaceholder>,
}

impl ImageMetaData {
    pub fn fallback_value(mime_type: Option<Mime>) -> Self {
        ImageMetaData {
            width: 100,
            height: 100,
            mime_type,
            blur_placeholder: Some(BlurPlaceholder::fallback()),
        }
    }
}

/// Options for generating a blur placeholder.
#[turbo_tasks::value(shared)]
pub struct BlurPlaceholderOptions {
    pub quality: u8,
    pub size: u32,
}

fn extension_to_image_format(extension: &str) -> Option<ImageFormat> {
    Some(match extension {
        "avif" => ImageFormat::Avif,
        "jpg" | "jpeg" => ImageFormat::Jpeg,
        "png" => ImageFormat::Png,
        "gif" => ImageFormat::Gif,
        "webp" => ImageFormat::WebP,
        "tif" | "tiff" => ImageFormat::Tiff,
        "tga" => ImageFormat::Tga,
        "dds" => ImageFormat::Dds,
        "bmp" => ImageFormat::Bmp,
        "ico" => ImageFormat::Ico,
        "hdr" => ImageFormat::Hdr,
        "exr" => ImageFormat::OpenExr,
        "pbm" | "pam" | "ppm" | "pgm" => ImageFormat::Pnm,
        "ff" | "farbfeld" => ImageFormat::Farbfeld,
        "qoi" => ImageFormat::Qoi,
        _ => return None,
    })
}

fn result_to_issue<T>(source: ResolvedVc<Box<dyn Source>>, result: Result<T>) -> Option<T> {
    match result {
        Ok(r) => Some(r),
        Err(err) => {
            ImageProcessingIssue {
                message: StyledString::Text(format!("{}", PrettyPrintError(&err)).into())
                    .resolved_cell(),
                issue_severity: None,
                title: None,
                source: IssueSource::from_source_only(source),
            }
            .resolved_cell()
            .emit();
            None
        }
    }
}

fn load_image(
    path: ResolvedVc<Box<dyn Source>>,
    bytes: &[u8],
    extension: Option<&str>,
) -> Option<(ImageBuffer, Option<ImageFormat>)> {
    result_to_issue(path, load_image_internal(path, bytes, extension))
}

/// Type of raw image buffer read by reader from `load_image`.
/// If the image could not be decoded, the raw bytes are returned.
enum ImageBuffer {
    Raw(Vec<u8>),
    Decoded(image::DynamicImage),
}

fn load_image_internal(
    image: ResolvedVc<Box<dyn Source>>,
    bytes: &[u8],
    extension: Option<&str>,
) -> Result<(ImageBuffer, Option<ImageFormat>)> {
    let reader = image::ImageReader::new(Cursor::new(&bytes));
    let mut reader = reader
        .with_guessed_format()
        .context("unable to determine image format from file content")?;
    let mut format = reader.format();
    if format.is_none()
        && let Some(ext) = extension
        && let Some(new_format) = extension_to_image_format(ext)
    {
        format = Some(new_format);
        reader.set_format(new_format);
    }

    // [NOTE]
    // Workaround for missing codec supports in Turbopack,
    // Instead of erroring out the whole build, emitting raw image bytes as-is
    // (Not applying resize, not applying optimization or anything else)
    // and expect a browser decodes it.
    // This is a stop gap until we have proper encoding/decoding in majority of the
    // platforms

    #[cfg(not(feature = "avif"))]
    if matches!(format, Some(ImageFormat::Avif)) {
        ImageProcessingIssue {
            source: IssueSource::from_source_only(image),
            message: StyledString::Text(rcstr!(
                "This version of Turbopack does not support AVIF images, will emit without \
                 optimization or encoding"
            ))
            .resolved_cell(),
            title: Some(StyledString::Text(rcstr!("AVIF image not supported")).resolved_cell()),
            issue_severity: Some(IssueSeverity::Warning),
        }
        .resolved_cell()
        .emit();
        return Ok((ImageBuffer::Raw(bytes.to_vec()), format));
    }

    #[cfg(not(feature = "webp"))]
    if matches!(format, Some(ImageFormat::WebP)) {
        ImageProcessingIssue {
            source: IssueSource::from_source_only(image),
            message: StyledString::Text(rcstr!(
                "This version of Turbopack does not support WEBP images, will emit without \
                 optimization or encoding"
            ))
            .resolved_cell(),
            title: Some(StyledString::Text(rcstr!("WEBP image not supported")).resolved_cell()),
            issue_severity: Some(IssueSeverity::Warning),
        }
        .resolved_cell()
        .emit();
        return Ok((ImageBuffer::Raw(bytes.to_vec()), format));
    }

    let image = reader.decode().context("unable to decode image data")?;
    Ok((ImageBuffer::Decoded(image), format))
}

fn compute_blur_data(
    path: ResolvedVc<Box<dyn Source>>,
    image: image::DynamicImage,
    format: ImageFormat,
    options: &BlurPlaceholderOptions,
) -> Option<BlurPlaceholder> {
    match compute_blur_data_internal(image, format, options)
        .context("unable to compute blur placeholder")
    {
        Ok(r) => Some(r),
        Err(err) => {
            ImageProcessingIssue {
                source: IssueSource::from_source_only(path),
                message: StyledString::Text(format!("{}", PrettyPrintError(&err)).into())
                    .resolved_cell(),
                issue_severity: None,
                title: None,
            }
            .resolved_cell()
            .emit();
            Some(BlurPlaceholder::fallback())
        }
    }
}

fn encode_image(image: DynamicImage, format: ImageFormat, quality: u8) -> Result<(Vec<u8>, Mime)> {
    let mut buf = Vec::new();
    let (width, height) = image.dimensions();

    Ok(match format {
        ImageFormat::Png => {
            PngEncoder::new_with_quality(
                &mut buf,
                CompressionType::Best,
                image::codecs::png::FilterType::NoFilter,
            )
            .write_image(image.as_bytes(), width, height, image.color().into())?;
            (buf, mime::IMAGE_PNG)
        }
        ImageFormat::Jpeg => {
            JpegEncoder::new_with_quality(&mut buf, quality).write_image(
                image.as_bytes(),
                width,
                height,
                image.color().into(),
            )?;
            (buf, mime::IMAGE_JPEG)
        }
        ImageFormat::Ico => {
            IcoEncoder::new(&mut buf).write_image(
                image.as_bytes(),
                width,
                height,
                image.color().into(),
            )?;
            // mime does not support typed IMAGE_X_ICO yet
            (buf, Mime::from_str("image/x-icon")?)
        }
        ImageFormat::Bmp => {
            BmpEncoder::new(&mut buf).write_image(
                image.as_bytes(),
                width,
                height,
                image.color().into(),
            )?;
            (buf, mime::IMAGE_BMP)
        }
        #[cfg(feature = "webp")]
        ImageFormat::WebP => {
            use image::codecs::webp::WebPEncoder;
            let encoder = WebPEncoder::new_lossless(&mut buf);
            encoder.encode(image.as_bytes(), width, height, image.color().into())?;

            (buf, Mime::from_str("image/webp")?)
        }
        #[cfg(feature = "avif")]
        ImageFormat::Avif => {
            use image::codecs::avif::AvifEncoder;
            AvifEncoder::new_with_speed_quality(&mut buf, 6, quality).write_image(
                image.as_bytes(),
                width,
                height,
                image.color().into(),
            )?;
            (buf, Mime::from_str("image/avif")?)
        }
        _ => bail!(
            "Encoding for image format {:?} has not been compiled into the current build",
            format
        ),
    })
}

fn compute_blur_data_internal(
    image: image::DynamicImage,
    format: ImageFormat,
    options: &BlurPlaceholderOptions,
) -> Result<BlurPlaceholder> {
    let small_image = image.resize(options.size, options.size, FilterType::Triangle);
    let width = small_image.width();
    let height = small_image.height();
    let (data, mime) = encode_image(small_image, format, options.quality)?;
    let data_url = format!(
        "data:{mime};base64,{}",
        Base64Display::new(&data, &STANDARD)
    );

    Ok(BlurPlaceholder {
        data_url,
        width,
        height,
    })
}

fn image_format_to_mime_type(format: ImageFormat) -> Result<Option<Mime>> {
    Ok(match format {
        ImageFormat::Png => Some(mime::IMAGE_PNG),
        ImageFormat::Jpeg => Some(mime::IMAGE_JPEG),
        ImageFormat::WebP => Some(Mime::from_str("image/webp")?),
        ImageFormat::Avif => Some(Mime::from_str("image/avif")?),
        ImageFormat::Bmp => Some(mime::IMAGE_BMP),
        ImageFormat::Dds => Some(Mime::from_str("image/vnd-ms.dds")?),
        ImageFormat::Farbfeld => Some(mime::APPLICATION_OCTET_STREAM),
        ImageFormat::Gif => Some(mime::IMAGE_GIF),
        ImageFormat::Hdr => Some(Mime::from_str("image/vnd.radiance")?),
        ImageFormat::Ico => Some(Mime::from_str("image/x-icon")?),
        ImageFormat::OpenExr => Some(Mime::from_str("image/x-exr")?),
        ImageFormat::Pnm => Some(Mime::from_str("image/x-portable-anymap")?),
        ImageFormat::Qoi => Some(mime::APPLICATION_OCTET_STREAM),
        ImageFormat::Tga => Some(Mime::from_str("image/x-tga")?),
        ImageFormat::Tiff => Some(Mime::from_str("image/tiff")?),
        _ => None,
    })
}

/// Analyze an image and return meta information about it.
/// Optionally computes a blur placeholder.
#[turbo_tasks::function]
pub async fn get_meta_data(
    image: ResolvedVc<Box<dyn Source>>,
    content: Vc<FileContent>,
    blur_placeholder: Option<Vc<BlurPlaceholderOptions>>,
) -> Result<Vc<ImageMetaData>> {
    let FileContent::Content(content) = &*content.await? else {
        bail!("Input image not found");
    };
    let bytes = content.content().to_bytes();
    let path = image.ident().path().await?;
    let extension = path.extension();

    if extension == Some("svg") {
        let content = result_to_issue(
            image,
            std::str::from_utf8(&bytes).context("Input image is not valid utf-8"),
        );
        let Some(content) = content else {
            return Ok(ImageMetaData::fallback_value(Some(mime::IMAGE_SVG)).cell());
        };
        let info = result_to_issue(
            image,
            calculate(content).context("Failed to parse svg source code for image dimensions"),
        );
        let Some((width, height)) = info else {
            return Ok(ImageMetaData::fallback_value(Some(mime::IMAGE_SVG)).cell());
        };
        return Ok(ImageMetaData {
            width,
            height,
            mime_type: Some(mime::IMAGE_SVG),
            blur_placeholder: None,
        }
        .cell());
    }
    let Some((image_buffer, format)) = load_image(image, &bytes, extension) else {
        return Ok(ImageMetaData::fallback_value(None).cell());
    };

    match image_buffer {
        ImageBuffer::Raw(..) => Ok(ImageMetaData::fallback_value(None).cell()),
        ImageBuffer::Decoded(image_data) => {
            let (width, height) = image_data.dimensions();
            let blur_placeholder = if let Some(blur_placeholder) = blur_placeholder {
                if matches!(
                    format,
                    // list should match next/client/image.tsx
                    Some(ImageFormat::Png)
                        | Some(ImageFormat::Jpeg)
                        | Some(ImageFormat::WebP)
                        | Some(ImageFormat::Avif)
                ) {
                    compute_blur_data(
                        image,
                        image_data,
                        format.unwrap(),
                        &*blur_placeholder.await?,
                    )
                } else {
                    None
                }
            } else {
                None
            };

            Ok(ImageMetaData {
                width,
                height,
                mime_type: if let Some(format) = format {
                    image_format_to_mime_type(format)?
                } else {
                    None
                },
                blur_placeholder,
            }
            .cell())
        }
    }
}

#[turbo_tasks::function]
pub async fn optimize(
    source: ResolvedVc<Box<dyn Source>>,
    content: Vc<FileContent>,
    max_width: u32,
    max_height: u32,
    quality: u8,
) -> Result<Vc<FileContent>> {
    let FileContent::Content(content) = &*content.await? else {
        return Ok(FileContent::NotFound.cell());
    };
    let bytes = content.content().to_bytes();
    let path = source.ident().path().await?;
    let extension = path.extension();

    let Some((image, format)) = load_image(source, &bytes, extension) else {
        return Ok(FileContent::NotFound.cell());
    };
    match image {
        ImageBuffer::Raw(buffer) => {
            #[cfg(not(feature = "avif"))]
            if matches!(format, Some(ImageFormat::Avif)) {
                return Ok(FileContent::Content(
                    File::from(buffer).with_content_type(Mime::from_str("image/avif")?),
                )
                .cell());
            }

            #[cfg(not(feature = "webp"))]
            if matches!(format, Some(ImageFormat::WebP)) {
                return Ok(FileContent::Content(
                    File::from(buffer).with_content_type(Mime::from_str("image/webp")?),
                )
                .cell());
            }

            let mime_type = if let Some(format) = format {
                image_format_to_mime_type(format)?
            } else {
                None
            };

            // Falls back to image/jpeg if the format is unknown, thouogh it is not
            // technically correct
            Ok(FileContent::Content(
                File::from(buffer).with_content_type(mime_type.unwrap_or(mime::IMAGE_JPEG)),
            )
            .cell())
        }
        ImageBuffer::Decoded(image) => {
            let (width, height) = image.dimensions();
            let image = if width > max_width || height > max_height {
                image.resize(max_width, max_height, FilterType::Lanczos3)
            } else {
                image
            };

            let format = format.unwrap_or(ImageFormat::Jpeg);
            let (data, mime_type) = encode_image(image, format, quality)?;

            Ok(FileContent::Content(File::from(data).with_content_type(mime_type)).cell())
        }
    }
}

#[turbo_tasks::value]
struct ImageProcessingIssue {
    message: ResolvedVc<StyledString>,
    title: Option<ResolvedVc<StyledString>>,
    issue_severity: Option<IssueSeverity>,
    source: IssueSource,
}

#[async_trait]
#[turbo_tasks::value_impl]
impl Issue for ImageProcessingIssue {
    fn severity(&self) -> IssueSeverity {
        self.issue_severity.unwrap_or(IssueSeverity::Error)
    }

    async fn file_path(&self) -> anyhow::Result<FileSystemPath> {
        self.source.file_path().owned().await
    }

    fn stage(&self) -> IssueStage {
        IssueStage::Transform
    }

    async fn title(&self) -> anyhow::Result<StyledString> {
        Ok(match self.title {
            Some(t) => (*t.await?).clone(),
            None => StyledString::Text(rcstr!("Processing image failed")),
        })
    }

    async fn description(&self) -> anyhow::Result<Option<StyledString>> {
        Ok(Some((*self.message.await?).clone()))
    }

    fn source(&self) -> Option<IssueSource> {
        Some(self.source)
    }
}