# Copyright (c) 2016-2025 Martin Donath <martin.donath@squidfunk.com>

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

from __future__ import annotations

import functools
import json
import logging
import os
import subprocess
import sys

from fnmatch import fnmatch
from colorama import Fore, Style
from concurrent.futures import Future
from concurrent.futures.thread import ThreadPoolExecutor
from hashlib import sha1
from mkdocs import utils
from mkdocs.config.defaults import MkDocsConfig
from mkdocs.exceptions import PluginError
from mkdocs.plugins import BasePlugin
from mkdocs.structure.files import File
from shutil import which
try:
    from PIL import Image
except ImportError:
    pass

from .config import OptimizeConfig

# -----------------------------------------------------------------------------
# Classes
# -----------------------------------------------------------------------------

# Optimize plugin
class OptimizePlugin(BasePlugin[OptimizeConfig]):
    supports_multiple_instances = True

    # Manifest
    manifest: dict[str, str] = {}

    # Initialize plugin
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Initialize incremental builds
        self.is_serve = False

    # Determine whether we're serving the site
    def on_startup(self, *, command, dirty):
        self.is_serve = command == "serve"

        # Initialize thread pool
        self.pool = ThreadPoolExecutor(self.config.concurrency)
        self.pool_jobs: dict[str, Future] = {}

    # Resolve and load manifest
    def on_config(self, config):
        if not self.config.enabled:
            return

        # Resolve cache directory (once) - this is necessary, so the cache is
        # always relative to the configuration file, and thus project, and not
        # relative to the current working directory, or it would not work with
        # the projects plugin.
        path = os.path.abspath(self.config.cache_dir)
        if path != self.config.cache_dir:
            self.config.cache_dir = os.path.join(
                os.path.dirname(config.config_file_path),
                os.path.normpath(self.config.cache_dir)
            )

            # Ensure cache directory exists
            os.makedirs(self.config.cache_dir, exist_ok = True)

        # Initialize manifest
        self.manifest_file = os.path.join(
            self.config.cache_dir, "manifest.json"
        )

        # Load manifest if it exists and the cache should be used
        if os.path.isfile(self.manifest_file) and self.config.cache:
            try:
                with open(self.manifest_file) as f:
                    self.manifest = json.load(f)
            except:
                pass

    # Initialize optimization pipeline
    def on_env(self, env, *, config, files):
        if not self.config.enabled:
            return

        # Skip if media files should not be optimized
        if not self.config.optimize:
            return

        # Filter all optimizable media files and steal reponsibility from MkDocs
        # by removing them from the files collection. Then, start a concurrent
        # job that checks if an image was already optimized and can be returned
        # from the cache, or optimize it accordingly.
        for file in files.media_files():
            if self._is_excluded(file):
                continue

            # Spawn concurrent job to optimize the given image and add future
            # to job dictionary, as it returns the file we need to copy later
            path = os.path.join(self.config.cache_dir, file.src_path)
            self.pool_jobs[file.abs_src_path] = self.pool.submit(
                self._optimize_image, file, path, config
            )

            # Steal responsibility from MkDocs
            files.remove(file)

    # Finish optimization pipeline
    def on_post_build(self, *, config):
        if not self.config.enabled:
            return

        # Skip if media files should not be optimized
        if not self.config.optimize:
            return

        # Reconcile concurrent jobs - we need to wait for all jobs to finish
        # before we can copy the optimized files to the output directory. If an
        # exception occurred in one of the jobs, we raise it here, so the build
        # fails and the author can fix the issue.
        for path, future in self.pool_jobs.items():
            if future.exception():
                raise future.exception()
            else:
                file: File = future.result()
                file.copy_file()

        # Save manifest if cache should be used
        if self.config.cache:
            with open(self.manifest_file, "w") as f:
                f.write(json.dumps(self.manifest, indent = 2, sort_keys = True))

        # Compute and print gains through optimization
        if self.config.print_gain_summary:
            print(Style.NORMAL)
            print(f"  Optimizations:")

            # Print summary for file extension
            for seek in [".png", ".jpg"]:
                size = size_opt = 0
                for path, future in self.pool_jobs.items():
                    file: File = future.result()

                    # Skip files that are not of the given type
                    _, extension = os.path.splitext(path)
                    extension = ".jpg" if extension == ".jpeg" else extension
                    if extension != seek:
                        continue

                    # Compute size before and after optimization
                    size     += os.path.getsize(path)
                    size_opt += os.path.getsize(file.abs_dest_path)

                # Compute absolute and relative gain
                if size and size_opt:
                    gain_abs = size - size_opt
                    gain_rel = (1 - size_opt / size) * 100

                    # Print summary for files
                    print(
                        f"    *{seek} {Fore.GREEN}{_size(size_opt)}"
                        f"{Fore.WHITE}{Style.DIM} ↓ "
                        f"{_size(gain_abs)} [{gain_rel:3.1f}%]"
                        f"{Style.RESET_ALL}"
                    )

            # Reset all styles
            print(Style.RESET_ALL)

    # Save manifest on shutdown
    def on_shutdown(self):
        if not self.config.enabled:
            return

        # Shutdown thread pool - if we're on Python 3.9 and above, cancel all
        # pending futures that have not yet been scheduled
        if sys.version_info >= (3, 9):
            self.pool.shutdown(cancel_futures = True)
        else:
            self.pool.shutdown()

        # Save manifest if cache should be used
        if self.manifest and self.config.cache:
            with open(self.manifest_file, "w") as f:
                f.write(json.dumps(self.manifest, indent = 2, sort_keys = True))

    # -------------------------------------------------------------------------

    # Check if a file can be optimized
    def _is_optimizable(self, file: File):

        # Check if PNG images should be optimized
        if file.url.endswith((".png")):
            return self.config.optimize_png

        # Check if JPG images should be optimized
        if file.url.endswith((".jpg", ".jpeg")):
            return self.config.optimize_jpg

        # File can not be optimized by the plugin
        return False

    # Check if the given file is excluded
    def _is_excluded(self, file: File):
        if not self._is_optimizable(file):
            return True

        # Check if file matches one of the inclusion patterns
        path = file.src_path
        if self.config.optimize_include:
            for pattern in self.config.optimize_include:
                if fnmatch(file.src_uri, pattern):
                    return False

            # File is not included
            log.debug(f"Excluding file '{path}' due to inclusion patterns")
            return True

        # Check if file matches one of the exclusion patterns
        for pattern in self.config.optimize_exclude:
            if fnmatch(file.src_uri, pattern):
                log.debug(f"Excluding file '{path}' due to exclusion patterns")
                return True

        # File is not excluded
        return False

    # Optimize image and write to cache
    def _optimize_image(self, file: File, path: str, config: MkDocsConfig):
        with open(file.abs_src_path, "rb") as f:
            data = f.read()
            hash = sha1(data).hexdigest()

            # Check if file hash changed, so we need to optimize again
            prev = self.manifest.get(file.url, "")
            if hash != prev or not os.path.isfile(path):
                os.makedirs(os.path.dirname(path), exist_ok = True)

                # Optimize PNG image using pngquant
                if file.url.endswith((".png")):
                    self._optimize_image_png(file, path, config)

                # Optimize JPG image using pillow
                if file.url.endswith((".jpg", ".jpeg")):
                    self._optimize_image_jpg(file, path, config)

                # Compute size before and after optimization
                size     = len(data)
                size_opt = os.path.getsize(path)

                # Compute absolute and relative gain
                gain_abs = size - size_opt
                gain_rel = (1 - size_opt / size) * 100

                # Print how much we gained, if we did and desired
                gain = ""
                if gain_abs and self.config.print_gain:
                    gain += " ↓ "
                    gain += " ".join([_size(gain_abs), f"[{gain_rel:3.1f}%]"])

                # Print summary for file
                log.info(
                    f"Optimized media file: {file.src_uri} "
                    f"{Fore.GREEN}{_size(size_opt)}"
                    f"{Fore.WHITE}{Style.DIM}{gain}"
                    f"{Style.RESET_ALL}"
                )

                # Update manifest by associating file with hash
                self.manifest[file.url] = hash

        # Compute project root
        root = os.path.dirname(config.config_file_path)

        # Compute source file system path
        file.abs_src_path = path
        file.src_path     = os.path.relpath(path, root)

        # Return file to be copied from cache
        return file

    # Optimize PNG image - we first tried to use libimagequant, but encountered
    # the occassional segmentation fault, which means it's probably not a good
    # choice. Instead, we just rely on pngquant which seems much more stable.
    def _optimize_image_png(self, file: File, path: str, config: MkDocsConfig):

        # Check if the required dependencies for optimizing are available, which
        # is, at the absolute minimum, the 'pngquant' binary, and raise an error
        # to the caller, so he can decide what to do with the error. The caller
        # can treat this as a warning or an error to abort the build.
        if not which("pngquant"):
            docs = os.path.relpath(config.docs_dir)
            path = os.path.relpath(file.abs_src_path, docs)
            raise PluginError(
                f"Couldn't optimize image '{path}' in '{docs}': 'pngquant' "
                f"not found. Make sure 'pngquant' is installed and in your path"
            )

        # Build command line arguments
        args = ["pngquant",
            "--force", "--skip-if-larger",
            "--output", path,
            "--speed", f"{self.config.optimize_png_speed}"
        ]

        # Add flag to remove optional metadata
        if self.config.optimize_png_strip:
            args.append("--strip")

        # Set input file and run, then check if pngquant actually wrote a file,
        # as we instruct it not to if the size of the optimized file is larger.
        # This can happen if files are already compressed and optimized by
        # the author. In that case, just copy the original file.
        subprocess.run([*args, file.abs_src_path])
        if not os.path.isfile(path):
            utils.copy_file(file.abs_src_path, path)

    # Optimize JPG image
    def _optimize_image_jpg(self, file: File, path: str, config: MkDocsConfig):

        # Check if the required dependencies for optimizing are available, which
        # is, at the absolute minimum, the 'pillow' package, and raise an error
        # to the caller, so he can decide what to do with the error. The caller
        # can treat this as a warning or an error to abort the build.
        if not _supports("Image"):
            docs = os.path.relpath(config.docs_dir)
            path = os.path.relpath(file.abs_src_path, docs)
            raise PluginError(
                f"Couldn't optimize image '{path}' in '{docs}': install "
                f"required dependencies – pip install 'mkdocs-material[imaging]'"
            )

        # Open and save optimized image
        image = Image.open(file.abs_src_path)
        image.save(path, "jpeg",
            quality     = self.config.optimize_jpg_quality,
            progressive = self.config.optimize_jpg_progressive
        )

# -----------------------------------------------------------------------------
# Helper functions
# -----------------------------------------------------------------------------

# Check for presence of optional imports
@functools.lru_cache(maxsize = None)
def _supports(name: str):
    return name in globals()

# -----------------------------------------------------------------------------

# Print human-readable size
def _size(value):
    for unit in ["B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB"]:
        if abs(value) < 1000.0:
            return f"{value:3.1f} {unit}"
        value /= 1000.0

# -----------------------------------------------------------------------------
# Data
# -----------------------------------------------------------------------------

# Set up logging
log = logging.getLogger("mkdocs.material.optimize")