//! Image handling.

mod pdf;
mod raster;
mod svg;

pub use self::pdf::PdfImage;
pub use self::raster::{
    ExchangeFormat, PixelEncoding, PixelFormat, RasterFormat, RasterImage,
};
pub use self::svg::SvgImage;

use std::fmt::{self, Debug, Formatter};
use std::num::NonZeroUsize;
use std::sync::Arc;

use ecow::EcoString;
use hayro_syntax::LoadPdfError;
use typst_syntax::{Spanned, VirtualPath};
use typst_utils::{LazyHash, NonZeroExt};

use crate::diag::{At, LoadedWithin, SourceResult, StrResult, bail, warning};
use crate::engine::Engine;
use crate::foundations::{
    Bytes, Cast, Derived, Packed, Smart, StyleChain, Synthesize, cast, elem,
};
use crate::introspection::{Locatable, Tagged};
use crate::layout::{Length, Rel, Sizing};
use crate::loading::{DataSource, Load, Loaded};
use crate::model::Figurable;
use crate::text::{LocalName, Locale, families};
use crate::visualize::image::pdf::PdfDocument;

/// A raster or vector graphic.
///
/// You can wrap the image in a @figure to give it a number and caption.
///
/// Like most elements, images are _block-level_ by default and thus do not
/// integrate themselves into adjacent paragraphs. To force an image to become
/// inline, put it into a @box.
///
/// = Example <example>
/// ```example
/// #figure(
///   image("molecular.jpg", width: 80%),
///   caption: [
///     A step in the molecular testing
///     pipeline of our lab.
///   ],
/// )
/// ```
#[elem(Locatable, Tagged, Synthesize, LocalName, Figurable)]
pub struct ImageElem {
    /// A path to an image file or raw bytes making up an image in one of the
    /// supported @image.format[formats].
    ///
    /// Bytes can be used to specify raw pixel data in a row-major,
    /// left-to-right, top-to-bottom format.
    ///
    /// ```example
    /// #let original = read("diagram.svg")
    /// #let changed = original.replace(
    ///   "#2B80FF", // blue
    ///   green.to-hex(),
    /// )
    ///
    /// #image(bytes(original))
    /// #image(bytes(changed))
    /// ```
    #[required]
    #[parse(
        let source = args.expect::<Spanned<DataSource>>("source")?;
        let loaded = source.load(engine.world)?;
        Derived::new(source.v, loaded)
    )]
    pub source: Derived<DataSource, Loaded>,

    /// The image's format.
    ///
    /// By default, the format is detected automatically. Typically, you thus
    /// only need to specify this when providing raw bytes as the
    /// @image.source[`source`] (even then, Typst will try to figure out the
    /// format automatically, but that's not always possible).
    ///
    /// Supported formats are `{"png"}`, `{"jpg"}`, `{"gif"}`, `{"svg"}`,
    /// `{"pdf"}`, `{"webp"}` as well as raw pixel data.
    ///
    /// Note that several restrictions apply when using PDF files as images:
    ///
    /// - When exporting to PDF, any PDF image file used must have a version
    ///   equal to or lower than the
    ///   @pdf:pdf-versions[export target PDF version].
    /// - PDF files as images are currently not supported when exporting with a
    ///   specific PDF standard, like PDF/A-3 or PDF/UA-1. In these cases, you
    ///   can instead use SVGs to embed vector images.
    /// - The image file must not be password-protected.
    /// - Tags in your PDF image will not be preserved. Instead, you must
    ///   provide an @image.alt[alternative description] to make the image
    ///   accessible.
    ///
    /// When providing raw pixel data as the `source`, you must specify a
    /// dictionary with the following keys as the `format`:
    /// - `encoding` (@str[str]): The encoding of the pixel data. One of:
    ///   - `{"rgb8"}` (three 8-bit channels: red, green, blue)
    ///   - `{"rgba8"}` (four 8-bit channels: red, green, blue, alpha)
    ///   - `{"luma8"}` (one 8-bit channel)
    ///   - `{"lumaa8"}` (two 8-bit channels: luma and alpha)
    /// - `width` (@int[int]): The pixel width of the image.
    /// - `height` (@int[int]): The pixel height of the image.
    ///
    /// The pixel width multiplied by the height multiplied by the channel count
    /// for the specified encoding must then match the `source` data.
    ///
    /// ```example
    /// #image(
    ///   read(
    ///     "tetrahedron.svg",
    ///     encoding: none,
    ///   ),
    ///   format: "svg",
    ///   width: 2cm,
    /// )
    ///
    /// #image(
    ///   bytes(range(16).map(x => x * 16)),
    ///   format: (
    ///     encoding: "luma8",
    ///     width: 4,
    ///     height: 4,
    ///   ),
    ///   width: 2cm,
    /// )
    /// ```
    pub format: Smart<ImageFormat>,

    /// The width of the image.
    pub width: Smart<Rel<Length>>,

    /// The height of the image.
    pub height: Sizing,

    /// An alternative description of the image.
    ///
    /// This text is used by Assistive Technology (AT) like screen readers to
    /// describe the image to users with visual impairments.
    ///
    /// When the image is wrapped in a @figure, use this parameter rather than
    /// the @figure.alt[figure's `alt` parameter] to describe the image. The
    /// only exception to this rule is when the image and the other contents in
    /// the figure form a single semantic unit. In this case, use the figure's
    /// `alt` parameter to describe the entire composition and do not use this
    /// parameter.
    ///
    /// You can learn how to write good alternative descriptions in the
    /// @guides:accessibility:textual-representations[Accessibility Guide].
    pub alt: Option<EcoString>,

    /// The page number that should be embedded as an image. This attribute only
    /// has an effect for PDF files.
    #[default(NonZeroUsize::ONE)]
    pub page: NonZeroUsize,

    /// How the image should adjust itself to a given area (the area is defined
    /// by the `width` and `height` fields). Note that `fit` doesn't visually
    /// change anything if the area's aspect ratio is the same as the image's
    /// one.
    ///
    /// ```example
    /// #set page(width: 300pt, height: 50pt, margin: 10pt)
    /// #image("tiger.jpg", width: 100%, fit: "cover")
    /// #image("tiger.jpg", width: 100%, fit: "contain")
    /// #image("tiger.jpg", width: 100%, fit: "stretch")
    /// ```
    #[default(ImageFit::Cover)]
    pub fit: ImageFit,

    /// A hint to viewers how they should scale the image.
    ///
    /// When set to `{auto}`, the default is left up to the viewer. For PNG
    /// export, Typst will default to smooth scaling, like most PDF and SVG
    /// viewers.
    ///
    /// _Note:_ The exact look may differ across PDF viewers.
    pub scaling: Smart<ImageScaling>,

    /// An ICC profile for the image.
    ///
    /// ICC profiles define how to interpret the colors in an image. When set to
    /// `{auto}`, Typst will try to extract an ICC profile from the image.
    #[parse(match args.named::<Spanned<Smart<DataSource>>>("icc")? {
        Some(Spanned { v: Smart::Custom(source), span }) => Some(Smart::Custom({
            let loaded = Spanned::new(&source, span).load(engine.world)?;
            Derived::new(source, loaded.data)
        })),
        Some(Spanned { v: Smart::Auto, .. }) => Some(Smart::Auto),
        None => None,
    })]
    pub icc: Smart<Derived<DataSource, Bytes>>,

    /// The locale of this element (used for the alternative description).
    #[internal]
    #[synthesized]
    pub locale: Locale,
}

