#!/usr/bin/env python3
"""Export the October UI PSD without changing existing application assets.

Requires: pip install psd-tools Pillow aggdraw
"""
import argparse
import hashlib
import json
from pathlib import Path

from PIL import Image, ImageDraw
from psd_tools import PSDImage

ROOT = Path(__file__).resolve().parents[1]
DESIGN = ROOT / "design/ui-20261009"
ASSETS = ROOT / "h5/assets/img/ui-20261009"

# PSD layer indices are recorded alongside descriptive filenames for traceability.
ASSET_SPECS = [
    ("bg-splash.jpg", "0/0", "0"),
    ("bg-profile.jpg", "3/0/0", "3"),
    ("bg-login.jpg", "login/0/0", "login"),
    ("bg-report.jpg", "8/0", "8"),
    ("bg-history.jpg", "2/0", "2"),
    ("brand-logo.png", "0/1/0", None),
    ("brand-login.png", "login/7/1", None),
    ("icon-close.png", "1/6", None),
    ("icon-dropdown.png", "3/1/2", None),
    ("icon-back.png", "3/6/0", None),
    ("icon-album.png", "5/1/1", None),
    ("icon-switch-camera.png", "5/1/0", None),
    ("icon-wechat.png", "login/7/6", None),
    ("icon-guest.png", "login/7/7", None),
    ("icon-shutter.png", "5/1/2", None),
    ("skin-0-dry.png", "3/3/1/0/2", None),
    ("skin-1-mixeddry.png", "3/3/1/1/2", None),
    ("skin-2-neutral.png", "3/3/1/2/2", None),
    ("skin-3-oily.png", "3/3/1/3/2", None),
    ("skin-4-mixedoily.png", "3/3/1/4/2", None),
    ("product-25.png", "8/1/0/1", None),
    ("product-35.png", "8/1/1/2", None),
    ("product-50.png", "8/1/2/4", None),
    ("product-badge.png", "8/1/0/3", None),
    ("score-ring-base.png", "8/12/5", None),
    ("score-ring-value.png", "8/12/6", None),
    ("sample-contour.png", "8/9/8", None),
    ("sample-lines.png", "8/9/10", None),
    ("sample-tone.png", "8/10/10", None),
    ("sample-redness.png", "8/11/8", None),
    ("camera-mask-front.png", "5/0", "5"),
    ("camera-mask-right.png", "6/0", "6"),
    ("camera-mask-left.png", "7/0", "7"),
    ("camera-guide-front.png", ["5/8", "5/9", "5/10"], "5"),
    ("camera-guide-right.png", ["6/8", "6/9", "6/10"], "6"),
    ("camera-guide-left.png", ["7/8", "7/9", "7/10"], "7"),
]


def layer_at(psd, indices):
    layer = psd
    for part in indices.split("/"):
        if part == "login":
            matches = [page for page in psd if page.name == "登录"]
            if len(matches) != 1:
                raise ValueError("Expected exactly one 登录 page; refusing incorrect asset mapping")
            layer = matches[0]
        else:
            layer = layer[int(part)]
    return layer


def bounds(layer):
    # Hidden groups can report an empty bbox; union their children explicitly.
    if layer.is_group():
        boxes = [bounds(child) for child in layer]
        boxes = [b for b in boxes if b[2] > b[0] and b[3] > b[1]]
        if boxes:
            return (min(b[0] for b in boxes), min(b[1] for b in boxes),
                    max(b[2] for b in boxes), max(b[3] for b in boxes))
    return tuple(layer.bbox)


