#!/usr/bin/env python3
import os
import struct
import zipfile
import xml.etree.ElementTree as ET

# Offsets for BevelEmboss style keys inside sultan.kra's Unnamed/annotations/layerstyles.asl:
# lagl (local angle): 747
# Lald (local altitude): 768
# srgR (depth): 790
# blur (size): 819
# bvlD (direction enum value, e.g. b"In  " or b"Out "): 851
# Sftn (soften): 1094
# hglO (highlight opacity): 512
# sdwO (shadow opacity): 652
# hglC (highlight R/G/B doubles): 438, 458, 478
# sdwC (shadow R/G/B doubles): 590, 610, 630
# hglM (highlight blend mode classID): 402
# sdwM (shadow blend mode classID): 544

BASE_KRA = "/mnt/extra/00_PROJECTS/hcie-rust-v4/_tmp/kra_tunes_simple/black_rect_simple.kra"
OUT_DIR = "/mnt/extra/00_PROJECTS/hcie-rust-v4/_tmp/kra_tunes_simple_black/bevel_emboss_kra_set"

def get_double_bytes(val):
    return struct.pack(">d", float(val))

def main():
    bases = [
        ("/mnt/extra/00_PROJECTS/hcie-rust-v4/_tmp/kra_tunes_simple/black_rect_simple.kra", "/mnt/extra/00_PROJECTS/hcie-rust-v4/_tmp/kra_tunes_simple_black/bevel_emboss_kra_set"),
        ("/mnt/extra/00_PROJECTS/hcie-rust-v4/_tmp/kra_tunes_simple/red_rect_simple.kra", "/mnt/extra/00_PROJECTS/hcie-rust-v4/_tmp/kra_tunes_simple_red/bevel_emboss_kra_set")
    ]

    for base_kra, out_dir in bases:
        if not os.path.exists(base_kra):
            print(f"Error: Base Krita file {base_kra} not found.")
            continue

        os.makedirs(out_dir, exist_ok=True)

        # Read base archive and cache entries in memory
        base_zip = zipfile.ZipFile(base_kra, 'r')
        files_cache = {}
        for name in base_zip.namelist():
            files_cache[name] = base_zip.read(name)
        
        asl_path_in_zip = "Unnamed/annotations/layerstyles.asl"
        if asl_path_in_zip in files_cache:
            asl_data = bytearray(files_cache[asl_path_in_zip])
        else:
            fallback_paths = [
                "/mnt/extra/00_PROJECTS/hcie-rust-v4/_images/kra_style_test/sultan.kra",
                "/mnt/extra/00_PROJECTS/hcie-rust-v4/_images/_psd_stil_test/sultan_test/sultan.kra",
                "_images/kra_style_test/sultan.kra",
                "_images/_psd_stil_test/sultan_test/sultan.kra"
            ]
            sultan_zip = None
            for path in fallback_paths:
                if os.path.exists(path):
                    sultan_zip = zipfile.ZipFile(path, 'r')
                    break
            if sultan_zip is None:
                print("Error: Could not find sultan.kra fallback to load layerstyles.asl")
                return
            asl_data = bytearray(sultan_zip.read(asl_path_in_zip))
            sultan_zip.close()

        xml_data = files_cache["maindoc.xml"]

        # Set visibility in XML: both target layer and Background should be visible
        # Also ensure the layerstyle attribute is attached to the target layer
        ET.register_namespace('', "http://www.calligra.org/DTD/krita")
        root = ET.fromstring(xml_data)
        for layer in root.findall('.//{http://www.calligra.org/DTD/krita}layer'):
            if layer.attrib.get('name') in ['Layer 2', 'Red Rectangle']:
                layer.set('visible', '1')
                layer.set('layerstyle', '{4bbe8f2c-7ad2-4e31-8af0-9ad6715cacf4}')
            elif layer.attrib.get('name') == 'Background':
                layer.set('visible', '1')
            else:
                layer.set('visible', '0')
        xml_patched = b'<?xml version="1.0" encoding="UTF-8"?>\n<!DOCTYPE DOC PUBLIC \'-//KDE//DTD krita 2.0//EN\' \'http://www.calligra.org/DTD/krita-2.0.dtd\'>\n' + ET.tostring(root, encoding='utf-8')

        # Parameter lists
        sizes = [5.0, 10.0, 20.0, 30.0, 50.0]
        softens = [0.0, 4.0, 8.0, 16.0]
        depths = [25.0, 50.0, 75.0, 100.0]
        angles = [90.0, 120.0]
        directions = [
            (b"In  ", "up"),
            (b"Out ", "down"),
        ]

        total = len(sizes) * len(softens) * len(depths) * len(angles) * len(directions)
        print(f"Generating {total} styled KRA files for Bevel & Emboss from {os.path.basename(base_kra)}...")

        def patch_untf(asl, key, value):
            idx = asl.find(key)
            if idx != -1 and asl[idx+4:idx+8] == b'UntF':
                asl[idx+12:idx+20] = get_double_bytes(value)

        def patch_enum(asl, key, value):
            idx = asl.find(key)
            if idx != -1 and asl[idx+4:idx+8] == b'enum':
                asl[idx+20:idx+24] = value

        count = 0
        for size in sizes:
            for soften in softens:
                for depth in depths:
                    for angle in angles:
                        for dir_bytes, dir_name in directions:
                            count += 1
                            filename = f"be_s{int(size)}_sf{int(soften)}_d{int(depth)}_a{int(angle)}_{dir_name}.kra"
                            kra_path = os.path.join(out_dir, filename)

                            # Clone and patch ASL
                            patched_asl = bytearray(asl_data)
                            
                            patch_untf(patched_asl, b'blur', size)
                            patch_untf(patched_asl, b'Sftn', soften)
                            patch_untf(patched_asl, b'srgR', depth)
                            patch_untf(patched_asl, b'lagl', angle)
                            patch_enum(patched_asl, b'bvlD', dir_bytes)

                            # Write new ZIP archive
                            with zipfile.ZipFile(kra_path, 'w') as out_zip:
                                has_asl = False
                                for item in base_zip.infolist():
                                    if item.filename == "Unnamed/annotations/layerstyles.asl":
                                        out_zip.writestr(item, patched_asl)
                                        has_asl = True
                                    elif item.filename == "maindoc.xml":
                                        out_zip.writestr(item, xml_patched)
                                    elif item.filename == "Unnamed/layers/layer6.defaultpixel":
                                        out_zip.writestr(item, b"\x80\x80\x80\xff") # 50% gray background
                                    else:
                                        out_zip.writestr(item, files_cache[item.filename])
                                if not has_asl:
                                    out_zip.writestr("Unnamed/annotations/layerstyles.asl", patched_asl)

                            if count % 10 == 0:
                                print(f"  [{count}/{total}] generated...")

        base_zip.close()
        print(f"Successfully generated {count} files under {out_dir}.")

if __name__ == "__main__":
    main()