#!/usr/bin/env python3
"""
gerber_to_pm4n.py  -  Anycubic Photon Mono 4 PCB exposure file generator

Usage:
    python3 gerber_to_pm4n.py <dummy.pm4n> <board.gbr> [options]

Options:
    -o OUTPUT       Output file path  [default: <board>.pm4n]
    --invert        Invert the image (for positive-working resist like Bungard standard)
    --mirror        Mirror X axis (for copper-side-down placement on FEP)
    --exposure SEC  Layer exposure time in seconds  [default: 60]
    --dpmm N        Render resolution in dots/mm    [default: 58.824, native 17µm/px]
    --pos X,Y       Place board at X,Y mm from top-left (default: centred on LCD)
    --help          Show this message

Photon Mono 4 specs:  9024 x 5120 px  |  153.408 x 87.040 mm  |  17.001 µm/px
"""

import argparse
import struct
import sys
import io
from pathlib import Path
from PIL import Image, ImageOps

# ---------------------------------------------------------------------------
# Printer constants
# ---------------------------------------------------------------------------
LCD_W_PX    = 9024
LCD_H_PX    = 5120
LCD_W_MM    = 153.408
LCD_H_MM    = 87.040
NATIVE_DPMM = LCD_W_PX / LCD_W_MM   # 58.824 dpmm  (1 px ≈ 17.001 µm)

# ---------------------------------------------------------------------------
# Photon Workshop RLE  (BW — 2 bytes per run)
#
#  Byte0  [7:4] = colour nibble  (0x0 = black,  0xF = white)
#  Byte0  [3:0] = high 4 bits of run length  (bits 11:8)
#  Byte1         = low 8 bits of run length  (bits 7:0)
#  Run length encodes (n-1): 0x000 = 1 pixel, 0xFFF = 4096 pixels
# ---------------------------------------------------------------------------
MAX_RUN = 4096

def encode_rle(pixels: bytes) -> bytes:
    """Encode flat 0x00/0xFF bytes → Photon Workshop BW RLE."""
    out = bytearray()
    i = 0
    n = len(pixels)
    while i < n:
        colour  = pixels[i]
        nibble  = 0xF if colour >= 0x80 else 0x0
        j = i + 1
        while j < n and pixels[j] == colour and (j - i) < MAX_RUN:
            j += 1
        run     = j - i
        encoded = run - 1
        out.append((nibble << 4) | ((encoded >> 8) & 0x0F))
        out.append(encoded & 0xFF)
        i = j
    return bytes(out)


def decode_rle(data: bytes, expected_pixels: int) -> bytes:
    """Decode PW RLE → raw pixel bytes (used for verification)."""
    out = bytearray()
    i = 0
    while i + 1 < len(data):
        b0, b1  = data[i], data[i + 1]
        nibble  = (b0 >> 4) & 0x0F
        colour  = 0xFF if nibble == 0xF else 0x00
        run     = (((b0 & 0x0F) << 8) | b1) + 1
        out.extend([colour] * run)
        i += 2
    return bytes(out[:expected_pixels])


# ---------------------------------------------------------------------------
# pm4n binary surgery
#
# Photon Workshop file = sequence of tagged sections:
#   tag:4  length:4  payload:length
#
# Sections we care about:
#   HEAD  – contains exposure time as a float somewhere in a packed struct
#   LAYE  – layer definition table: count:u32 then N × entry(28 bytes)
#             entry[0:4]  = absolute file offset of RLE blob
#             entry[4:8]  = RLE blob length in bytes
#   After the sections: raw RLE layer image blobs (referenced by LAYE offsets)
# ---------------------------------------------------------------------------
SECTION_HDR = 8   # 4-byte tag + 4-byte length

def read_sections(data: bytes) -> list:
    sections = []
    i = 0
    while i + SECTION_HDR <= len(data):
        tag    = data[i:i+4]
        length = struct.unpack_from('<I', data, i+4)[0]
        sections.append((tag, i + SECTION_HDR, length))
        i += SECTION_HDR + length
    return sections

def find_section(data: bytes, tag: bytes) -> tuple:
    for t, off, ln in read_sections(data):
        if t == tag:
            return off, ln
    raise ValueError(f"Section {tag!r} not found in file")

def patch_u32(data: bytearray, offset: int, value: int):
    struct.pack_into('<I', data, offset, value)

def patch_f32(data: bytearray, offset: int, value: float):
    struct.pack_into('<f', data, offset, value)


# ---------------------------------------------------------------------------
# Gerber → PIL Image
# ---------------------------------------------------------------------------

def render_gerber(gbr_path: Path, dpmm: float,
                  invert: bool, mirror: bool,
                  pos_mm: tuple | None) -> Image.Image:
    """
    Render a Gerber file to a binary PIL image sized to the Photon Mono 4 LCD.
    copper = white on black background before any transforms.
    """
    try:
        from pygerber.gerberx3.api.v2 import (
            GerberFile, ColorScheme, PixelFormatEnum, ImageFormatEnum
        )
    except ImportError:
        sys.exit(
            "ERROR: pygerber not found.\n"
            "Activate the venv first:  source .venv/bin/activate\n"
            "Or install:  pip install pygerber Pillow numpy"
        )

    buf = io.BytesIO()
    (GerberFile
        .from_file(str(gbr_path))
        .parse()
        .render_raster(
            buf,
            dpmm=int(round(dpmm)),
            color_scheme=ColorScheme.DEFAULT_GRAYSCALE,
            pixel_format=PixelFormatEnum.RGB,
            image_format=ImageFormatEnum.PNG,
        )
    )
    buf.seek(0)
    layer_img = Image.open(buf).convert('L')

    # Place onto full LCD canvas
    cw, ch  = layer_img.size
    canvas  = Image.new('L', (LCD_W_PX, LCD_H_PX), 0)

    if pos_mm is not None:
        px = int(round(pos_mm[0] * dpmm))
        py = int(round(pos_mm[1] * dpmm))
    else:
        px = (LCD_W_PX - cw) // 2
        py = (LCD_H_PX - ch) // 2

    canvas.paste(layer_img, (max(0, px), max(0, py)))

    if mirror:
        canvas = ImageOps.mirror(canvas)
    if invert:
        canvas = ImageOps.invert(canvas)

    # Hard-binarise: no antialiasing artefacts in the RLE stream
    canvas = canvas.point(lambda v: 255 if v >= 128 else 0)

    return canvas


