"""Derive the Web Blender RE 5000PBR asset library from the desktop catalog.

Usage:
  build_blender_native/bin/blender --background --factory-startup \
    --python tools/web/generate-re-pbr-library.py -- \
    "/path/to/5000+ PBR catalog.blend" OUTPUT_DIRECTORY \
    [--base-color-size 1024] [--normal-size 1024] [--scalar-size 512] [--preview-size 256]

The source .blend is never modified. The derived library keeps the original
shader graphs (FRAME/REROUTE/DISPLACEMENT/MAPPING/VALUE included) so the Web
compiler/reader must support them; this script only marks assets, downsamples
packed textures, repacks them and renders material previews.

Outputs (inside OUTPUT_DIRECTORY):
  re-pbr-library.blend    six marked MATERIAL assets, no scene objects
  blender_assets.cats.txt  Materials/RE5000PBR catalog
  previews/<sha256>.png    256x256 EEVEE material previews
  preview-manifest.json
  build-report.json
"""

import hashlib
import importlib.util
import json
import pathlib
import sys

import bpy

CATALOG_ID = "a8000000-0000-4000-8000-000000000001"
CATALOG_PATH = "Materials/RE5000PBR"
CATALOG_NAME = "RE5000PBR"
LIBRARY_NAME = "re-pbr-library.blend"
EXPECTED_MATERIALS = 6