impl Synthesize for Packed<ImageElem> {
    fn synthesize(&mut self, _: &mut Engine, styles: StyleChain) -> SourceResult<()> {
        self.locale = Some(Locale::get_in(styles));
        Ok(())
    }
}

impl Packed<ImageElem> {
    /// Decodes the image.
    pub fn decode(&self, engine: &mut Engine, styles: StyleChain) -> SourceResult<Image> {
        let span = self.span();
        let loaded = &self.source.derived;
        let format = self.determine_format(styles).at(span)?;

        // Construct the image itself.
        let kind = match format {
            ImageFormat::Raster(format) => ImageKind::Raster(
                RasterImage::new(
                    loaded.data.clone(),
                    format,
                    self.icc.get_ref(styles).as_ref().map(|icc| icc.derived.clone()),
                )
                .at(span)?,
            ),
            ImageFormat::Vector(VectorFormat::Svg) => {
                // Warn the user if the image contains a foreign object. Not
                // perfect because the svg could also be encoded, but that's an
                // edge case.
                if memchr::memmem::find(&loaded.data, b"<foreignObject").is_some() {
                    engine.sink.warn(warning!(
                        span,
                        "image contains foreign object";
                        hint: "SVG images with foreign objects might render incorrectly \
                               in Typst";
                        hint: "see https://github.com/typst/typst/issues/1421 for more \
                               information";
                    ));
                }

                // Identify the SVG file in case contained hrefs need to be resolved.
                let svg_file = match &self.source.source {
                    DataSource::Path(path) => {
                        path.resolve_if_some(span.id()).ok().map(|v| v.intern())
                    }
                    DataSource::Bytes(_) => span.id(),
                };
                ImageKind::Svg(
                    SvgImage::with_fonts_images(
                        loaded.data.clone(),
                        engine.world,
                        &families(styles).map(|f| f.as_str()).collect::<Vec<_>>(),
                        svg_file,
                    )
                    .within(loaded)?,
                )
            }
            ImageFormat::Vector(VectorFormat::Pdf) => {
                let document = match PdfDocument::new(loaded.data.clone()) {
                    Ok(doc) => doc,
                    Err(e) => match e {
                        // TODO: the `DecyptionError` is currently not public
                        LoadPdfError::Decryption(_) => {
                            bail!(
                                span,
                                "the PDF is encrypted or password-protected";
                                hint: "such PDFs are currently not supported";
                                hint: "preprocess the PDF to remove the encryption";
                            );
                        }
                        LoadPdfError::Invalid => {
                            bail!(
                                span,
                                "the PDF could not be loaded";
                                hint: "perhaps the PDF file is malformed";
                            );
                        }
                    },
                };

                // See https://github.com/LaurenzV/hayro/issues/141.
                if document.pdf().xref().has_optional_content_groups() {
                    engine.sink.warn(warning!(
                        span,
                        "PDF contains optional content groups";
                        hint: "the image might display incorrectly in PDF export";
                        hint: "preprocess the PDF to flatten or remove optional content \
                               groups";
                    ));
                }

                // The user provides the page number start from 1, but further
                // down the pipeline, page numbers are 0-based.
                let page_num = self.page.get(styles).get();
                let page_idx = page_num - 1;
                let num_pages = document.num_pages();

                let Some(pdf_image) = PdfImage::new(document, page_idx) else {
                    let s = if num_pages == 1 { "" } else { "s" };
                    bail!(
                        span,
                        "page {page_num} does not exist";
                        hint: "the document only has {num_pages} page{s}";
                    );
                };

                ImageKind::Pdf(pdf_image)
            }
        };

        Ok(Image::new(kind, self.alt.get_cloned(styles), self.scaling.get(styles)))
    }

