#!/usr/bin/env python3

import subprocess
import zipfile
import tempfile
import os.path as path
import os
import sys
import time
import shutil
import math

import tqdm
import gerber
from gerber.render import GerberCairoContext
import numpy as np
import cv2

def paste_image(
        target_gerber:str,
        outline_gerber:str,
        source_img:np.ndarray,
        subtract_gerber:list=[],
        extend_overlay_r_mil:float=6,
        extend_picture_r_mil:float=2,
        status_print=lambda *args:None,
        debugdir:str=None):
    debugctr = 0
    def debugimg(img, name):
        nonlocal debugctr
        if debugdir:
            cv2.imwrite(path.join(debugdir, '{:02d}{}.png'.format(debugctr, name)), img)
        debugctr += 1

    # Parse outline layer to get bounds of gerber file
    status_print('Parsing outline gerber')
    outline = gerber.loads(outline_gerber)
    (minx, maxx), (miny, maxy) = outline.bounds
    grbw, grbh = maxx - minx, maxy - miny
    status_print('  * outline has offset {}, size {}'.format((minx, miny), (grbw, grbh)))

    # Read source image
    imgh, imgw = source_img.shape
    scale = math.ceil(max(imgw/grbw, imgh/grbh)) # scale is in dpi
    status_print('  * source image has size {}, going for scale {}dpmm'.format((imgw, imgh), scale))

    # Parse target layer
    status_print('Parsing target gerber')
    target = gerber.loads(target_gerber)
    (tminx, tmaxx), (tminy, tmaxy) = target.bounds
    status_print('  * target layer has offset {}, size {}'.format((tminx, tminy), (tmaxx-tminx, tmaxy-tminy)))

    # Render all gerber layers whose features are to be excluded from the target image, such as board outline, the
    # original silk layer and the solder paste layer to binary images.
    with tempfile.TemporaryDirectory() as tmpdir:
        img_file = path.join(tmpdir, 'target.png')

        status_print('Combining keepout composite')
        fg, bg = gerber.render.RenderSettings((1, 1, 1)), gerber.render.RenderSettings((0, 0, 0))
        ctx = GerberCairoContext(scale=scale)
        status_print('  * outline')
        ctx.render_layer(outline, settings=fg, bgsettings=bg)
        status_print('  * target layer')
        ctx.render_layer(target, settings=fg, bgsettings=bg)
        for fn, sub in subtract_gerber:
            status_print('  * extra layer', os.path.basename(fn))
            layer = gerber.loads(sub)
            ctx.render_layer(layer, settings=fg, bgsettings=bg)
        status_print('Rendering keepout composite')
        ctx.dump(img_file)

        # Vertically flip exported image
        original_img = cv2.imread(img_file, cv2.IMREAD_GRAYSCALE)[::-1, :]

    r = 1+2*max(1, int(extend_overlay_r_mil/1000 * 25.4 * scale))
    status_print('Expanding keepout composite by', r, extend_overlay_r_mil/1000 * 25.4 * scale, scale, grbw, grbh)

    # Extend image by a few pixels and flood-fill from (0, 0) to mask out the area outside the outermost outline
    # This ensures no polygons are generated outside the board even for non-rectangular boards.
    border = 10
    outh, outw = original_img.shape
    extended_img = np.zeros((outh + 2*border, outw + 2*border), dtype=np.uint8)
    extended_img[border:outh+border, border:outw+border] = original_img
    cv2.floodFill(extended_img, None, (0, 0), (255,))
    original_img = extended_img[border:outh+border, border:outw+border]
    debugimg(extended_img, 'flooded')

    # Dilate the white areas of the image using gaussian blur and threshold. Use these instead of primitive dilation
    # here for their non-directionality.
    target_img = cv2.blur(original_img, (r, r))
    _, target_img = cv2.threshold(target_img, 255//(1+r), 255, cv2.THRESH_BINARY)

    # Threshold source image. Ideally, the source image is already binary but in case it's not, or in case it's not
    # exactly binary (having a few very dark or very light grays e.g. due to JPEG compression) we're thresholding here.
    status_print('Thresholding source image')
    qr = 1+2*max(1, int(extend_picture_r_mil/1000 * scale))
    source_img = source_img[::-1]
    _, source_img = cv2.threshold(source_img, 127, 255, cv2.THRESH_BINARY)
    debugimg(source_img, 'thresh')

    # Pad image to size of target layer images generated above. After this, `scale` applies to the padded image as well
    # as the gerber renders. For padding, zoom or shrink the image to completely fit the gerber's rectangular bounding
    # box. Center the image vertically or horizontally if it has a different aspect ratio.
    status_print('Padding source image')
    tgth, tgtw = target_img.shape
    padded_img = np.zeros(shape=target_img.shape, dtype=source_img.dtype)
    offx = int((minx-tminx if tminx < minx else 0)*scale)
    offy = int((miny-tminy if tminy < miny else 0)*scale)
    offx += int(grbw*scale - imgw) // 2
    offy += int(grbh*scale - imgh) // 2
    endx, endy = min(offx+imgw, tgtw), min(offy+imgh, tgth)
    print('off', (offx, offy), 'end', (endx, endy), 'img', (imgw, imgh), 'tgt', (tgtw, tgth))
    padded_img[offy:endy, offx:endx] = source_img[:endy-offy, :endx-offx]
    debugimg(padded_img, 'padded')
    debugimg(target_img, 'target')

    # Mask out excluded gerber features (source silk, holes, solder mask etc.) from the target image
    status_print('Masking source image')
    out_img = (np.multiply((padded_img/255.0), (target_img/255.0) * -1 + 1) * 255).astype(np.uint8)

    debugimg(out_img, 'multiplied')

    # Calculate contours from masked target image and plot them to the target gerber context
    status_print('Calculating contour lines')
    plot_contours(out_img,
            target,
            offx=(tminx, tminy),
            scale=scale,
            status_print=lambda *args: status_print('   ', *args))

    # Write target gerber context to disk
    status_print('Generating output gerber')
    from gerber.render import rs274x_backend
    ctx = rs274x_backend.Rs274xContext(target.settings)
    target.render(ctx)
    out = ctx.dump().getvalue()
    status_print('Done.')
    return out


def plot_contours(
        img:np.ndarray,
        layer:gerber.rs274x.GerberFile,
        offx:tuple,
        scale:float,
        debug=lambda *args:None,
        status_print=lambda *args:None):
    imgh, imgw = img.shape

    # Extract contours
    status_print('Extracting contours')
    img_cont_out, contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_TC89_KCOS)

    aperture = list(layer.apertures)[0]

    from gerber.primitives import Line, Region
    status_print('offx', offx, 'scale', scale)

    xbias, ybias = offx
    def map(coord):
        x, y = coord
        return (x/scale + xbias, y/scale + ybias)
    def contour_lines(c):
        return [ Line(map(start), map(end), aperture, level_polarity='dark', units=layer.settings.units)
            for start, end in zip(c, np.vstack((c[1:], c[:1]))) ]

    done = []
    process_stack = [-1]
    next_process_stack = []
    parents = [ (i, first_child != -1, parent) for i, (_1, _2, first_child, parent) in enumerate(hierarchy[0]) ]
    is_dark = True
    status_print('Converting contours to gerber primitives')
    with tqdm.tqdm(total=len(contours)) as progress:
        while len(done) != len(contours):
            for i, has_children, parent in parents[:]:
                if parent in process_stack:
                    contour = contours[i]
                    polarity = 'dark' if is_dark else 'clear'
                    debug('rendering {} with parent {} as {} with {} vertices'.format(i, parent, polarity, len(contour)))
                    debug('process_stack is', process_stack)
                    debug()
                    layer.primitives.append(Region(contour_lines(contour[:,0]), level_polarity=polarity, units=layer.settings.units))
                    if has_children:
                        next_process_stack.append(i)
                    done.append(i)
                    parents.remove((i, has_children, parent))
                    progress.update(1)
            debug('skipping to next level')
            process_stack, next_process_stack = next_process_stack, []
            is_dark = not is_dark
    debug('done', done)

