#!/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)

Photon Mono 4 specs:  9024 × 5120 px  |  153.408 × 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)

# ---------------------------------------------------------------------------
# pm4n / ANYCUBIC file format constants (reverse-engineered from Dummy.pm4n)
#
# File layout:
#   0x00  ANYCUBIC magic (8 bytes)
#   0x08  unknown u32 (0)
#   0x0C  version u32 (0x00000206)
#   0x10  section_count u32 (5)
#   0x14  header_size u32 (0x40 = 64)
#   0x18  section table: 5 × (length:u32, offset:u32)
#           entry 0: PREV preview images
#           entry 1: LAYE layer definitions
#           entry 2: MACH machine params
#           entry 3: Mode layer image data (header + RLE blocks)
#           entry 4: HEAD printer config
#
# Sections located by scanning for 4-byte tags from known offsets.
# Tags: HEAD(0x40), PREV(0xB0), LAYE(0x0126D4), MACH(0x0127F0), Mode(0x012974)
# Note: actual offsets vary per dummy file — we scan for tags.
#
# LAYE section layout (at its offset):
#   +0x00  'LAYE' tag
#   +0x04  'REDF' sub-tag
#   +0x08  u32=0
#   +0x0C  u32 payload_size
#   +0x10  u32 (AA level or mode)
#   +0x14  u32 composite_image_offset  (absolute file offset of block 0)
#   +0x18  u32 image_block_size        (size of EACH block, all identical)
#   +0x1C  layer entries begin (N entries × 32 bytes):
#     +0x00  f32 bottom_exposure
#     +0x04  f32 normal_exposure   ← patched with --exposure value
#     +0x08  f32 lift_mm
#     +0x0C  f32 layer_height
#     +0x10  u32 unknown
#     +0x14  u32 unknown
#     +0x18  u32 image_data_offset  (absolute file offset for this layer)
#     +0x1C  u32 image_data_size    (patched when RLE size changes)
#
# Mode section layout:
#   +0x00  'Model\0\0\0' tag (8 bytes)
#   +0x04  u32 sub-header size (108)
#   +0x08..+0x2B  bounding box / Z params as floats
#   +0x2C  'SUBIMGS\0\0\0\0\0' sub-tag (12 bytes)
#   +0x38  u32 SUBIMGS table size
#   +0x3C  u32 layer_count (N)
#   +0x40  u32 bytes_per_pixel (1)
#   +0x44  u32 first_image_offset  (= composite_image_offset, same as LAYE+0x14)
#   +0x48  u32 first_image_size    (= image_block_size, same as LAYE+0x18)
#   +0x4C  u32 unknown
#   then N RLE image blocks follow (at offsets stored in LAYE entries)
# ---------------------------------------------------------------------------

LAYE_TAG      = b'LAYE'
MODE_TAG      = b'Mode'
ENTRY_STRIDE  = 0x20   # 32 bytes per layer entry in LAYE
LAYE_HDR_SIZE = 0x1C   # bytes before first entry


def find_tag(data: bytes, tag: bytes, start: int = 0) -> int:
    """Return file offset of first occurrence of tag (exact 4-byte match at 4-byte boundary)."""
    i = start
    while i + 4 <= len(data):
        if data[i:i+4] == tag:
            return i
        i += 4
    # Fall back to unaligned search
    pos = data.find(tag, start)
    if pos < 0:
        raise ValueError(f"Tag {tag!r} not found in file")
    return pos


def count_laye_entries(data: bytes, laye_off: int) -> int:
    """Count layer entries by scanning until we hit the 'EXTR' sub-section or end of section."""
    entry_start = laye_off + LAYE_HDR_SIZE
    n = 0
    while True:
        pos = entry_start + n * ENTRY_STRIDE
        if pos + 4 > len(data):
            break
        word = data[pos:pos+4]
        # Stop if we hit a known sub-tag marker
        if word in (b'EXTR', b'MACH', b'Mode', b'HEAD', b'PREV'):
            break
        # Stop if the f32 at this position is not a plausible exposure time
        v = struct.unpack_from('<f', data, pos)[0]
        if not (0.0 < v < 1000.0):
            break
        n += 1
    return n


def pack_f32(v: float) -> bytes:
    return struct.pack('<f', v)


def pack_u32(v: int) -> bytes:
    return struct.pack('<I', v)


def unpack_u32(data: bytes, off: int) -> int:
    return struct.unpack_from('<I', data, off)[0]


# ---------------------------------------------------------------------------
# Photon Workshop RLE  (BW — 2 bytes per run)
#   Byte0 [7:4] colour nibble: 0x0=black, 0xF=white
#   Byte0 [3:0] + Byte1: run length - 1  (12-bit, max run=4096)
# ---------------------------------------------------------------------------
MAX_RUN = 4096

