#!/usr/bin/env python3
"""
PSD Layer Effects Test File Generator v2
=========================================
Uses ImageMagick to create correct PSD base, then patches in lfx2 effects.
"""

import struct, os, sys, subprocess
from PIL import Image

OUTPUT_DIR = "/home/hc/Pictures/_psd_stil_test"
W, H = 512, 512
PREFIX = "fx_"
TMP_PNG = '/tmp/_fx_shape.png'
TMP_PSD = '/tmp/_fx_base.psd'


# ============================================================================
# Descriptor helpers
# ============================================================================

def dk(k):
    return struct.pack('>I', 0) + k[:4].ljust(4, b'\x00')

def db(k, v):
    return dk(k) + b'bool' + (b'\x01' if v else b'\x00')

def de(k, t, v):
    return dk(k) + b'enum' + struct.pack('>I', 0) + t[:4].ljust(4, b'\x00') + struct.pack('>I', 0) + v[:4].ljust(4, b'\x00')

def duf(k, u, v):
    return dk(k) + b'UntF' + u[:4].ljust(4, b'\x00') + struct.pack('>d', v)

def dd(k, v):
    return dk(k) + b'doub' + struct.pack('>d', v)

def dl(k, v):
    return dk(k) + b'long' + struct.pack('>I', v)

def dobject(k, cid, n, fields):
    body = struct.pack('>I', 0) + cid + struct.pack('>I', n) + fields
    if k:
        return dk(k) + b'Objc' + body
    return b'Objc' + body

def rgb_color(k, r, g, b):
    return dobject(k, b'RGBC', 3, dd(b'Rd  ', float(r)) + dd(b'Grn ', float(g)) + dd(b'Bl  ', float(b)))


# ============================================================================
# Effect lfx2 builders
# ============================================================================

def lfx2_drop_shadow():
    f = db(b'enab', True)
    f += de(b'Md  ', b'BlnM', b'Nrml')
    f += rgb_color(b'Clr ', 0, 0, 0)
    f += duf(b'Opct', b'#Prc', 75)
    f += duf(b'lagl', b'#Ang', 135)
    f += duf(b'Dstn', b'#Pxl', 15)
    f += duf(b'Ckmt', b'#Pxl', 0)
    f += duf(b'blur', b'#Pxl', 10)
    f += db(b'Nose', False)
    fx = dk(b'DrSh') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 9) + f
    return _wrap_lfx2(1, fx)

def lfx2_inner_shadow():
    f = db(b'enab', True)
    f += de(b'Md  ', b'BlnM', b'Nrml')
    f += rgb_color(b'Clr ', 0, 0, 0)
    f += duf(b'Opct', b'#Prc', 75)
    f += duf(b'lagl', b'#Ang', 135)
    f += duf(b'Dstn', b'#Pxl', 10)
    f += duf(b'Ckmt', b'#Pxl', 0)
    f += duf(b'blur', b'#Pxl', 8)
    f += db(b'Nose', False)
    fx = dk(b'IrSh') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 9) + f
    return _wrap_lfx2(1, fx)

def lfx2_outer_glow():
    f = db(b'enab', True)
    f += de(b'Md  ', b'BlnM', b'Scrn')
    f += rgb_color(b'Clr ', 255, 200, 0)
    f += duf(b'Opct', b'#Prc', 75)
    f += duf(b'Ckmt', b'#Pxl', 0)
    f += duf(b'blur', b'#Pxl', 20)
    f += db(b'Nose', False)
    fx = dk(b'OrGl') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 7) + f
    return _wrap_lfx2(1, fx)

def lfx2_inner_glow():
    f = db(b'enab', True)
    f += de(b'Md  ', b'BlnM', b'Scrn')
    f += rgb_color(b'Clr ', 255, 100, 0)
    f += duf(b'Opct', b'#Prc', 80)
    f += duf(b'Ckmt', b'#Pxl', 0)
    f += duf(b'blur', b'#Pxl', 15)
    f += db(b'Nose', False)
    f += de(b'glwS', b'BETE', b'Intr')
    fx = dk(b'IrGl') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 8) + f
    return _wrap_lfx2(1, fx)

