"""
Prepare voxel textures from FreeStylized.com ZIPs.
Outputs per material:
  - *_albedo.png  : RGBA (RGB=albedo, A=heightmap)
  - *_normal.png  : RGB normal map (OpenGL convention, Y-up)
"""
import io
import os
import zipfile
from PIL import Image, ImageEnhance

# (zip_name, color_pattern, height_pattern, normal_pattern, brightness_factor)
# brightness_factor: <1 = darken, >1 = brighten, 1.0 = unchanged
MATERIALS = [
    ("grass_01_1k",         "color",     "height", "normal_gl", 1.0),
    ("ground_02_1k",        "color",     "height", "normal_gl", 0.75),  # dirt: darkened
    ("ground_stones_01_1k", "baseColor", "height", "normal_gl", 1.0),
    ("sand_01_1k",          "color",     "height", "normal_gl", 1.0),
    ("snow_01_1k",          "color",     "height", "normal_gl", 1.0),
    ("rock_01_1k",          "color",     "height", "normal_gl", 1.0),
]

OUTPUT_NAMES = [
    "grass",
    "dirt",
    "stone",
    "sand",
    "snow",
    "smoothstone",
]

TARGET_SIZE = 512
RAW_DIR = os.path.join(os.path.dirname(__file__), "..", "assets", "raw")
OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "assets", "voxel")


def find_file_in_zip(zf, pattern):
    """Find a file in the zip matching a pattern substring."""
    for name in zf.namelist():
        basename = os.path.basename(name).lower()
        if pattern.lower() in basename and basename.endswith(".png"):
            return name
    return None


def load_image_from_zip(zf, filename, mode="RGB"):
    data = zf.read(filename)
    img = Image.open(io.BytesIO(data))
    # Handle 16-bit heightmaps: Pillow's .convert("L") on I;16 images
    # doesn't scale properly. We must manually scale 0-65535 → 0-255.
    if img.mode in ("I;16", "I") and mode == "L":
        # Convert to 32-bit int first, then scale down
        img = img.convert("I")
        img = img.point(lambda v: v / 256)
        return img.convert("L")
    return img.convert(mode)


def process_material(zip_path, color_pat, height_pat, normal_pat, brightness, out_name):
    with zipfile.ZipFile(zip_path, "r") as zf:
        color_file = find_file_in_zip(zf, color_pat)
        height_file = find_file_in_zip(zf, height_pat)
        normal_file = find_file_in_zip(zf, normal_pat)

        if not color_file:
            print(f"  ERROR: no color file matching '{color_pat}' in {zip_path}")
            return False

        # ── Albedo + Heightmap → RGBA ──
        color_img = load_image_from_zip(zf, color_file, "RGB")

        if brightness != 1.0:
            color_img = ImageEnhance.Brightness(color_img).enhance(brightness)

        if height_file:
            height_img = load_image_from_zip(zf, height_file, "L")
        else:
            print(f"  WARNING: no height map, deriving from luminance")
            height_img = color_img.convert("L")

        color_img = color_img.resize((TARGET_SIZE, TARGET_SIZE), Image.LANCZOS)
        height_img = height_img.resize((TARGET_SIZE, TARGET_SIZE), Image.LANCZOS)

        r, g, b = color_img.split()
        rgba = Image.merge("RGBA", (r, g, b, height_img))

        albedo_path = os.path.join(OUT_DIR, f"{out_name}_albedo.png")
        rgba.save(albedo_path, "PNG")
        print(f"  OK: {out_name}_albedo.png ({TARGET_SIZE}x{TARGET_SIZE})")

        # ── Normal map → RGB ──
        if normal_file:
            normal_img = load_image_from_zip(zf, normal_file, "RGB")
            normal_img = normal_img.resize((TARGET_SIZE, TARGET_SIZE), Image.LANCZOS)
            normal_path = os.path.join(OUT_DIR, f"{out_name}_normal.png")
            normal_img.save(normal_path, "PNG")
            print(f"  OK: {out_name}_normal.png ({TARGET_SIZE}x{TARGET_SIZE})")
        else:
            print(f"  WARNING: no normal map found")

        return True


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    print(f"Output directory: {os.path.abspath(OUT_DIR)}")
    print()

    success = 0
    for i, (zip_name, color_pat, height_pat, normal_pat, brightness) in enumerate(MATERIALS):
        zip_path = os.path.join(RAW_DIR, zip_name + ".zip")
        print(f"[{i+1}/6] {OUTPUT_NAMES[i]} <- {zip_name}.zip")

        if not os.path.exists(zip_path):
            print(f"  ERROR: {zip_path} not found")
            continue

        if process_material(zip_path, color_pat, height_pat, normal_pat, brightness, OUTPUT_NAMES[i]):
            success += 1

    print(f"\nDone: {success}/6 materials generated in {os.path.abspath(OUT_DIR)}")


if __name__ == "__main__":
    main()