def save_image(image, target):
    target.parent.mkdir(parents=True, exist_ok=True)
    if target.suffix == ".jpg":
        background = Image.new("RGB", image.size, "white")
        background.paste(image, mask=image.getchannel("A") if image.mode == "RGBA" else None)
        background.save(target, quality=92, optimize=True, progressive=True)
    else:
        image.save(target, optimize=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("psd", type=Path)
    parser.add_argument("--assets-only", action="store_true", help="Reuse existing page exports")
    args = parser.parse_args()
    psd = PSDImage.open(args.psd)
    DESIGN.mkdir(parents=True, exist_ok=True)
    manifest = {"source": args.psd.name, "document_size": psd.size,
                "pages": [], "assets": [], "layers": []}

    def inventory(group, prefix=""):
        for index, layer in enumerate(group):
            key = f"{prefix}/{index}".strip("/")
            entry = {"path": key, "name": layer.name, "kind": layer.kind,
                     "bbox": list(layer.bbox), "visible": layer.visible}
            if layer.kind == "type":
                entry["text"] = layer.text
            manifest["layers"].append(entry)
            if layer.is_group():
                inventory(layer, key)
    inventory(psd)

    previews = []
    for index, page in enumerate(psd):
        print(f"Page {index}: {page.name}", flush=True)
        # Cached layer pixels preserve Photoshop text appearance without local fonts.
        target = DESIGN / "pages" / f"{index + 1:02d}-{page.name}.png"
        if args.assets_only and target.exists():
            image = Image.open(target)
        else:
            image = page.composite(viewport=page.bbox)
            save_image(image, target)
        preview = image.copy()
        preview.thumbnail((216, 1553))
        previews.append((page.name, preview))
        manifest["pages"].append({"file": str(target.relative_to(ROOT)),
                                  "name": page.name, "size": image.size,
                                  "bbox": list(page.bbox)})

    thumbs = []
    for filename, indices, crop_page in ASSET_SPECS:
        source_indices = indices if isinstance(indices, list) else [indices]
        layers = [layer_at(psd, key) for key in source_indices]
        layer = layers[0]
        viewport = layer_at(psd, crop_page).bbox if crop_page else bounds(layer)
        print(f"Asset: {filename}", flush=True)
        # A Photoshop clipping layer composites to empty if its base is omitted.
        # Disable only this root layer's clipping while exporting it independently.
        image = Image.new("RGBA", (viewport[2]-viewport[0], viewport[3]-viewport[1]))
        for source in layers:
            was_clipping = source.clipping
            if was_clipping:
                source.clipping = False
            part = source.composite(viewport=viewport, layer_filter=lambda _: True)
            if was_clipping:
                source.clipping = True
            if part is not None:
                image = Image.alpha_composite(image, part.convert("RGBA"))
        if image is None:
            raise RuntimeError(f"No image for {indices}: {layer.name}")
        if image.mode == "RGBA" and image.getchannel("A").getextrema()[1] == 0:
            raise RuntimeError(f"Empty transparent export for {indices}: {filename}")
        target = ASSETS / filename
        save_image(image, target)
        manifest["assets"].append({"file": str(target.relative_to(ROOT)),
                                   "source_layer": indices, "layer_name": layer.name,
                                   "size": image.size, "bbox": list(viewport),
                                   "bytes": target.stat().st_size,
                                   "sha256": hashlib.sha256(target.read_bytes()).hexdigest(),
                                   "original_visible": layer.is_visible(),
                                   "usage": "reference-only" if filename.startswith("sample-")
                                            else "frontend-asset"})
        thumb = image.copy()
        thumb.thumbnail((180, 190))
        thumbs.append((filename, thumb))

    sheet = Image.new("RGB", (5 * 236, ((len(previews)+4)//5) * 1620), "#eeeeee")
    draw = ImageDraw.Draw(sheet)
    for i, (_, preview) in enumerate(previews):
        x, y = (i % 5) * 236 + 10, (i // 5) * 1620 + 35
        draw.text((x, y - 25), f"Page {i+1:02d}", fill="black")
        sheet.paste(preview, (x, y), preview if preview.mode == "RGBA" else None)
    save_image(sheet, DESIGN / "pages-overview.jpg")

    sheet = Image.new("RGB", (6 * 205, ((len(thumbs)+5)//6)*230), "#e3e3e3")
    draw = ImageDraw.Draw(sheet)
    for i, (name, thumb) in enumerate(thumbs):
        x, y = (i % 6)*205 + 10, (i // 6)*230 + 30
        draw.text((x, y - 22), name, fill="black")
        sheet.paste(thumb, (x, y), thumb if thumb.mode == "RGBA" else None)
    save_image(sheet, DESIGN / "assets-overview.jpg")
    (DESIGN / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2)+"\n")
    print(f"Exported {len(previews)} pages and {len(thumbs)} assets", flush=True)


if __name__ == "__main__":
    main()