def lfx2_bevel_inner():
    f = db(b'enab', True)
    f += de(b'bvlS', b'BESl', b'InrB')
    f += de(b'bvlT', b'BETE', b'SfBL')
    f += duf(b'srgR', b'#Prc', 200)
    f += de(b'bvlD', b'BESs', b'In  ')
    f += duf(b'Sftn', b'#Pxl', 0)
    f += duf(b'blur', b'#Pxl', 10)
    f += dl(b'Hlag', 1)
    f += duf(b'lagl', b'#Ang', 120)
    f += duf(b'Lald', b'#Ang', 30)
    f += de(b'hglM', b'BlnM', b'Scrn')
    f += rgb_color(b'hglC', 255, 255, 255)
    f += duf(b'hglO', b'#Prc', 75)
    f += de(b'sdwM', b'BlnM', b'Mltp')
    f += rgb_color(b'sdwC', 0, 0, 0)
    f += duf(b'sdwO', b'#Prc', 75)
    fx = dk(b'ebbl') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 18) + f
    return _wrap_lfx2(1, fx)

def lfx2_bevel_outer():
    f = db(b'enab', True)
    f += de(b'bvlS', b'BESl', b'OtrB')
    f += de(b'bvlT', b'BETE', b'SfBL')
    f += duf(b'srgR', b'#Prc', 200)
    f += de(b'bvlD', b'BESs', b'In  ')
    f += duf(b'Sftn', b'#Pxl', 0)
    f += duf(b'blur', b'#Pxl', 10)
    f += dl(b'Hlag', 1)
    f += duf(b'lagl', b'#Ang', 120)
    f += duf(b'Lald', b'#Ang', 30)
    f += de(b'hglM', b'BlnM', b'Scrn')
    f += rgb_color(b'hglC', 255, 255, 255)
    f += duf(b'hglO', b'#Prc', 75)
    f += de(b'sdwM', b'BlnM', b'Mltp')
    f += rgb_color(b'sdwC', 0, 0, 0)
    f += duf(b'sdwO', b'#Prc', 75)
    fx = dk(b'ebbl') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 18) + f
    return _wrap_lfx2(1, fx)

def lfx2_emboss():
    f = db(b'enab', True)
    f += de(b'bvlS', b'BESl', b'Embs')
    f += de(b'bvlT', b'BETE', b'SfBL')
    f += duf(b'srgR', b'#Prc', 200)
    f += de(b'bvlD', b'BESs', b'In  ')
    f += duf(b'Sftn', b'#Pxl', 0)
    f += duf(b'blur', b'#Pxl', 8)
    f += dl(b'Hlag', 1)
    f += duf(b'lagl', b'#Ang', 120)
    f += duf(b'Lald', b'#Ang', 30)
    f += de(b'hglM', b'BlnM', b'Scrn')
    f += rgb_color(b'hglC', 255, 255, 255)
    f += duf(b'hglO', b'#Prc', 75)
    f += de(b'sdwM', b'BlnM', b'Mltp')
    f += rgb_color(b'sdwC', 0, 0, 0)
    f += duf(b'sdwO', b'#Prc', 75)
    fx = dk(b'ebbl') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 18) + f
    return _wrap_lfx2(1, fx)

def lfx2_pillow_emboss():
    f = db(b'enab', True)
    f += de(b'bvlS', b'BESl', b'PlEb')
    f += de(b'bvlT', b'BETE', b'SfBL')
    f += duf(b'srgR', b'#Prc', 200)
    f += de(b'bvlD', b'BESs', b'In  ')
    f += duf(b'Sftn', b'#Pxl', 0)
    f += duf(b'blur', b'#Pxl', 8)
    f += dl(b'Hlag', 1)
    f += duf(b'lagl', b'#Ang', 120)
    f += duf(b'Lald', b'#Ang', 30)
    f += de(b'hglM', b'BlnM', b'Scrn')
    f += rgb_color(b'hglC', 255, 255, 255)
    f += duf(b'hglO', b'#Prc', 75)
    f += de(b'sdwM', b'BlnM', b'Mltp')
    f += rgb_color(b'sdwC', 0, 0, 0)
    f += duf(b'sdwO', b'#Prc', 75)
    fx = dk(b'ebbl') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 18) + f
    return _wrap_lfx2(1, fx)