# ---------------------------------------------------------------------------
# pm4n surgery
# ---------------------------------------------------------------------------

def patch_pm4n(dummy_path: Path, image: Image.Image,
               exposure_sec: float, output_path: Path):
    """Replace layer RLE + exposure time in a dummy .pm4n, write output."""
    raw = bytearray(dummy_path.read_bytes())

    # Encode new layer image
    pixels  = image.convert('L').tobytes()
    new_rle = encode_rle(pixels)

    # Patch exposure time: scan HEAD for any float in 0.5–600 s range
    hdr_off, hdr_len = find_section(raw, b'HEAD')
    for off in range(hdr_off, hdr_off + hdr_len - 3):
        val = struct.unpack_from('<f', raw, off)[0]
        if 0.5 <= val <= 600.0:
            patch_f32(raw, off, exposure_sec)

    # Locate layer image via LAYERDEF
    ld_off, _      = find_section(raw, b'LAYE')
    layer_count    = struct.unpack_from('<I', raw, ld_off)[0]
    if layer_count != 1:
        print(f"WARNING: dummy has {layer_count} layers; only layer 0 will be replaced.")
    entry_off      = ld_off + 4
    img_offset     = struct.unpack_from('<I', raw, entry_off)[0]
    img_len        = struct.unpack_from('<I', raw, entry_off + 4)[0]

    print(f"  Dummy RLE:  offset=0x{img_offset:08X}  {img_len} bytes")
    print(f"  New RLE:    {len(new_rle)} bytes")

    # Splice new RLE in place
    old_end = img_offset + img_len
    raw[img_offset:old_end] = new_rle

    # Update LAYERDEF length field
    patch_u32(raw, entry_off + 4, len(new_rle))

    # Update any enclosing section's length field
    for tag, sec_off, sec_len in read_sections(bytes(raw)):
        if tag not in (b'HEAD', b'LAYE', b'PREV') and sec_off <= img_offset < sec_off + sec_len:
            patch_u32(raw, sec_off - 4, sec_len + len(new_rle) - img_len)
            break

    output_path.write_bytes(raw)
    print(f"  Written: {output_path}  ({len(raw):,} bytes)")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(
        description='Convert Gerber → Anycubic Photon Mono 4 .pm4n PCB exposure file',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument('dummy',          help='Dummy .pm4n template (from Photon Workshop)')
    p.add_argument('gerber',         help='Input Gerber file')
    p.add_argument('-o', '--output', default=None,    help='Output .pm4n path')
    p.add_argument('--invert',       action='store_true', help='Invert image (positive-working resist)')
    p.add_argument('--mirror',       action='store_true', help='Mirror X axis (copper-side-down placement)')
    p.add_argument('--exposure',     type=float, default=60.0,         help='Exposure seconds (default: 60)')
    p.add_argument('--dpmm',         type=float, default=NATIVE_DPMM,  help=f'Dots/mm (default: {NATIVE_DPMM:.3f})')
    p.add_argument('--pos',          default=None,    help='Board position X,Y mm from top-left')
    return p.parse_args()


def main():
    args  = parse_args()
    dummy = Path(args.dummy)
    gbr   = Path(args.gerber)

    for p, label in [(dummy, 'dummy'), (gbr, 'gerber')]:
        if not p.exists():
            sys.exit(f"ERROR: {label} file not found: {p}")

    out    = Path(args.output) if args.output else gbr.with_suffix('.pm4n')
    pos_mm = None
    if args.pos:
        try:
            x, y   = map(float, args.pos.split(','))
            pos_mm = (x, y)
        except Exception:
            sys.exit("ERROR: --pos must be X,Y  e.g.  --pos 10.5,8.0")

    print(f"Gerber:   {gbr}")
    print(f"Dummy:    {dummy}")
    print(f"Output:   {out}")
    print(f"Invert:   {args.invert}  Mirror: {args.mirror}  Exposure: {args.exposure}s  dpmm: {args.dpmm:.3f}")
    print()

    print("Rendering Gerber...")
    img = render_gerber(gbr, dpmm=args.dpmm, invert=args.invert,
                        mirror=args.mirror, pos_mm=pos_mm)
    print(f"Canvas: {img.size[0]}×{img.size[1]} px")

    preview = out.with_suffix('.preview.png')
    img.resize((img.size[0] // 4, img.size[1] // 4), Image.NEAREST).save(preview)
    print(f"Preview: {preview}")

    print("Patching .pm4n...")
    patch_pm4n(dummy, img, args.exposure, out)
    print("Done.")


if __name__ == '__main__':
    main()