#!/usr/bin/env python3
"""Prepares the game files of the Endless Sky web build (fredygames.com).

    prepare-assets.py <endless-sky source dir> <redrawn font dir> <output dir>

Works from the original release tree and from the pruned tarball hosted on fredygames.com
(endless-sky-0.11.3-pruned.tar.gz, the same tree without the image files listed in
removed-images.txt). In both cases the left-out images are worked out from the `copyright` file:
in the original tree the result must be exactly the list in removed-images.txt; in the pruned tree
none of those files may be present and no other image may fall into a left-out category. Any
other state of the tree stops the script.

Copies what the game reads at run time (data/, images/, sounds/, shaders/, credits.txt, keys.txt)
from the release tree into <output dir>, with these changes:

  - Images whose licence is not free for this site, or doubtful, are left out ("when in doubt, leave
    it out"). Each image's licence is the last stanza of the release's `copyright` file that
    matches it (Debian copyright format: later stanzas override earlier ones). Left out:
      * "unsplash": stanzas whose licence is "Unsplash" (photos uploaded to unsplash.com after
        June 2017; the Unsplash License is not a free licence);
      * "morguefile": images whose stanza says they were "taken from morguefile.com" (the catch-all
        stanzas for images/land/* and images/scene/*); morgueFile's own licence is not a public
        domain dedication, and the file does not name the photographers;
      * "third-party-assets": screenshots or renders of purchased third-party assets (the sivael
        landscapes, Unity Asset Store scenes).
    The data lines that point to a removed image (planet landscapes, conversation scenes, start
    thumbnails, logbook scenes) are removed too; any other kind of reference stops the script. The
    expected count of each category is fixed below, so a different source tree stops the script.
  - images/font/ubuntu14r.png and ubuntu18r.png (pictures of the Ubuntu font) are replaced by the
    DejaVu Sans versions drawn by make-font.py.
  - Every other image is converted to WebP to keep the download small (the web build of the game
    reads WebP; see the patch): lossy at quality 80 with a lossless alpha channel, except the
    interface pictures drawn as PNG (images/ui, icon, label) and the swizzle masks (@sw), which are
    lossless.
    Images that the game blends as premultiplied alpha ("=" before the frame number) keep the colour
    under transparent pixels ("exact"). The pictures stay under their licences (CC BY-SA, CC BY, CC0,
    public domain; see `copyright`), converted.
  - sounds/ambient/ (MP3 music, which the web build cannot play) is left out.

Needs Pillow 12.3.0 with libwebp 1.6.0 (artifacts/tools/venv). The output is the same for the same
input and the same Pillow; the build script checks a SHA-256 of the whole package.
"""
import fnmatch
import io
import os
import re
import shutil
import sys
from multiprocessing import Pool

import PIL
from PIL import Image, features

LOSSLESS_DIRS = ('ui/', 'icon/', 'label/')
QUALITY = 80
METHOD = 5


# Expected number of image files left out per category (Endless Sky 0.11.3).
EXPECTED = {'unsplash': 37, 'morguefile': 90, 'third-party-assets': 10}
# Expected number of data lines removed because they point to a left-out image.
EXPECTED_LINES = {'landscape': 155, 'scene': 5, 'thumbnail': 3}
# Keywords of the data lines that may point to a removed image; the whole line goes. A sprite name
# alone on a line is a weighted child of "landscape".
REFERENCE_KEYS = {'landscape', 'scene', 'thumbnail'}


def copyright_stanzas(path):
    """Yields (file patterns, licence, stanza text) for every Files stanza of a Debian copyright file."""
    with open(path, encoding='utf-8') as f:
        text = f.read()
    for para in re.split(r'\n\s*\n', text):
        lines = para.split('\n')
        files, licence, in_files = [], None, False
        for line in lines:
            if line.startswith('Files:'):
                in_files = True
                files += line[len('Files:'):].split()
            elif line.startswith((' ', '\t')) and in_files:
                files += line.split()
            else:
                in_files = False
                if line.startswith('License:'):
                    licence = line[len('License:'):].strip()
        if files and licence:
            yield files, licence, para


def category(licence, text):
    if licence == 'Unsplash':
        return 'unsplash'
    if re.search(r'morgue-?file', text, re.I):
        return 'morguefile'
    if re.search(r'asset store|assets purchased|purchased assets', text, re.I):
        return 'third-party-assets'
    return None


def matches(rel, pattern):
    # Debian copyright globs: "*" and "?" match anything, "/" included. A few stanzas name a file
    # without its extension ("images/land/station20"), so the name without extension counts too.
    return fnmatch.fnmatchcase(rel, pattern) or fnmatch.fnmatchcase(os.path.splitext(rel)[0], pattern)