def lfx2_satin():
    f = db(b'enab', True)
    f += de(b'Md  ', b'BlnM', b'Mltp')
    f += rgb_color(b'Clr ', 0, 0, 0)
    f += duf(b'Opct', b'#Prc', 50)
    f += duf(b'lagl', b'#Ang', 19)
    f += duf(b'Dstn', b'#Pxl', 11)
    f += duf(b'blur', b'#Pxl', 20)
    f += db(b'MpgS', True)
    f += db(b'Invr', False)
    fx = dk(b'lagl') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 9) + f
    return _wrap_lfx2(1, fx)

def lfx2_color_overlay():
    f = db(b'enab', True)
    f += de(b'Md  ', b'BlnM', b'Nrml')
    f += rgb_color(b'Clr ', 220, 50, 50)
    f += duf(b'Opct', b'#Prc', 80)
    fx = dk(b'SoFi') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 4) + f
    return _wrap_lfx2(1, fx)

def lfx2_gradient_overlay():
    grad_fields = dl(b'Intr', 4096)
    grad_fields += de(b'GrdF', b'GrdF', b'Lnr ')
    grad_fields += dk(b'Clrs') + b'VlLs' + struct.pack('>I', 2)
    grad_fields += b'Objc' + struct.pack('>I', 0) + b'Clrt' + struct.pack('>I', 4) + rgb_color(b'Clr ', 0, 0, 0) + dl(b'In  ', 0) + dl(b'Lctn ', 0) + dl(b'Mdpn ', 50)
    grad_fields += b'Objc' + struct.pack('>I', 0) + b'Clrt' + struct.pack('>I', 4) + rgb_color(b'Clr ', 255, 255, 255) + dl(b'In  ', 0) + dl(b'Lctn ', 4096) + dl(b'Mdpn ', 50)
    grad_fields += dk(b'Trns') + b'VlLs' + struct.pack('>I', 2)
    grad_fields += b'Objc' + struct.pack('>I', 0) + b'Trns' + struct.pack('>I', 2) + dd(b'Opct', 100.0) + dl(b'Lctn ', 0) + dl(b'Mdpn ', 50)
    grad_fields += b'Objc' + struct.pack('>I', 0) + b'Trns' + struct.pack('>I', 2) + dd(b'Opct', 100.0) + dl(b'Lctn ', 4096) + dl(b'Mdpn ', 50)
    grad = dobject(b'Grad', b'Grdn', 7, grad_fields)

    f = db(b'enab', True)
    f += de(b'Md  ', b'BlnM', b'Nrml')
    f += grad
    f += duf(b'Opct', b'#Prc', 100)
    f += dl(b'Angl', 90)
    f += de(b'Type ', b'GrdT', b'Lnr ')
    f += db(b'Algn', True)
    f += duf(b'Scl ', b'#Prc', 100)
    fx = dk(b'GrFl') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 8) + f
    return _wrap_lfx2(1, fx)

def lfx2_stroke_inside():
    f = db(b'enab', True)
    f += de(b'Stly', b'FStl', b'InsF')
    f += de(b'PntT', b'FrFl', b'SClr')
    f += de(b'Md  ', b'BlnM', b'Nrml')
    f += duf(b'Opct', b'#Prc', 100)
    f += duf(b'Sz  ', b'#Pxl', 6)
    f += rgb_color(b'Clr ', 255, 0, 0)
    fx = dk(b'FrFX') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 7) + f
    return _wrap_lfx2(1, fx)

def lfx2_stroke_center():
    f = db(b'enab', True)
    f += de(b'Stly', b'FStl', b'CtrF')
    f += de(b'PntT', b'FrFl', b'SClr')
    f += de(b'Md  ', b'BlnM', b'Nrml')
    f += duf(b'Opct', b'#Prc', 100)
    f += duf(b'Sz  ', b'#Pxl', 6)
    f += rgb_color(b'Clr ', 0, 255, 0)
    fx = dk(b'FrFX') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 7) + f
    return _wrap_lfx2(1, fx)