def encode_rle(pixels: bytes) -> bytes:
    out = bytearray()
    i, n = 0, 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
        enc = run - 1
        out.append((nibble << 4) | ((enc >> 8) & 0x0F))
        out.append(enc & 0xFF)
        i = j
    return bytes(out)


# ---------------------------------------------------------------------------
# Gerber → PIL Image at LCD resolution
# ---------------------------------------------------------------------------

def render_gerber(gbr_path: Path, dpmm: float,
                  invert: bool, mirror: bool,
                  pos_mm: tuple | None) -> Image.Image:
    try:
        from pygerber.gerberx3.api.v2 import (
            GerberFile, ColorScheme, PixelFormatEnum, ImageFormatEnum
        )
    except ImportError:
        sys.exit(
            "ERROR: pygerber not found.\n"
            "Activate the venv:  source .venv/bin/activate\n"
            "Or install:         pip install pygerber Pillow"
        )

    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')

    cw, ch = layer_img.size
    canvas  = Image.new('L', (LCD_W_PX, LCD_H_PX), 0)

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

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

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

    # Hard-binarise to strict 0/255
    canvas = canvas.point(lambda v: 255 if v >= 128 else 0)
    return canvas


# ---------------------------------------------------------------------------
# pm4n surgery — rewrite with exact format knowledge
# ---------------------------------------------------------------------------

def patch_pm4n(dummy_path: Path, image: Image.Image,
               exposure_sec: float, output_path: Path):

    raw = bytearray(dummy_path.read_bytes())

    # --- encode new RLE ---
    new_rle = encode_rle(image.convert('L').tobytes())
    new_rle_size = len(new_rle)

    # --- locate LAYE section ---
    laye_off = find_tag(raw, LAYE_TAG)
    n_entries = count_laye_entries(raw, laye_off)
    print(f"  LAYE at 0x{laye_off:06X}, {n_entries} layer entries")

    # Read composite image offset and original block size
    composite_off = unpack_u32(raw, laye_off + 0x14)
    old_block_size = unpack_u32(raw, laye_off + 0x18)
    print(f"  Image blocks: first=0x{composite_off:06X}, old_size={old_block_size}, new_size={new_rle_size}")

    # --- Read all existing image offsets from LAYE entries ---
    # Block 0 is the composite (at composite_off), blocks 1..N from entries
    old_offsets = [composite_off]
    for i in range(n_entries):
        entry_base = laye_off + LAYE_HDR_SIZE + i * ENTRY_STRIDE
        old_offsets.append(unpack_u32(raw, entry_base + 0x18))

    # All blocks should be contiguous and equal-sized; verify
    expected = composite_off
    for i, off in enumerate(old_offsets):
        if off != expected:
            print(f"  WARNING: block {i} offset 0x{off:06X} != expected 0x{expected:06X}")
        expected = off + old_block_size

    # --- Patch exposure time in all LAYE entries ---
    for i in range(n_entries):
        entry_base = laye_off + LAYE_HDR_SIZE + i * ENTRY_STRIDE
        struct.pack_into('<f', raw, entry_base + 0x04, exposure_sec)
    print(f"  Patched exposure to {exposure_sec}s in {n_entries} entries")

    # --- Build new image data section ---
    # All N+1 blocks (composite + N layers) get the SAME new RLE
    # (single-layer exposure: every layer shows the same image)
    n_blocks = n_entries + 1
    new_image_section = new_rle * n_blocks

    # --- Reconstruct file ---
    # Everything before the first image block stays unchanged
    prefix = bytes(raw[:composite_off])

    # --- Update LAYE: block size field ---
    struct.pack_into('<I', raw, laye_off + 0x18, new_rle_size)

    # --- Update Mode section: first_image_size field ---
    # Mode sub-header: tag 'Mode' + u32(108) + ... + 'SUBIMGS' + ...
    # first_image_size is at Mode_off + 0x48
    mode_off = find_tag(raw, MODE_TAG)
    struct.pack_into('<I', raw, mode_off + 0x48, new_rle_size)
    print(f"  Mode at 0x{mode_off:06X}, patched first_image_size")

    # --- Update LAYE entries: image offsets and sizes ---
    new_composite_off = composite_off  # composite block stays at same position
    for i in range(n_entries):
        entry_base = laye_off + LAYE_HDR_SIZE + i * ENTRY_STRIDE
        new_layer_off = composite_off + (i + 1) * new_rle_size
        struct.pack_into('<I', raw, entry_base + 0x18, new_layer_off)
        struct.pack_into('<I', raw, entry_base + 0x1C, new_rle_size)

    # --- Splice new image data ---
    old_images_end = composite_off + n_blocks * old_block_size
    raw[composite_off:old_images_end] = new_image_section

    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',
    )
    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)
    p.add_argument('--invert',       action='store_true', help='Invert image (positive-working resist)')
    p.add_argument('--mirror',       action='store_true', help='Mirror X (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)
    p.add_argument('--pos',          default=None, help='Board 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()

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


if __name__ == '__main__':
    main()