#!/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  Exposure time in seconds  [default: 60]
    --dpmm N        Render resolution in dots/mm  [default: 58.824, native 17µm/px]
    --pos X,Y       Board position mm from top-left (default: centred)
    --verbose       Print detailed progress

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 format constants  (reverse-engineered from Dummy.pm4n)
# ---------------------------------------------------------------------------
LAYE_TAG      = b'LAYE'
MODE_TAG      = b'Mode'
ENTRY_STRIDE  = 0x20   # 32 bytes per layer entry
LAYE_HDR_SIZE = 0x20   # bytes before first entry within LAYE section
#
# LAYE header layout (verified from binary probe):
#   +0x00  tag  "LAYE"
#   +0x04  tag  "RDEF"
#   +0x08  u32  0
#   +0x0C  u32  0xC4
#   +0x10  u32  n_entries
#   +0x14  u32  composite_image_offset
#   +0x18  u32  block_size  (RLE bytes per block)
#   +0x1C  f32  lift_height_mm
#
# LAYE entry layout (0x20 bytes each, immediately after header):
#   +0x00  f32  exposure_sec        <-- was wrongly assumed at +0x04
#   +0x04  f32  z_position_mm
#   +0x08  f32  layer_thickness_mm
#   +0x0C  u32  unknown
#   +0x10  u32  unknown
#   +0x14  u32  image_data_offset   <-- was wrongly assumed at +0x18
#   +0x18  u32  image_data_size     <-- was wrongly assumed at +0x1C
#   +0x1C  f32  lift_speed


def find_tag(data: bytes, tag: bytes, start: int = 0) -> int:
    """Return file offset of first occurrence of tag aligned to 4 bytes."""
    i = (start + 3) & ~3
    while i + 4 <= len(data):
        if data[i:i+4] == tag:
            return i
        i += 4
    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:
    """Read n_entries directly from LAYE header at +0x10."""
    return unpack_u32(data, laye_off + 0x10)


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,
                  verbose: bool = False) -> 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)

    canvas = canvas.point(lambda v: 255 if v >= 128 else 0)

    if verbose:
        print(f"  Gerber rendered: {layer_img.size[0]}×{layer_img.size[1]} px"
              f"  placed at ({px},{py})  invert={invert}  mirror={mirror}")

    return canvas


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

def patch_pm4n(dummy_path: Path, image: Image.Image,
               exposure_sec: float, output_path: Path,
               verbose: bool = False):

    def log(*a):
        if verbose:
            print(*a)

    raw = bytearray(dummy_path.read_bytes())

    new_rle      = encode_rle(image.convert('L').tobytes())
    new_rle_size = len(new_rle)

    laye_off   = find_tag(raw, LAYE_TAG)
    n_entries  = count_laye_entries(raw, laye_off)
    log(f"  LAYE at 0x{laye_off:06X}, {n_entries} layer entries")

    composite_off  = unpack_u32(raw, laye_off + 0x14)
    old_block_size = unpack_u32(raw, laye_off + 0x18)
    log(f"  Image blocks: first=0x{composite_off:06X}, "
        f"old_size={old_block_size}, new_size={new_rle_size}")

    # Patch exposure in all entries  (entry+0x00 = exposure_sec float)
    for i in range(n_entries):
        base = laye_off + LAYE_HDR_SIZE + i * ENTRY_STRIDE
        struct.pack_into('<f', raw, base + 0x00, exposure_sec)
    log(f"  Exposure patched to {exposure_sec}s in {n_entries} entries")

    # Update block size in LAYE header and Mode header
    struct.pack_into('<I', raw, laye_off + 0x18, new_rle_size)
    mode_off = find_tag(raw, MODE_TAG)
    struct.pack_into('<I', raw, mode_off + 0x48, new_rle_size)
    log(f"  Mode at 0x{mode_off:06X}")

    # Update image offsets and sizes in all entries.
    # Layout: [composite block][layer-0 block][layer-1 block]...
    # entry+0x14 = data offset, entry+0x18 = data size
    for i in range(n_entries):
        base = laye_off + LAYE_HDR_SIZE + i * ENTRY_STRIDE
        struct.pack_into('<I', raw, base + 0x14, composite_off + (i + 1) * new_rle_size)
        struct.pack_into('<I', raw, base + 0x18, new_rle_size)

    # Splice new image data (composite block + one block per layer, all identical)
    # Layout: composite block at composite_off, then one block per entry.
    # The file contains exactly n_entries blocks total (composite counts as block 0;
    # the last entry's data_off is therefore one block past the last stored block,
    # which Chitubox tolerates). We mirror the same layout.
    n_blocks  = n_entries          # composite + (n_entries-1) layer blocks = n_entries total
    old_end   = composite_off + n_blocks * old_block_size
    raw[composite_off:old_end] = new_rle * n_blocks

    output_path.write_bytes(raw)
    log(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')
    p.add_argument('gerber',         help='Input Gerber file')
    p.add_argument('-o', '--output', default=None)
    p.add_argument('--invert',       action='store_true')
    p.add_argument('--mirror',       action='store_true')
    p.add_argument('--exposure',     type=float, default=60.0)
    p.add_argument('--dpmm',         type=float, default=NATIVE_DPMM)
    p.add_argument('--pos',          default=None)
    p.add_argument('--verbose',      action='store_true')
    p.add_argument('--preview-scale', type=int, default=1, metavar='N',
                   help='Downsample preview PNG by factor N (default: 1 = full resolution)')
    return p.parse_args()


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

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

    if args.output:
        out = Path(args.output)
    else:
        # If in 'output/gerbers/', default to 'output/pm4n/'
        parent = gbr.parent
        if parent.name == 'gerbers' and parent.parent.name == 'output':
            out = parent.parent / 'pm4n' / gbr.with_suffix('.pm4n').name
        else:
            out = 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")

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

    img = render_gerber(gbr, dpmm=args.dpmm, invert=args.invert,
                        mirror=args.mirror, pos_mm=pos_mm, verbose=v)

    preview = out.with_suffix('.preview.png')
    scale = args.preview_scale
    if scale <= 1:
        img.save(preview)
    else:
        img.resize((img.size[0] // scale, img.size[1] // scale), Image.NEAREST).save(preview)

    patch_pm4n(dummy, img, args.exposure, out, verbose=v)

    # Always print the output path (quiet mode only output)
    print(out)
    print(preview)


if __name__ == '__main__':
    main()