def lfx2_stroke_outside():
    f = db(b'enab', True)
    f += de(b'Stly', b'FStl', b'OutF')
    f += de(b'PntT', b'FrFl', b'SClr')
    f += de(b'Md  ', b'BlnM', b'Nrml')
    f += duf(b'Opct', b'#Prc', 100)
    f += duf(b'Sz  ', b'#Pxl', 6)
    f += rgb_color(b'Clr ', 0, 0, 255)
    fx = dk(b'FrFX') + b'Objc' + struct.pack('>I', 0) + b'null' + struct.pack('>I', 7) + f
    return _wrap_lfx2(1, fx)


def _wrap_lfx2(count, *effects):
    # Unicode name: [4B char_count=1][2B U+0000 null char]
    name = struct.pack('>I', 1) + struct.pack('>H', 0)
    # ClassID: [4B length=0][4B key='null']
    classid = struct.pack('>I', 0) + b'null'
    desc = name + classid + struct.pack('>I', count)
    for fx in effects:
        desc += fx
    return struct.pack('>I', 0) + struct.pack('>I', 16) + desc


# ============================================================================
# Shape generators
# ============================================================================

def create_l_shape():
    img = Image.new('RGBA', (W, H), (0, 0, 0, 0))
    for y in range(H):
        for x in range(W):
            if (100 <= x < 380) and (180 <= y < 230):
                img.putpixel((x, y), (70, 130, 180, 255))
            elif (140 <= x < 190) and (180 <= y < 400):
                img.putpixel((x, y), (70, 130, 180, 255))
    return img

def create_diamond():
    import math
    img = Image.new('RGBA', (W, H), (0, 0, 0, 0))
    cx, cy, rx, ry = 256, 250, 160, 120
    for y in range(H):
        for x in range(W):
            if abs(x - cx) / rx + abs(y - cy) / ry <= 1.0:
                img.putpixel((x, y), (180, 80, 60, 255))
    return img

def create_triangle():
    img = Image.new('RGBA', (W, H), (0, 0, 0, 0))
    x0, y0, x1, y1, x2, y2 = 200, 140, 380, 320, 150, 350
    def sign(px, py, ax, ay, bx, by):
        return (px - bx) * (ay - by) - (ax - bx) * (py - by)
    for y in range(H):
        for x in range(W):
            d1 = sign(x, y, x0, y0, x1, y1)
            d2 = sign(x, y, x1, y1, x2, y2)
            d3 = sign(x, y, x2, y2, x0, y0)
            has_neg = (d1 < 0) or (d2 < 0) or (d3 < 0)
            has_pos = (d1 > 0) or (d2 > 0) or (d3 > 0)
            if not (has_neg and has_pos):
                img.putpixel((x, y), (60, 160, 80, 255))
    return img


# ============================================================================
# PSD patcher: adds lfx2 block to ImageMagick-generated PSD
# ============================================================================

def patch_psd_with_lfx2(base_psd_path, out_path, lfx2_data, layer_name=None):
    data = open(base_psd_path, 'rb').read()

    # Parse to find layer record boundaries
    pos = 26
    cm = struct.unpack('>I', data[pos:pos+4])[0]; pos += 4 + cm
    ir = struct.unpack('>I', data[pos:pos+4])[0]; pos += 4 + ir
    lami_pos = pos
    lami_len = struct.unpack('>I', data[pos:pos+4])[0]; pos += 4
    lami_start = pos
    lami_end = lami_pos + 4 + lami_len
    li_len = struct.unpack('>I', data[pos:pos+4])[0]; pos += 4
    li_start = pos
    lc = struct.unpack('>h', data[pos:pos+2])[0]; pos += 2

    # Parse first layer record
    layer_rec_start = pos
    pos += 16  # bounds
    ch_cnt = struct.unpack('>H', data[pos:pos+2])[0]; pos += 2
    for c in range(ch_cnt):
        pos += 6
    pos += 8  # blend
    pos += 4  # opacity/flags
    extra_len = struct.unpack('>I', data[pos:pos+4])[0]; pos += 4
    extra_start = pos
    extra_end = extra_start + extra_len

    # Find channel data start
    record_size = extra_end - layer_rec_start
    ch_data_start = layer_rec_start + record_size

    # Build tagged block
    tagged = bytearray(b'8BIMlfx2')
    tagged.extend(struct.pack('>I', len(lfx2_data)))
    tagged.extend(lfx2_data)
    if len(tagged) % 2:
        tagged.append(0)

    # Build new extra data
    new_extra = bytearray(data[extra_start:extra_end])

    # If layer_name provided, rebuild name in extra data
    # For now, just append lfx2
    new_extra.extend(tagged)

    # Build new layer record
    new_rec = bytearray()
    new_rec.extend(data[layer_rec_start:layer_rec_start + 16 + 2 + ch_cnt * 6 + 8 + 4])
    new_rec.extend(struct.pack('>I', len(new_extra)))
    new_rec.extend(new_extra)

    # Build new LI
    new_li = bytearray()
    new_li.extend(struct.pack('>h', lc))
    new_li.extend(new_rec)
    new_li.extend(data[ch_data_start:lami_end])
    if len(new_li) % 2:
        new_li.append(0)

    # Build new LAMI
    glmi = struct.pack('>I', 0)
    new_lami_len = 4 + len(new_li) + len(glmi)

    # Assemble new file
    new_buf = bytearray()
    new_buf.extend(data[:lami_pos])
    new_buf.extend(struct.pack('>I', new_lami_len))
    new_buf.extend(struct.pack('>I', len(new_li)))
    new_buf.extend(new_li)
    new_buf.extend(glmi)
    new_buf.extend(data[lami_end:])

    with open(out_path, 'wb') as f:
        f.write(new_buf)
    return len(new_buf)