    /// Tries to determine the image format based on the format that was
    /// explicitly defined, or else the extension, or else the data.
    fn determine_format(&self, styles: StyleChain) -> StrResult<ImageFormat> {
        if let Smart::Custom(v) = self.format.get(styles) {
            return Ok(v);
        };

        let Derived { source, derived: loaded } = &self.source;
        if let DataSource::Path(path) = source
            && let Ok(id) = path.resolve_if_some(self.span().id())
            && let Some(format) = determine_format_from_path(id.vpath())
        {
            return Ok(format);
        }

        Ok(ImageFormat::detect(&loaded.data).ok_or("unknown image format")?)
    }
}

/// Derive the image format from the file extension of a path.
fn determine_format_from_path(path: &VirtualPath) -> Option<ImageFormat> {
    match path.extension()? {
        // Raster formats
        "png" => Some(ExchangeFormat::Png.into()),
        "jpg" | "jpeg" => Some(ExchangeFormat::Jpg.into()),
        "gif" => Some(ExchangeFormat::Gif.into()),
        "webp" => Some(ExchangeFormat::Webp.into()),
        // Vector formats
        "svg" | "svgz" => Some(VectorFormat::Svg.into()),
        "pdf" => Some(VectorFormat::Pdf.into()),
        _ => None,
    }
}

