#!/usr/bin/env python3
"""
Minimal PSD writer — byte-exact per Adobe Photoshop File Formats Specification.
Creates a 256x256 RGBA document with one layer (red rectangle).
"""
import struct, os

def be16(v): return struct.pack(">H", v)
def be32(v): return struct.pack(">I", v)
def bei16(v): return struct.pack(">h", v)

def packbits_compress(row):
    out = bytearray()
    n = len(row)
    i = 0
    while i < n:
        # find run
        run = 1
        while run < 128 and i + run < n and row[i+run] == row[i]:
            run += 1
        if run > 1:
            out.append((1 - run) & 0xFF)
            out.append(row[i])
            i += run
        else:
            # literal
            start = i
            lit = 0
            while lit < 128 and i + lit < n:
                nxt = 0
                if i + lit + 1 < n and row[i+lit] == row[i+lit+1]:
                    r = 2
                    while r < 128 and i+lit+r < n and row[i+lit+r] == row[i+lit]:
                        r += 1
                    nxt = r
                if nxt >= 3:
                    break
                lit += 1
            if lit == 0:
                lit = 1
            out.append(lit - 1)
            out.extend(row[start:start+lit])
            i += lit
    return bytes(out)

def write_resource(buf, res_id, name, data):
    buf.extend(b"8BIM")
    buf.extend(be16(res_id))
    name_padded = name + b'\x00' if (1 + len(name)) % 2 == 1 else name
    buf.append(len(name))
    buf.extend(name)
    if (1 + len(name)) % 2 == 1:
        buf.append(0)
    buf.extend(be32(len(data)))
    buf.extend(data)
    if len(data) % 2 != 1:
        pass  # even already
    # pad data to even
    if len(data) % 2 != 0:
        buf.append(0)

def write_pascal_padded4(buf, s):
    """Write Pascal string padded to multiple of 4 bytes."""
    b = s.encode("ascii") if isinstance(s, str) else s
    buf.append(len(b))
    buf.extend(b)
    total = 1 + len(b)
    pad = (4 - total % 4) % 4
    buf.extend(b'\x00' * pad)

def write_tagged_block(buf, key, data):
    """Write 8BIM tagged block."""
    buf.extend(b"8BIM")
    buf.extend(key)
    buf.extend(be32(len(data)))
    buf.extend(data)
    if len(data) % 2 != 0:
        buf.append(0)