def classify(src):
    """Returns {image path relative to src: category} for the images present in src that are left out."""
    stanzas = list(copyright_stanzas(os.path.join(src, 'copyright')))
    out = {}
    for root, dirs, names in os.walk(os.path.join(src, 'images')):
        dirs.sort()
        for name in sorted(names):
            rel = os.path.relpath(os.path.join(root, name), src).replace(os.sep, '/')
            last = None
            for files, licence, text in stanzas:
                if any(matches(rel, p) for p in files):
                    last = (licence, text)
            kind = category(*last) if last else None
            if kind:
                out[rel] = kind
    return out


def removed_list():
    """The checked-in list of left-out image files (removed-images.txt, next to this script)."""
    out = {}
    with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'removed-images.txt'),
              encoding='utf-8') as f:
        for line in f:
            if line.strip() and not line.startswith('#'):
                kind, rel = line.rstrip('\n').split('\t')
                out[rel] = kind
    counts = {k: sum(1 for v in out.values() if v == k) for k in EXPECTED}
    if counts != EXPECTED or len(out) != sum(EXPECTED.values()):
        sys.exit(f'removed-images.txt has counts {counts}, expected {EXPECTED}')
    return out


def excluded_images(src):
    """Returns ({left-out image path: category}, tree kind), checking the tree against the list."""
    listed = removed_list()
    found = classify(src)
    present = [rel for rel in listed if os.path.exists(os.path.join(src, rel))]
    if len(present) == len(listed):
        # Original release tree: the copyright file must give exactly the checked-in list.
        if found != listed:
            extra = sorted(set(found.items()) ^ set(listed.items()))[:10]
            sys.exit(f'copyright classification differs from removed-images.txt: {extra}')
        return listed, 'original'
    if not present:
        # Pruned tree: the listed files are gone, and nothing else may fall into a left-out category.
        if found:
            sys.exit(f'pruned tree still has left-out images: {sorted(found)[:10]}')
        return listed, 'pruned'
    sys.exit(f'{len(present)} of {len(listed)} left-out images present; expected all or none: {present[:5]}')


def sprite_name(rel):
    """The name the game gives an image (source/image/ImageFileData.cpp), without "images/"."""
    name = os.path.splitext(rel[len('images/'):])[0]
    for suffix in ('@2x', '@1x', '@sw'):
        if name.endswith(suffix):
            name = name[:-len(suffix)]
    m = re.search(r'[-+~=^](\d*)$', name)
    if m and m.start() > 0:
        name = name[:m.start()]
    return name


TOKEN = re.compile(r'"([^"]*)"|`([^`]*)`|(#.*)|(\S+)')


def tokens(line):
    out = []
    for m in TOKEN.finditer(line):
        if m.group(3) is not None:
            break
        out.append(next(g for g in (m.group(1), m.group(2), m.group(4)) if g is not None))
    return out


def indent(line):
    return len(line) - len(line.lstrip('\t '))


def convert(job):
    src, dst, rel = job
    name = os.path.basename(rel)
    stem = os.path.splitext(name)[0]
    im = Image.open(src)
    im.load()
    if im.mode not in ('RGB', 'RGBA'):
        im = im.convert('RGBA' if ('A' in im.getbands() or 'transparency' in im.info) else 'RGB')
    if im.mode == 'RGBA' and im.getchannel('A').getextrema() == (255, 255):
        im = im.convert('RGB')
    lossless = (rel.startswith(LOSSLESS_DIRS) and name.lower().endswith('.png')) or stem.endswith('@sw')
    exact = bool(re.search(r'=\d*$', stem))
    buf = io.BytesIO()
    if lossless:
        im.save(buf, 'WEBP', lossless=True, quality=100, method=METHOD, exact=True)
    else:
        im.save(buf, 'WEBP', quality=QUALITY, alpha_quality=100, method=METHOD, exact=exact)
    data = buf.getvalue()
    if data[12:16] not in (b'VP8 ', b'VP8L', b'VP8X') or b'ANIM' in data[:64]:
        raise RuntimeError(f'{rel}: unexpected WebP layout')
    with open(dst, 'wb') as f:
        f.write(data)
    return len(data)


