#!/usr/bin/env python3
"""Pack a folder into a PrimoEngine .livewallpaper file (package format v1).

Folder layout:
    my-wallpaper/
      manifest.json
      content/...          (shader.metal | web/index.html + assets | video.mp4)
      thumbnail.png        (optional: shown in the app's grid; not part of the checksum)

What it does:
  1. Runs the same basic checks the app runs on import (so you see errors here, not in the app).
  2. Computes the v1 `checksum` over content/ and writes it into the packaged manifest.
  3. Zips manifest.json + content/ (+ thumbnail.png) into <out>.livewallpaper, skipping .DS_Store.

Usage:  python3 pack.py <folder> [out.livewallpaper]
Needs only Python 3 (preinstalled on macOS via the Command Line Tools).
"""
import hashlib, json, os, re, sys, zipfile

MAX_ENTRY = 256 * 1024 * 1024      # 256 MB per file (uncompressed)
MAX_ENTRIES = 4096                 # files per package
SKIP = {".DS_Store"}

# Mirrors the app's Metal safety gate (textual, case-sensitive, comments ignored).
BANNED = [
    ("kernel ", "compute kernels are not allowed (fragment shaders only)"),
    ("[[kernel", "compute kernels are not allowed"),
    ("device ", "writable `device` buffers are not allowed"),
    ("threadgroup ", "threadgroup memory is not allowed"),
    ("atomic_", "atomic operations are not allowed"),
    ('#include "', "local #include is not allowed"),
]


def list_content(content_dir):
    rels = []
    for root, _dirs, files in os.walk(content_dir):
        for name in files:
            if name in SKIP:
                continue
            rels.append(os.path.relpath(os.path.join(root, name), content_dir).replace(os.sep, "/"))
    return sorted(rels, key=lambda p: p.encode("utf-8"))


def checksum(content_dir, rels):
    h = hashlib.sha256()
    for rel in rels:
        h.update(rel.encode("utf-8"))
        h.update(b"\x00")
        with open(os.path.join(content_dir, rel), "rb") as f:
            h.update(f.read())
    return "sha256-" + h.hexdigest()


def strip_comments(src):
    src = re.sub(r"/\*.*?\*/", "", src, flags=re.S)
    return re.sub(r"//[^\n]*", "", src)


def lint(src_dir, m, rels):
    errors, warnings = [], []
    if m.get("schemaVersion") != 1:
        errors.append("schemaVersion must be 1")
    for k in ("id", "version", "title", "type", "entry"):
        if not m.get(k):
            errors.append(f"'{k}' is missing or empty")
    typ, entry = m.get("type"), m.get("entry", "")
    if typ not in ("metal", "web", "video"):
        errors.append("type must be \"metal\", \"web\" or \"video\"")
    if entry and not entry.startswith("content/"):
        errors.append("entry must be a path inside content/ (e.g. content/shader.metal)")
    entry_path = os.path.join(src_dir, entry)
    if entry and not os.path.isfile(entry_path):
        errors.append(f"entry file '{entry}' does not exist")
    ext = os.path.splitext(entry)[1].lower()
    if typ == "metal" and ext != ".metal":
        errors.append("metal entry must be a .metal file")
    if typ == "web" and ext not in (".html", ".htm"):
        errors.append("web entry must be an .html file")
    if typ == "video" and ext not in (".mp4", ".mov", ".m4v"):
        errors.append("video entry must be .mp4, .mov or .m4v")
    if typ == "metal" and os.path.isfile(entry_path):
        code = strip_comments(open(entry_path, encoding="utf-8").read())
        for token, why in BANNED:
            if token in code:
                errors.append(f"Shader rejected: {why} (found '{token.strip()}')")
        if "f_main" not in code:
            errors.append("Shader is missing the required 'f_main' fragment function")
        if "v_main" not in code:
            errors.append("Shader is missing the required 'v_main' vertex function (it would render black)")
    for c in m.get("config") or []:
        if c.get("type") == "enum":
            warnings.append(f"config '{c.get('key')}': type 'enum' is not supported yet and will be hidden")
        elif c.get("type") not in ("float", "int", "bool", "color"):
            warnings.append(f"config '{c.get('key')}': unknown type '{c.get('type')}' will be ignored")
    if typ == "metal":
        extra = [c.get("key") for c in m.get("config") or [] if c.get("key") not in ("speed", "tint")]
        if extra:
            warnings.append(f"metal shaders only receive 'speed' and 'tint'; {extra} will show in Settings but not reach the shader")
    if len(rels) + 1 > MAX_ENTRIES:
        errors.append(f"too many files ({len(rels)}); limit is {MAX_ENTRIES}")
    for rel in rels:
        if os.path.getsize(os.path.join(src_dir, "content", rel)) > MAX_ENTRY:
            errors.append(f"content/{rel} is larger than 256 MB")
        if ".." in rel:
            errors.append(f"unsafe path content/{rel}")
    return errors, warnings


def main():
    if len(sys.argv) not in (2, 3):
        print(__doc__)
        return 2
    src = sys.argv[1].rstrip("/")
    out = sys.argv[2] if len(sys.argv) == 3 else os.path.basename(src) + ".livewallpaper"
    content_dir = os.path.join(src, "content")
    if not os.path.isdir(content_dir):
        print(f"error: missing folder {content_dir}", file=sys.stderr)
        return 1
    try:
        manifest = json.load(open(os.path.join(src, "manifest.json"), encoding="utf-8"))
    except (OSError, ValueError) as e:
        print(f"error: manifest.json unreadable or invalid JSON: {e}", file=sys.stderr)
        return 1

    rels = list_content(content_dir)
    errors, warnings = lint(src, manifest, rels)
    for w in warnings:
        print(f"warning: {w}")
    if errors:
        for e in errors:
            print(f"error: {e}", file=sys.stderr)
        return 1

    manifest["checksum"] = checksum(content_dir, rels)
    with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as z:
        z.writestr("manifest.json", json.dumps(manifest, indent=2, ensure_ascii=False))
        for rel in rels:
            z.write(os.path.join(content_dir, rel), "content/" + rel)
        thumb = os.path.join(src, "thumbnail.png")
        if os.path.isfile(thumb):
            z.write(thumb, "thumbnail.png")
    print(f"ok: wrote {out}  ({len(rels)} content file(s), {manifest['checksum'][:19]}…)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())