def make_descriptor_string(s):
    """Write a descriptor UTF-16 string: length(4) + chars(2*N)."""
    enc = s.encode("utf-16-be")
    return be32(len(enc) // 2) + enc

def main():
    W, H = 256, 256
    MARGIN = 64

    # ── Create layer pixels ──
    rgba = bytearray(W * H * 4)
    for y in range(H):
        for x in range(W):
            i = (y * W + x) * 4
            if MARGIN <= x < W - MARGIN and MARGIN <= y < H - MARGIN:
                rgba[i]   = 255  # R
                rgba[i+1] = 0    # G
                rgba[i+2] = 0    # B
                rgba[i+3] = 255  # A

    # ── Header (26 bytes) ──
    header = bytearray()
    header.extend(b"8BPS")
    header.extend(be16(1))          # version
    header.extend(b'\x00' * 6)      # reserved
    header.extend(be16(4))          # channels (RGBA)
    header.extend(be32(H))          # height
    header.extend(be32(W))          # width
    header.extend(be16(8))          # depth
    header.extend(be16(3))          # color mode (RGB)

    # ── Color Mode Data ──
    color_mode = be32(0)

    # ── Image Resources ──
    img_res = bytearray()
    # 0x03ED: Resolution info
    res_data = bytearray()
    res_data.extend(be32(72 << 16))  # hRes (72 DPI fixed)
    res_data.extend(be16(1))         # hResUnit (1=pixels/inch)
    res_data.extend(be16(1))         # vResUnit
    res_data.extend(be32(72 << 16))  # vRes
    res_data.extend(be16(1))         # widthUnit
    res_data.extend(be16(1))         # heightUnit
    write_resource(img_res, 0x03ED, b"", bytes(res_data))

    # 0x0419: Global Angle
    write_resource(img_res, 0x0419, b"", be32(120))

    # pad img_res to even
    if len(img_res) % 2 != 0:
        img_res.append(0)

    # ── Layer & Mask Information ──
    # Build layer info section first
    layer_info = bytearray()

    # Layer count (positive = no groups)
    layer_info.extend(bei16(1))

    # Layer record
    layer_rec = bytearray()
    # Rect: top, left, bottom, right
    layer_rec.extend(be32(0))   # top
    layer_rec.extend(be32(0))   # left
    layer_rec.extend(be32(H))   # bottom
    layer_rec.extend(be32(W))   # right

    # Channel count
    layer_rec.extend(be16(4))

    # Channel info: (-1=alpha, 0=R, 1=G, 2=B)
    # We'll compute channel data lengths later
    ch_ids = [-1, 0, 1, 2]
    ch_data_list = []
    for ch_id in ch_ids:
        plane = bytearray()
        for y in range(H):
            for x in range(W):
                idx = (y * W + x) * 4
                if ch_id == -1:
                    plane.append(rgba[idx + 3])
                elif ch_id == 0:
                    plane.append(rgba[idx])
                elif ch_id == 1:
                    plane.append(rgba[idx + 1])
                elif ch_id == 2:
                    plane.append(rgba[idx + 2])
        ch_data_list.append(plane)

    # Compress each channel
    ch_compressed = []
    for plane in ch_data_list:
        row_lengths = []
        compressed_rows = []
        for y in range(H):
            row = plane[y * W : (y + 1) * W]
            compressed = packbits_compress(row)
            row_lengths.append(len(compressed))
            compressed_rows.append(compressed)
        ch_compressed.append((row_lengths, compressed_rows))

    # Channel data = compression(2) + row_lengths(H*2) + compressed_data
    ch_packed = []
    for row_lengths, compressed_rows in ch_compressed:
        buf = bytearray()
        buf.extend(be16(1))  # compression type = RLE
        for rl in row_lengths:
            buf.extend(be16(rl))
        for cr in compressed_rows:
            buf.extend(cr)
        ch_packed.append(bytes(buf))

    # Write channel info records
    for i, ch_id in enumerate(ch_ids):
        layer_rec.extend(bei16(ch_id))
        layer_rec.extend(be32(len(ch_packed[i])))

    # Blend mode
    layer_rec.extend(b"8BIMnorm")
    layer_rec.append(255)   # opacity
    layer_rec.append(0)     # clipping (base)
    layer_rec.append(0)     # flags (visible)
    layer_rec.append(0)     # filler

    # Extra data
    extra = bytearray()
    # Layer mask data (4 bytes, length=0)
    extra.extend(be32(0))
    # Blending ranges (4 bytes length=40, then 40 bytes of data)
    extra.extend(be32(40))
    for _ in range(10):
        extra.extend(be16(0))
        extra.extend(be16(65535))
    # Layer name (Pascal string padded to 4)
    write_pascal_padded4(extra, "Red Rectangle")

    layer_rec.extend(be32(len(extra)))
    layer_rec.extend(extra)

    # Append channel data
    layer_info.extend(layer_rec)
    for ch_data in ch_packed:
        layer_info.extend(ch_data)

    # Pad layer info to even
    if len(layer_info) % 2 != 0:
        layer_info.append(0)

    # Global Layer Mask Info (empty, 4 bytes)
    glmi = be32(0)

    # Layer & Mask Information section
    lami = bytearray()
    lami.extend(be32(len(layer_info) + len(glmi)))  # length of rest
    lami.extend(layer_info)
    lami.extend(glmi)

    # ── Merged Image Data ──
    merged = bytearray()
    merged.extend(be16(1))  # compression = RLE

    # Row lengths for all channels (per-channel order: RRR...GGG...BBB...AAA...)
    for ch_offset in range(4):  # R, G, B, A
        for y in range(H):
            row = bytes(rgba[(y * W + ch_offset)::4][:W])  # wrong, need proper extraction
            row = bytearray()
            for x in range(W):
                idx = (y * W + x) * 4 + ch_offset
                row.append(rgba[idx])
            compressed = packbits_compress(row)
            merged.extend(be16(len(compressed)))

    # Compressed data for all channels
    for ch_offset in range(4):
        for y in range(H):
            row = bytearray()
            for x in range(W):
                idx = (y * W + x) * 4 + ch_offset
                row.append(rgba[idx])
            compressed = packbits_compress(row)
            merged.extend(compressed)

    # ── Assemble PSD ──
    psd = bytearray()
    psd.extend(header)
    psd.extend(color_mode)
    psd.extend(be32(len(img_res)))
    psd.extend(img_res)
    psd.extend(lami)
    psd.extend(merged)

    out_path = "_tmp/python_ref_test.psd"
    with open(out_path, "wb") as f:
        f.write(psd)
    print(f"Wrote {len(psd)} bytes to {out_path}")

    # Also verify roundtrip
    assert psd[:4] == b"8BPS"
    assert len(psd) == 26 + 4 + 4 + len(img_res) + len(lami) + len(merged)
    print("Roundtrip check: OK")

if __name__ == "__main__":
    main()