impl LocalName for Packed<ImageElem> {
    const KEY: &'static str = "figure";
}

impl Figurable for Packed<ImageElem> {}

/// How an image should adjust itself to a given area,
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Cast)]
pub enum ImageFit {
    /// The image should completely cover the area (preserves aspect ratio by
    /// cropping the image only horizontally or vertically). This is the
    /// default.
    Cover,
    /// The image should be fully contained in the area (preserves aspect
    /// ratio; doesn't crop the image; one dimension can be narrower than
    /// specified).
    Contain,
    /// The image should be stretched so that it exactly fills the area, even if
    /// this means that the image will be distorted (doesn't preserve aspect
    /// ratio and doesn't crop the image).
    Stretch,
}

/// A loaded raster or vector image.
///
/// Values of this type are cheap to clone and hash.
#[derive(Clone, Eq, PartialEq, Hash)]
pub struct Image(Arc<LazyHash<ImageInner>>);

/// The internal representation of an [`Image`].
#[derive(Hash)]
struct ImageInner {
    /// The raw, undecoded image data.
    kind: ImageKind,
    /// A text describing the image.
    alt: Option<EcoString>,
    /// The scaling algorithm to use.
    scaling: Smart<ImageScaling>,
}

impl Image {
    /// When scaling an image to it's natural size, we default to this DPI
    /// if the image doesn't contain DPI metadata.
    pub const DEFAULT_DPI: f64 = 72.0;

    /// Should always be the same as the default DPI used by usvg.
    pub const USVG_DEFAULT_DPI: f64 = 96.0;

    /// Create an image from a `RasterImage` or `SvgImage`.
    pub fn new(
        kind: impl Into<ImageKind>,
        alt: Option<EcoString>,
        scaling: Smart<ImageScaling>,
    ) -> Self {
        Self::new_impl(kind.into(), alt, scaling)
    }

    /// Create an image with optional properties set to the default.
    pub fn plain(kind: impl Into<ImageKind>) -> Self {
        Self::new(kind, None, Smart::Auto)
    }

    /// The internal, non-generic implementation. This is memoized to reuse
    /// the `Arc` and `LazyHash`.
    #[comemo::memoize]
    fn new_impl(
        kind: ImageKind,
        alt: Option<EcoString>,
        scaling: Smart<ImageScaling>,
    ) -> Image {
        Self(Arc::new(LazyHash::new(ImageInner { kind, alt, scaling })))
    }

    /// The format of the image.
    pub fn format(&self) -> ImageFormat {
        match &self.0.kind {
            ImageKind::Raster(raster) => raster.format().into(),
            ImageKind::Svg(_) => VectorFormat::Svg.into(),
            ImageKind::Pdf(_) => VectorFormat::Pdf.into(),
        }
    }

    /// The width of the image in pixels.
    pub fn width(&self) -> f64 {
        match &self.0.kind {
            ImageKind::Raster(raster) => raster.width() as f64,
            ImageKind::Svg(svg) => svg.width(),
            ImageKind::Pdf(pdf) => pdf.width() as f64,
        }
    }

    /// The height of the image in pixels.
    pub fn height(&self) -> f64 {
        match &self.0.kind {
            ImageKind::Raster(raster) => raster.height() as f64,
            ImageKind::Svg(svg) => svg.height(),
            ImageKind::Pdf(pdf) => pdf.height() as f64,
        }
    }