# Utility foo
# ===========

def find_gerber_in_dir(dir_path, file_or_ext):
    lname = path.join(dir_path, file_or_ext)
    if path.isfile(lname):
        with open(lname, 'r') as f:
            return lname, f.read()

    contents = os.listdir(dir_path)
    exts = file_or_ext.split(',')
    for entry in contents:
        if any(entry.lower().endswith(ext.lower()) for ext in exts):
            lname = path.join(dir_path, entry)
            if not path.isfile(lname):
                continue
            with open(lname, 'r') as f:
                return lname, f.read()

    raise ValueError('Cannot find file or suffix "{}" in dir {}'.format(file_or_ext, dir_path))

def find_gerber_in_zip(zip_path, file_or_ext):
    with zipfile.ZeipFile(zip_path, 'r') as lezip:
        nlist = [ item.filename for item in zipin.infolist() ]
        if file_or_ext in nlist:
            return file_or_ext, lezip.read(file_or_ext)

        exts = file_or_ext.split(',')
        for n in nlist:
            if any(n.lower().endswith(ext.lower()) for ext in exts):
                return n, lezip.read(n)

    raise ValueError('Cannot find file or suffix "{}" in zip {}'.format(file_or_ext, dir_path))

def replace_file_in_zip(zip_path, filename, contents):
    with tempfile.TemporaryDirectory() as tmpdir:
        tempname = path.join(tmpdir, 'out.zip')
        with zipfile.ZipFile(zip_path, 'r') as zipin, zipfile.ZipFile(tempname, 'w') as zipout:
            for item in zipin.infolist():
                if item.filename != filename:
                    zipout.writestr(item, zipin.read(item.filename))
            zipout.writestr(filename, contents)
        shutil.move(tempname, zip_path)

