import argparse
import os
import zipfile
import tempfile
import shutil
import numpy as np
from PIL import Image, ImageChops
import xml.etree.ElementTree as ET

NUM_LAYERS = 4
MASKS_PER_LAYER = 3
NOISE_SCALES = [4, 16, 64]


def noise_mask(w, h, scale):
    sw = max(1, w // scale)
    sh = max(1, h // scale)

    noise = np.random.rand(sh, sw)
    img = Image.fromarray((noise * 255).astype(np.uint8), "L")
    img = img.resize((w, h), Image.BILINEAR)

    return img


def build_stack_xml(w, h, groups):
    image = ET.Element(
        "image",
        {"version": "0.0.3", "w": str(w), "h": str(h)}
    )

    root_stack = ET.SubElement(image, "stack")

    for group in groups:

        group_el = ET.SubElement(
            root_stack,
            "stack",
            {"name": group["name"]}
        )

        for layer in group["layers"]:
            layer_attrs = {
                "name": layer["name"],
                "src": layer["src"],
                "opacity": "1.0"
            }
            ET.SubElement(
                group_el,
                "layer",
                layer_attrs
            )

    return ET.tostring(image, encoding="utf-8", xml_declaration=True)


def build_ora(input_png, output_ora):

    base = Image.open(input_png).convert("RGBA")
    w, h = base.size

    tmp = tempfile.mkdtemp()

    try:

        data_dir = os.path.join(tmp, "data")
        thumb_dir = os.path.join(tmp, "Thumbnails")

        os.makedirs(data_dir)
        os.makedirs(thumb_dir)

        groups = []

        for i in range(NUM_LAYERS):

            group_layers = []

            for j in range(MASKS_PER_LAYER):
                mask = noise_mask(w, h, NOISE_SCALES[j])

                # Apply mask to alpha channel of base image
                masked_image = base.copy()
                r, g, b, a = masked_image.split()
                # Multiply existing alpha by mask
                new_alpha = ImageChops.multiply(a, mask)
                masked_image.putalpha(new_alpha)

                layer_path = f"data/layer_{i}_{j}.png"
                masked_image.save(os.path.join(tmp, layer_path))

                group_layers.append({
                    "name": f"layer {j}",
                    "src": layer_path
                })

            groups.append({
                "name": f"group {i}",
                "layers": group_layers
            })

        stack_xml = build_stack_xml(w, h, groups)

        with open(os.path.join(tmp, "stack.xml"), "wb") as f:
            f.write(stack_xml)

        with open(os.path.join(tmp, "mimetype"), "w") as f:
            f.write("image/openraster")

        base.save(os.path.join(tmp, "mergedimage.png"))

        thumb = base.copy()
        thumb.thumbnail((256, 256))
        thumb.save(os.path.join(thumb_dir, "thumbnail.png"))

        with zipfile.ZipFile(output_ora, "w") as z:

            z.write(
                os.path.join(tmp, "mimetype"),
                "mimetype",
                compress_type=zipfile.ZIP_STORED
            )

            for root, _, files in os.walk(tmp):
                for file in files:

                    if file == "mimetype":
                        continue

                    full = os.path.join(root, file)
                    rel = os.path.relpath(full, tmp)

                    z.write(full, rel)

    finally:
        shutil.rmtree(tmp)


def main():

    parser = argparse.ArgumentParser(
        description="Generate ORA with grouped layers and mask layers"
    )

    parser.add_argument("input_png")
    parser.add_argument("output_ora")

    args = parser.parse_args()

    build_ora(args.input_png, args.output_ora)


if __name__ == "__main__":
    main()