def sha256_file(path: pathlib.Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def load_common():
    path = pathlib.Path(__file__).with_name("generate-pbr-library.py")
    spec = importlib.util.spec_from_file_location("pbr_library_common", path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"cannot load {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def parse_arguments():
    arguments = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    if len(arguments) < 2:
        raise SystemExit(
            "usage: blender -b --python generate-re-pbr-library.py -- "
            "SOURCE_BLEND OUTPUT_DIR [--base-color-size N] [--normal-size N] "
            "[--scalar-size N] [--preview-size N]"
        )
    source = pathlib.Path(arguments[0]).resolve()
    output = pathlib.Path(arguments[1]).resolve()
    options = {
        "base_color_size": 1024,
        "normal_size": 1024,
        "scalar_size": 512,
        "preview_size": 256,
        "preserve_original": False,
        "material_suffix": "",
        "library_id": "library:re-pbr-v1",
    }
    index = 2
    while index < len(arguments):
        flag = arguments[index]
        if flag == "--preserve-original":
            options["preserve_original"] = True
            index += 1
            continue
        if flag not in {"--base-color-size", "--normal-size", "--scalar-size", "--preview-size", "--texture-size", "--material-suffix", "--library-id"}:
            raise SystemExit(f"unknown option {flag}")
        if index + 1 >= len(arguments):
            raise SystemExit(f"missing value for {flag}")
        key = flag.removeprefix("--").replace("-", "_")
        value = arguments[index + 1]
        if key == "texture_size":
            size = int(value)
            options["base_color_size"] = size
            options["normal_size"] = size
            options["scalar_size"] = size
        elif key in {"base_color_size", "normal_size", "scalar_size", "preview_size"}:
            options[key] = int(value)
        else:
            options[key] = value
        index += 2
    if not source.is_file():
        raise SystemExit(f"source blend not found: {source}")
    return source, output, options


def target_size(image_name: str, options) -> int:
    lowered = image_name.lower()
    if "normal" in lowered:
        return options["normal_size"]
    if "base_color" in lowered or "basecolor" in lowered or "albedo" in lowered:
        return options["base_color_size"]
    return options["scalar_size"]


def resample_image(image, size: int, temporary: pathlib.Path, preserve_original: bool = False):
    if image.size[0] == 0 or image.size[1] == 0:
        return
    if preserve_original and image.size[0] == size and image.size[1] == size:
        return
    if image.size[0] != size or image.size[1] != size:
        image.scale(size, size)
    image.file_format = "JPEG"
    temporary.parent.mkdir(parents=True, exist_ok=True)
    try:
        image.save(filepath=str(temporary), quality=84)
    except TypeError:
        image.filepath_raw = str(temporary)
        image.save()
    if not temporary.is_file() or temporary.stat().st_size == 0:
        raise RuntimeError(f"failed to write resampled image {temporary}")
    data = temporary.read_bytes()
    # Image.pack() with no data re-packs the original packed file. Pack the
    # resampled JPEG bytes explicitly so the library actually shrinks.
    image.pack(data=data, data_len=len(data))


def material_uses_principled(material) -> bool:
    if not material.use_nodes or material.node_tree is None:
        return False
    return any(node.type == "BSDF_PRINCIPLED" for node in material.node_tree.nodes)


def collect_materials():
    materials = [
        material for material in bpy.data.materials
        if material.name != "Dots Stroke" and material_uses_principled(material)
    ]
    materials.sort(key=lambda material: material.name)
    if len(materials) != EXPECTED_MATERIALS:
        raise RuntimeError(
            f"expected {EXPECTED_MATERIALS} Principled materials, found {len(materials)}: "
            + ", ".join(material.name for material in materials)
        )
    return materials


def write_catalog_definition(path: pathlib.Path):
    lines = [
        "# This is an Asset Catalog Definition file for Blender.",
        "# Derived from the desktop RE 5000PBR catalog by generate-re-pbr-library.py",
        "",
        "VERSION 1",
        "",
        f"{CATALOG_ID}:{CATALOG_PATH}:{CATALOG_NAME}",
    ]
    path.write_text("\n".join(lines) + "\n", encoding="utf-8", newline="\n")


def mark_materials(materials):
    for material in materials:
        if material.asset_data is None:
            material.asset_mark()
        material.asset_data.catalog_id = CATALOG_ID
        material.asset_data.description = "RE 5000PBR material for the Web Blender render lesson"
        material.asset_data.tags.new("PBR")
        material.asset_data.tags.new("RE5000PBR")


def resample_material_images(materials, options, temporary_dir: pathlib.Path):
    processed = set()
    resampled = []
    for material in materials:
        if material.node_tree is None:
            continue
        for node in material.node_tree.nodes:
            if node.type != "TEX_IMAGE" or node.image is None:
                continue
            image = node.image
            if image.name in processed or not image.packed_file:
                continue
            processed.add(image.name)
            size = target_size(image.name, options)
            temporary = temporary_dir / f"{sha256_bytes(image.name.encode('utf-8'))[:16]}.jpg"
            before = image.size[:]
            colorspace = image.colorspace_settings.name
            resample_image(image, size, temporary, options.get("preserve_original", False))
            image.colorspace_settings.name = colorspace
            resampled.append({
                "name": image.name,
                "from": list(before),
                "to": list(image.size[:]),
                "colorspace": colorspace,
                "packedBytes": image.packed_file.size if image.packed_file else 0,
            })
    return resampled


def remove_scene_objects():
    for obj in list(bpy.data.objects):
        bpy.data.objects.remove(obj, do_unlink=True)


def remove_unused_data(materials):
    keep_material = {material.name for material in materials}
    for material in list(bpy.data.materials):
        if material.name not in keep_material:
            bpy.data.materials.remove(material)
    used_images = set()
    for material in materials:
        if material.node_tree is None:
            continue
        for node in material.node_tree.nodes:
            if node.type == "TEX_IMAGE" and node.image is not None:
                used_images.add(node.image.name)
    for image in list(bpy.data.images):
        if image.name not in used_images and not image.name.startswith("Render Result"):
            bpy.data.images.remove(image)


def render_previews(common, materials, preview_dir: pathlib.Path, preview_size: int):
    common.PREVIEW_SIZE = preview_size
    scene = bpy.context.scene
    camera = common.configure_preview_scene(scene)
    scene.render.resolution_x = preview_size
    scene.render.resolution_y = preview_size
    entries = []
    for material in materials:
        bpy.ops.mesh.primitive_uv_sphere_add(segments=96, ring_count=64, radius=1.0, location=(0.0, 0.0, 0.0))
        sphere = bpy.context.object
        sphere.name = f"Preview {material.name}"
        bpy.ops.object.shade_smooth()
        sphere.data.materials.append(material)
        common.frame_object(scene, camera, sphere)
        temporary = preview_dir / "tmp.png"
        common.render_preview(scene, temporary)
        entry = common.store_preview(temporary, preview_dir)
        bpy.data.objects.remove(sphere, do_unlink=True)
        entry.update({"idType": "MATERIAL", "name": material.name})
        entries.append(entry)
    return entries


def main():
    source, output, options = parse_arguments()
    output.mkdir(parents=True, exist_ok=True)
    preview_dir = output / "previews"
    preview_dir.mkdir(parents=True, exist_ok=True)
    for stale in preview_dir.glob("*.png"):
        stale.unlink()

    common = load_common()
    bpy.ops.wm.open_mainfile(filepath=str(source))
    materials = collect_materials()
    suffix = str(options.get("material_suffix", ""))
    if suffix:
        for material in materials:
            material.name = f"{material.name}{suffix}"
    mark_materials(materials)
    temporary_dir = output / ".resample"
    resampled = resample_material_images(materials, options, temporary_dir)
    remove_scene_objects()
    remove_unused_data(materials)

    catalog_path = output / "blender_assets.cats.txt"
    write_catalog_definition(catalog_path)
    library_path = output / LIBRARY_NAME
    bpy.ops.wm.save_as_mainfile(filepath=str(library_path), compress=True)

    preview_entries = render_previews(common, materials, preview_dir, options["preview_size"])
    preview_manifest = output / "preview-manifest.json"
    preview_manifest.write_text(
        json.dumps({"schemaVersion": 1, "previews": preview_entries}, indent=2, ensure_ascii=True) + "\n",
        encoding="utf-8",
        newline="\n",
    )

    for item in temporary_dir.glob("*"):
        if item.is_file():
            item.unlink()
    if temporary_dir.is_dir():
        temporary_dir.rmdir()

    source_sha = sha256_file(source)
    library_sha = sha256_file(library_path)
    preview_bytes = sum(entry["byteLength"] for entry in preview_entries)
    report = {
        "schemaVersion": 1,
        "generator": "tools/web/generate-re-pbr-library.py",
        "scriptSha256": sha256_file(pathlib.Path(__file__).resolve()),
        "source": {"path": str(source), "sha256": source_sha, "byteLength": source.stat().st_size},
        "library": {
            "path": str(library_path),
            "sha256": library_sha,
            "byteLength": library_path.stat().st_size,
            "libraryId": "library:re-pbr-v1",
            "name": LIBRARY_NAME,
        },
        "catalog": {"id": CATALOG_ID, "path": CATALOG_PATH, "name": CATALOG_NAME},
        "materials": [material.name for material in materials],
        "textures": resampled,
        "textureOptions": options,
        "previewCount": len(preview_entries),
        "previewBytes": preview_bytes,
        "blenderVersion": ".".join(str(value) for value in bpy.app.version),
    }
    (output / "build-report.json").write_text(
        json.dumps(report, indent=2, ensure_ascii=True) + "\n", encoding="utf-8", newline="\n"
    )
    print(json.dumps({
        "materials": len(materials),
        "previews": len(preview_entries),
        "librarySha256": library_sha,
        "libraryBytes": library_path.stat().st_size,
        "textureBytes": sum(item["packedBytes"] for item in resampled),
    }, indent=2))


if __name__ == "__main__":
    main()