def main(src, fontdir, out):
    if PIL.__version__ != '12.3.0' or features.version('webp') != '1.6.0':
        sys.exit(f'Needs Pillow 12.3.0 / libwebp 1.6.0, found {PIL.__version__} / {features.version("webp")}')
    if os.path.exists(out):
        shutil.rmtree(out)
    os.makedirs(out)

    removed, tree = excluded_images(src)
    print(f'source tree: {tree}', file=sys.stderr)
    removed_sprites = {sprite_name(r) for r in removed}
    # Every image of a removed sprite must be removed (no frame or variant left behind).
    for root, _, names in os.walk(os.path.join(src, 'images')):
        for name in names:
            rel = os.path.relpath(os.path.join(root, name), src).replace(os.sep, '/')
            if rel not in removed and sprite_name(rel) in removed_sprites:
                sys.exit(f'{rel} belongs to a removed sprite but is not removed')

    # data/: the game's text files, minus the lines that point to removed images.
    dropped = {}
    for root, dirs, names in os.walk(os.path.join(src, 'data')):
        dirs.sort()
        for name in sorted(names):
            path = os.path.join(root, name)
            rel_path = os.path.relpath(path, src)
            dst = os.path.join(out, rel_path)
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            with open(path, 'rb') as f:
                lines = f.read().split(b'\n')
            kept = []
            for i, line in enumerate(lines):
                text = line.decode('utf-8', 'replace')
                toks = tokens(text)
                hits = [t for t in toks if t in removed_sprites]
                if not hits:
                    kept.append(line)
                    continue
                # Allowed forms: "<key> <sprite>", "add|remove <key> <sprite>", "log scene <sprite>",
                # and a sprite alone (weighted child of "landscape"), with an optional weight.
                key = [t for t in toks if t not in ('add', 'remove', 'log')][0]
                if not (key in REFERENCE_KEYS or (key in removed_sprites and len(toks) <= 2)):
                    sys.exit(f'{rel_path}:{i + 1}: unexpected reference to a removed image: {text.strip()}')
                if len([t for t in toks if t in removed_sprites or t in REFERENCE_KEYS
                        or t in ('add', 'remove', 'log') or re.fullmatch(r'[\d.]+', t)]) != len(toks):
                    sys.exit(f'{rel_path}:{i + 1}: line has more than the removed image: {text.strip()}')
                following = next((l for l in lines[i + 1:] if l.strip()), b'')
                if following and indent(following.decode('utf-8', 'replace')) > indent(text):
                    sys.exit(f'{rel_path}:{i + 1}: removed line has children: {text.strip()}')
                kind = key if key in REFERENCE_KEYS else 'landscape (weighted)'
                dropped[kind] = dropped.get(kind, 0) + 1
            with open(dst, 'wb') as f:
                f.write(b'\n'.join(kept))
    for kind in sorted(set(removed.values())):
        print(f'removed {sum(1 for v in removed.values() if v == kind)} images: {kind}', file=sys.stderr)
    print(f'removed data lines: {dropped}', file=sys.stderr)
    if dropped != EXPECTED_LINES:
        sys.exit(f'unexpected removed data lines {dropped}, expected {EXPECTED_LINES}')

    for name in ('credits.txt', 'keys.txt'):
        shutil.copyfile(os.path.join(src, name), os.path.join(out, name))
    shutil.copytree(os.path.join(src, 'shaders'), os.path.join(out, 'shaders'))

    # sounds/: the WAV effects as they are; no ambient music.
    for root, dirs, names in os.walk(os.path.join(src, 'sounds')):
        dirs.sort()
        rel_root = os.path.relpath(root, src).replace(os.sep, '/')
        if rel_root == 'sounds/ambient' or rel_root.startswith('sounds/ambient/'):
            continue
        for name in sorted(names):
            if not name.endswith('.wav'):
                sys.exit(f'unexpected sound {rel_root}/{name}')
            dst = os.path.join(out, rel_root, name)
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            shutil.copyfile(os.path.join(root, name), dst)

    # images/: WebP, except the redrawn fonts.
    jobs, targets = [], set()
    for root, dirs, names in os.walk(os.path.join(src, 'images')):
        dirs.sort()
        for name in sorted(names):
            path = os.path.join(root, name)
            rel = os.path.relpath(path, src).replace(os.sep, '/')
            sub = rel[len('images/'):]
            if rel in removed or sub.startswith('font/'):
                continue
            target = os.path.join(out, 'images', os.path.splitext(sub)[0] + '.webp')
            if target in targets:
                sys.exit(f'two images would become {target}')
            targets.add(target)
            os.makedirs(os.path.dirname(target), exist_ok=True)
            jobs.append((path, target, sub))
    os.makedirs(os.path.join(out, 'images', 'font'))
    for name in ('ubuntu14r.png', 'ubuntu18r.png'):
        shutil.copyfile(os.path.join(fontdir, name), os.path.join(out, 'images', 'font', name))

    with Pool(os.cpu_count()) as pool:
        total = sum(pool.imap(convert, jobs, chunksize=8))
    print(f'{len(jobs)} images, {total / 1e6:.1f} MB of WebP', file=sys.stderr)


if __name__ == '__main__':
    if len(sys.argv) != 4:
        sys.exit(__doc__)
    main(*sys.argv[1:])