    /// The image's pixel density in pixels per inch, if known.
    pub fn dpi(&self) -> Option<f64> {
        match &self.0.kind {
            ImageKind::Raster(raster) => raster.dpi(),
            ImageKind::Svg(_) => Some(Image::USVG_DEFAULT_DPI),
            ImageKind::Pdf(_) => Some(Image::DEFAULT_DPI),
        }
    }

    /// A text describing the image.
    pub fn alt(&self) -> Option<&str> {
        self.0.alt.as_deref()
    }

    /// The image scaling algorithm to use for this image.
    pub fn scaling(&self) -> Smart<ImageScaling> {
        self.0.scaling
    }

    /// The decoded image.
    pub fn kind(&self) -> &ImageKind {
        &self.0.kind
    }
}

impl Debug for Image {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.debug_struct("Image")
            .field("format", &self.format())
            .field("width", &self.width())
            .field("height", &self.height())
            .field("alt", &self.alt())
            .field("scaling", &self.scaling())
            .finish()
    }
}

/// A kind of image.
#[derive(Clone, Hash)]
pub enum ImageKind {
    /// A raster image.
    Raster(RasterImage),
    /// An SVG image.
    Svg(SvgImage),
    /// A PDF image.
    Pdf(PdfImage),
}

impl From<RasterImage> for ImageKind {
    fn from(image: RasterImage) -> Self {
        Self::Raster(image)
    }
}

impl From<SvgImage> for ImageKind {
    fn from(image: SvgImage) -> Self {
        Self::Svg(image)
    }
}

/// A raster or vector image format.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub enum ImageFormat {
    /// A raster graphics format.
    Raster(RasterFormat),
    /// A vector graphics format.
    Vector(VectorFormat),
}

impl ImageFormat {
    /// Try to detect the format of an image from data.
    pub fn detect(data: &[u8]) -> Option<Self> {
        if let Some(format) = ExchangeFormat::detect(data) {
            return Some(Self::Raster(RasterFormat::Exchange(format)));
        }

        if is_svg(data) {
            return Some(Self::Vector(VectorFormat::Svg));
        }

        if is_pdf(data) {
            return Some(Self::Vector(VectorFormat::Pdf));
        }

        None
    }
}

/// Checks whether the data looks like a PDF file.
fn is_pdf(data: &[u8]) -> bool {
    let head = &data[..data.len().min(2048)];
    memchr::memmem::find(head, b"%PDF-").is_some()
}

/// Checks whether the data looks like an SVG or a compressed SVG.
fn is_svg(data: &[u8]) -> bool {
    // Check for the gzip magic bytes. This check is perhaps a bit too
    // permissive as other formats than SVGZ could use gzip.
    if data.starts_with(&[0x1f, 0x8b]) {
        return true;
    }

    // If the first 2048 bytes contain the SVG namespace declaration, we assume
    // that it's an SVG. Note that, if the SVG does not contain a namespace
    // declaration, usvg will reject it.
    let head = &data[..data.len().min(2048)];
    memchr::memmem::find(head, b"http://www.w3.org/2000/svg").is_some()
}

/// A vector graphics format.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Cast)]
pub enum VectorFormat {
    /// The vector graphics format of the web.
    Svg,
    /// High-fidelity document and graphics format, with focus on exact
    /// reproduction in print.
    Pdf,
}

impl<R> From<R> for ImageFormat
where
    R: Into<RasterFormat>,
{
    fn from(format: R) -> Self {
        Self::Raster(format.into())
    }
}

impl From<VectorFormat> for ImageFormat {
    fn from(format: VectorFormat) -> Self {
        Self::Vector(format)
    }
}

cast! {
    ImageFormat,
    self => match self {
        Self::Raster(v) => v.into_value(),
        Self::Vector(v) => v.into_value(),
    },
    v: RasterFormat => Self::Raster(v),
    v: VectorFormat => Self::Vector(v),
}

/// The image scaling algorithm a viewer should use.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Cast)]
pub enum ImageScaling {
    /// Scale with a smoothing algorithm such as bilinear interpolation.
    Smooth,
    /// Scale with nearest neighbor or a similar algorithm to preserve the
    /// pixelated look of the image.
    Pixelated,
}