def paste_image_file(zip_or_dir, target, outline, source_img, subtract=[], status_print=lambda *args:None, debugdir=None):
    if path.isdir(zip_or_dir):
        tname, target = find_gerber_in_dir(zip_or_dir, target)
        status_print('Target layer file {}'.format(os.path.basename(tname)))
        oname, outline  = find_gerber_in_dir(zip_or_dir, outline)
        status_print('Outline layer file {}'.format(os.path.basename(oname)))
        subtract = [ (fn, layer) for fn, layer in (find_gerber_in_dir(zip_or_dir, elem) for elem in subtract) ]
        
        out = paste_image(target, outline, source_img, subtract, debugdir=debugdir, status_print=status_print)

        if not tname.endswith('.bak'):
            shutil.copy(tname, tname+'.bak')
            with open(tname, 'w') as f:
                f.write(out)
        else:
            with open(tname[:-4], 'w') as f:
                f.write(out)
    elif zipfile.is_zipfile(zip_or_dir):
        _fn, outline  = find_gerber_in_zip(zip_or_dir, outline)
        subtract = [ (fn, layer) for fn, layer in (find_gerber_in_zip(zip_or_dir, elem) for elem in subtract) ]
        
        out = paste_image(target, outline, source_img, subtract, debugdir=debugdir, status_print=status_print)
        replace_file_in_zip(zip_or_dir, tname, out)
    else:
        raise ValueError('{} does not look like either a folder or a zip file')

# Command line interface
# ======================

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-b', '--bottom', action='store_true', help='Default to bottom layer file names')
    parser.add_argument('-t', '--target', help='Target layer. Filename or extension in target folder/zip')
    parser.add_argument('-s', '--subtract', nargs='*', help='Layer to subtract. Filename or extension in target folder/zip')
    parser.add_argument('-o', '--outline', default='.GKO,.GM1', help='Target outline layer. Filename or extension in target folder/zip')
    parser.add_argument('-d', '--debug', type=str, help='Directory to place debug files into')
    parser.add_argument('zip_or_dir', default='.', nargs='?', help='Optional folder or zip with target files')
    parser.add_argument('source', help='Source image')
    args = parser.parse_args()

    if not args.target:
        args.target   = '.GBO.bak,.GBO' if args.bottom else '.GTO.bak,.GTO'
    if not args.subtract:
        args.subtract = ['.GBS', '.TXT'] if args.bottom else ['.GTS', '.TXT']

    source_img = cv2.imread(args.source, cv2.IMREAD_GRAYSCALE)
    paste_image_file(
            args.zip_or_dir,
            args.target,
            args.outline,
            source_img,
            args.subtract,
            status_print=lambda *args: print(*args, flush=True),
            debugdir=args.debug)