import argparse
import re
import sys
from pathlib import Path

from PIL import Image, ImageOps

import jsonfile
from jsonfile import PROJECT_ROOT

SOURCE_DIR = Path.home() / "dotfiles" / "wallpapers"
OUTPUT_DIR = PROJECT_ROOT / "static" / "wallpapers" / "gallery"
THUMB_DIR = OUTPUT_DIR / "thumbs"
PUBLIC_PATH = "/wallpapers/gallery"
MANIFEST_FILE = "wallpapers.json"
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
FULL_SIZE = (2560, 1440)
THUMB_SIZE = (480, 270)
FULL_QUALITY = 82
THUMB_QUALITY = 70
SMALL_IMAGE_EDGE = 1000


def slug_of(path: Path) -> str:
    return re.sub(r"[^a-z0-9]+", "-", path.stem.lower()).strip("-")


def display_name(path: Path) -> str:
    words = re.sub(r"[_\-]+", " ", path.stem).strip()
    return re.sub(r"^\d+\s+", "", words).title()


def source_images(source_dir: Path) -> list[Path]:
    return sorted(path for path in source_dir.iterdir() if path.suffix.lower() in IMAGE_EXTENSIONS and path.is_file())


def flatten(image: Image.Image) -> Image.Image:
    image = ImageOps.exif_transpose(image)
    if image.mode in ("RGBA", "LA", "P"):
        image = image.convert("RGBA")
        background = Image.new("RGB", image.size, (0, 0, 0))
        background.paste(image, mask=image.getchannel("A"))
        return background
    return image.convert("RGB")


def resample_for(image: Image.Image) -> int:
    is_pixel_art = max(image.size) <= SMALL_IMAGE_EDGE
    return Image.Resampling.NEAREST if is_pixel_art else Image.Resampling.LANCZOS


def fit_within(image: Image.Image, size: tuple[int, int]) -> Image.Image:
    scale = min(size[0] / image.width, size[1] / image.height)
    if scale >= 1 and max(image.size) > SMALL_IMAGE_EDGE:
        return image
    target = (max(1, round(image.width * scale)), max(1, round(image.height * scale)))
    return image.resize(target, resample_for(image))


def save_jpeg(image: Image.Image, path: Path, quality: int):
    path.parent.mkdir(parents=True, exist_ok=True)
    image.save(path, "JPEG", quality=quality, optimize=True, progressive=True)


def is_up_to_date(source: Path, outputs: list[Path]) -> bool:
    return all(output.exists() and output.stat().st_mtime >= source.stat().st_mtime for output in outputs)


def convert(source: Path, force: bool) -> dict:
    slug = slug_of(source)
    full_path = OUTPUT_DIR / f"{slug}.jpg"
    thumb_path = THUMB_DIR / f"{slug}.jpg"

    if force or not is_up_to_date(source, [full_path, thumb_path]):
        with Image.open(source) as opened:
            image = flatten(opened)
        save_jpeg(fit_within(image, FULL_SIZE), full_path, FULL_QUALITY)
        save_jpeg(ImageOps.fit(image, THUMB_SIZE, resample_for(image)), thumb_path, THUMB_QUALITY)

    return {
        "id": slug,
        "name": display_name(source),
        "image": f"{PUBLIC_PATH}/{slug}.jpg",
        "thumb": f"{PUBLIC_PATH}/thumbs/{slug}.jpg",
    }


def remove_stale(kept_slugs: set[str]) -> int:
    removed = 0
    for folder in (OUTPUT_DIR, THUMB_DIR):
        for path in folder.glob("*.jpg"):
            if path.stem not in kept_slugs:
                path.unlink()
                removed += 1
    return removed


def folder_size(folder: Path) -> int:
    return sum(path.stat().st_size for path in folder.rglob("*") if path.is_file())


def main() -> int:
    parser = argparse.ArgumentParser(description="Compress wallpapers from the dotfiles repo into JPEGs for the website")
    parser.add_argument("source", nargs="?", default=str(SOURCE_DIR), help="folder with the original wallpapers")
    parser.add_argument("--force", action="store_true", help="convert every wallpaper again")
    arguments = parser.parse_args()

    source_dir = Path(arguments.source).expanduser()
    if not source_dir.is_dir():
        print(f"{source_dir} is not a folder", file=sys.stderr)
        return 1

    wallpapers = []
    for source in source_images(source_dir):
        try:
            wallpapers.append(convert(source, arguments.force))
            print(f"ok      {source.name}")
        except OSError as error:
            print(f"skipped {source.name}: {error}", file=sys.stderr)

    removed = remove_stale({wallpaper["id"] for wallpaper in wallpapers})
    jsonfile.save(MANIFEST_FILE, wallpapers)

    original_size = sum(path.stat().st_size for path in source_images(source_dir))
    compressed_size = folder_size(OUTPUT_DIR)
    print(f"\n{len(wallpapers)} wallpapers, {removed} stale files removed")
    print(f"{original_size / 1024 / 1024:.1f} MB -> {compressed_size / 1024 / 1024:.1f} MB (including thumbnails)")
    return 0


if __name__ == "__main__":
    sys.exit(main())