# ============================================================================
# Test file definitions
# ============================================================================

TESTS = [
    ("drop_shadow",      "Drop Shadow",      lfx2_drop_shadow,      create_l_shape),
    ("inner_shadow",     "Inner Shadow",     lfx2_inner_shadow,     create_l_shape),
    ("outer_glow",       "Outer Glow",       lfx2_outer_glow,       create_l_shape),
    ("inner_glow",       "Inner Glow",       lfx2_inner_glow,       create_l_shape),
    ("bevel_inner",      "Bevel Inner",      lfx2_bevel_inner,      create_l_shape),
    ("bevel_outer",      "Bevel Outer",      lfx2_bevel_outer,      create_l_shape),
    ("emboss",           "Emboss",           lfx2_emboss,           create_l_shape),
    ("pillow_emboss",    "Pillow Emboss",    lfx2_pillow_emboss,    create_l_shape),
    ("satin",            "Satin",            lfx2_satin,            create_l_shape),
    ("color_overlay",    "Color Overlay",    lfx2_color_overlay,    create_l_shape),
    ("gradient_overlay", "Gradient Overlay", lfx2_gradient_overlay, create_l_shape),
    ("stroke_inside",    "Stroke Inside",    lfx2_stroke_inside,    create_l_shape),
    ("stroke_center",    "Stroke Center",    lfx2_stroke_center,    create_l_shape),
    ("stroke_outside",   "Stroke Outside",   lfx2_stroke_outside,   create_l_shape),
]


def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # Generate individual test files
    for fname, lname, fx_fn, shape_fn in TESTS:
        out_name = f"{PREFIX}{fname}.psd"
        out_path = os.path.join(OUTPUT_DIR, out_name)

        print(f"  {out_name}...", end=" ", flush=True)

        # Create shape PNG
        img = shape_fn()
        img.save(TMP_PNG)

        # Create base PSD with ImageMagick
        subprocess.run(['convert', TMP_PNG, TMP_PSD], check=True, capture_output=True)

        # Patch with lfx2
        lfx2 = fx_fn()
        size = patch_psd_with_lfx2(TMP_PSD, out_path, lfx2)
        print(f"{size} bytes")

    # Baseline (no effects)
    print(f"  {PREFIX}baseline.psd...", end=" ", flush=True)
    create_l_shape().save(TMP_PNG)
    subprocess.run(['convert', TMP_PNG, TMP_PSD], check=True, capture_output=True)
    import shutil
    shutil.copy2(TMP_PSD, os.path.join(OUTPUT_DIR, f"{PREFIX}baseline.psd"))
    print(f"{os.path.getsize(TMP_PSD)} bytes")

    # Clean up
    os.unlink(TMP_PNG)
    os.unlink(TMP_PSD)

    print(f"\nDone! {len(TESTS) + 1} files in {OUTPUT_DIR}")


if __name__ == '__main__':
    main()