#!/usr/bin/env python3

import os
import sys
import time
from os import path
from textwrap import dedent
import pkgutil
import subprocess
import xml.etree.ElementTree as xe

import ezdxf


__version__ = '0.1'

PIN_TS_BASE = 0x23420000
TEDIT_BASE = 0x23430000
PATH_BASE = 0x23440000


def sch_template(name, num_pins, yspace=200):
    templ = f'''
        EESchema Schematic File Version 5
        EELAYER 30 0
        EELAYER END
        $Descr A3 16535 11693
        encoding utf-8
        Sheet 1 1
        Title "{name}"
        Date "{time.strftime("%d %b %Y")}"
        Rev ""
        Comp ""
        Comment1 ""
        Comment2 ""
        Comment3 ""
        Comment4 ""
        Comment5 ""
        Comment6 ""
        Comment7 ""
        Comment8 ""
        Comment9 ""
        $EndDescr
        {{components}}
        $EndSCHEMATC
    '''

    components = []
    for i in range(num_pins):
        identifier = f'TP{i}'
        value = 'pogopin'
        x, y = 1000, 1000 + i*yspace
        components.append(dedent(f'''
            $Comp
            L Connector:Conn_01x01_Female {identifier}
            U 1 1 {PIN_TS_BASE + i:08X}
            P {x} {y}
            F 0 "{identifier}" H {x-50} {y+50} 50  0000 R CNN
            F 1 "{value}" H {x+50} {y} 50  0000 L CNN
            F 2 "Pogopin:AutogeneratedPogopinFootprint" H {x} {y} 50  0001 C CNN
            F 3 "~" H {x} {y} 50  0001 C CNN
                    1    {x} {y}
                    -1   0    0    1   
            $EndComp
            ''').strip())

    return dedent(templ).lstrip().format(components='\n'.join(components))

def pcb_template(outline, pins, annular=0.5):
    pcb_templ = f'''
        (kicad_pcb (version 20190605) (host pogojig "({__version__})")

          (general
            (thickness 1.6)
            (drawings {len(pins)})
            (tracks 0)
            (modules {len(pins)})
            (nets {len(pins)+1})
          )

          (page "A4")
          (layers
            (0 "F.Cu" signal)
            (31 "B.Cu" signal)
            (32 "B.Adhes" user)
            (33 "F.Adhes" user)
            (34 "B.Paste" user)
            (35 "F.Paste" user)
            (36 "B.SilkS" user)
            (37 "F.SilkS" user)
            (38 "B.Mask" user)
            (39 "F.Mask" user)
            (40 "Dwgs.User" user)
            (41 "Cmts.User" user)
            (42 "Eco1.User" user)
            (43 "Eco2.User" user)
            (44 "Edge.Cuts" user)
            (45 "Margin" user)
            (46 "B.CrtYd" user)
            (47 "F.CrtYd" user)
            (48 "B.Fab" user)
            (49 "F.Fab" user)
          )

          (setup
            (last_trace_width 0.25)
            (trace_clearance 0.2)
            (zone_clearance 0.508)
            (zone_45_only no)
            (trace_min 0.2)
            (via_size 0.8)
            (via_drill 0.4)
            (via_min_size 0.4)
            (via_min_drill 0.3)
            (uvia_size 0.3)
            (uvia_drill 0.1)
            (uvias_allowed no)
            (uvia_min_size 0.2)
            (uvia_min_drill 0.1)
            (max_error 0.005)
            (defaults
              (edge_clearance 0.01)
              (edge_cuts_line_width 0.05)
              (courtyard_line_width 0.05)
              (copper_line_width 0.2)
              (copper_text_dims (size 1.5 1.5) (thickness 0.3) keep_upright)
              (silk_line_width 0.12)
              (silk_text_dims (size 1 1) (thickness 0.15) keep_upright)
              (other_layers_line_width 0.1)
              (other_layers_text_dims (size 1 1) (thickness 0.15) keep_upright)
            )
            (pad_size 3.14159 3.14159)
            (pad_drill 1.41421)
            (pad_to_mask_clearance 0.051)
            (solder_mask_min_width 0.25)
            (aux_axis_origin 0 0)
            (visible_elements FFFFFF7F)
            (pcbplotparams
              (layerselection 0x010fc_ffffffff)
              (usegerberextensions false)
              (usegerberattributes false)
              (usegerberadvancedattributes false)
              (creategerberjobfile false)
              (excludeedgelayer true)
              (linewidth 0.100000)
              (plotframeref false)
              (viasonmask false)
              (mode 1)
              (useauxorigin false)
              (hpglpennumber 1)
              (hpglpenspeed 20)
              (hpglpendiameter 15.000000)
              (psnegative false)
              (psa4output false)
              (plotreference true)
              (plotvalue true)
              (plotinvisibletext false)
              (padsonsilk false)
              (subtractmaskfromsilk false)
              (outputformat 1)
              (mirror false)
              (drillshape 1)
              (scaleselection 1)
              (outputdirectory ""))
          )

          (net 0 "")
          {{net_defs}}

          (net_class "Default" "This is the default net class."
            (clearance 0.2)
            (trace_width 0.25)
            (via_dia 0.8)
            (via_drill 0.4)
            (uvia_dia 0.3)
            (uvia_drill 0.1)
            {{net_class_defs}}
          )
          
          {{module_defs}}

          {{edge_cuts}}
        )'''

    module_defs = []
    for i, pin in enumerate(pins):
        (x, y), hole_dia = pin # all dimensions in mm here
        pad_dia = hole_dia + 2*annular
        mod = f'''
          (module "Pogopin:AutogeneratedPogopinFootprint" (layer "F.Cu") (tedit {TEDIT_BASE + i:08X}) (tstamp {PIN_TS_BASE + i:08X})
            (at {x} {y})
            (descr "Pogo pin {i}")
            (tags "test point plated hole")
            (path "/{PATH_BASE + i:08X}")
            (attr virtual)
            (fp_text reference "TP{i}" (at 0 -{pad_dia/2 + 1}) (layer "F.SilkS")
              (effects (font (size 1 1) (thickness 0.15)))
            )
            (fp_text value "pogo pin {i}" (at 0 {pad_dia/2 + 1}) (layer "F.Fab")
              (effects (font (size 1 1) (thickness 0.15)))
            )
            (fp_text user "%R" (at 0 -{pad_dia/2 + 1}) (layer "F.Fab")
              (effects (font (size 1 1) (thickness 0.15)))
            )
            (fp_circle (center 0 0) (end {pad_dia} 0) (layer "F.CrtYd") (width 0.05))
            (fp_circle (center 0 0) (end 0 -{pad_dia}) (layer "F.SilkS") (width 0.12))
            (pad "1" thru_hole circle (at 0 0) (size {pad_dia} {pad_dia}) (drill {hole_dia}) (layers *.Cu *.Mask)
              (net {i+1} "pogo{i}"))
          )'''
        module_defs.append(mod)

    edge_cuts = [ f'(gr_line (start {x1} {y1}) (end {x2} {y2}) (layer "Edge.Cuts") (width 0.05))'
                  for (x1, y1), (x2, y2) in outline ]

    net_defs = [ f'(net {i+1} "pogo{i}")' for i, _pin in enumerate(pins) ]
    net_class_defs = [ f'(add_net "pogo{i}")' for i, _pin in enumerate(pins) ]
    return pcb_templ.format(
            net_defs='\n'.join(net_defs),
            net_class_defs='\n'.join(net_class_defs),
            module_defs='\n'.join(module_defs),
            edge_cuts='\n'.join(edge_cuts))

def inkscape_query_all(filename):
    proc = subprocess.run([ os.environ.get('INKSCAPE', 'inkscape'), filename, '--query-all'], capture_output=True)
    proc.check_returncode()
    data = [ line.split(',') for line in proc.stdout.decode().splitlines() ]
    return { id: (float(x), float(y), float(w), float(h)) for id, x, y, w, h in data }

SVG_NS = {
        'svg': 'http://www.w3.org/2000/svg',
        'inkscape': 'http://www.inkscape.org/namespaces/inkscape',
        'sodipodi': 'http://sodipodi.sourceforge.net/DTD/sodipodi-0.dtd'
}

def svg_find_elements(doc, tag, layer=None):
    for i, g in enumerate(doc.findall('svg:g', SVG_NS)):
        if g.attrib.get(f'{{{SVG_NS["inkscape"]}}}groupmode') != 'layer':
            continue

        label = g.attrib.get(f'{{{SVG_NS["inkscape"]}}}label', '')
        if not layer or label == layer:
            yield from g.iter(tag)

# def svg_get_scale(doc):
#     w = doc.attrib['width']
#     h = doc.attrib['height']
# 
#     if not w.endswith('mm') and h.endswith('mm'):
#         raise ValueError('Document dimensions in SVG must be set to millimeters')
# 
#     w, h = float(w[:-2]), float(h[:-2])
#     _x, _y, vb_w, vb_h = map(float, doc.attrib['viewBox'].split())
#     scale_x, scale_y = vb_w / w, vb_h / h
#     assert abs(1 - scale_x/scale_y) < 0.001
#     return scale_x

def svg_get_viewbox_mm(doc):
    w = doc.attrib['width']
    h = doc.attrib['height']

    if not w.endswith('mm') and h.endswith('mm'):
        raise ValueError('Document dimensions in SVG must be set to millimeters')

    w, h = float(w[:-2]), float(h[:-2])
    x, y, vb_w, vb_h = map(float, doc.attrib['viewBox'].split())
    scale_x, scale_y = vb_w / w, vb_h / h
    return x/scale_x, y/scale_y, w, h

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('svg', metavar='pogo_map.svg', help='Input inkscape SVG pogo pin map (use provided template!)')
    parser.add_argument('outline', metavar='outline.dxf', help='Board outline DXF generated by OpenSCAD')
    parser.add_argument('output', default='kicad', help='Output directory/project name and path')
    parser.add_argument('-y', '--yspace', type=int, default=200, help='Schematic pin Y spacing in mil (default: 200)')
    parser.add_argument('-a', '--annular', type=float, default=0.5, help='Pogo pin annular ring width in mm (default: 0.5)')
    parser.add_argument('-l', '--svg-layer', type=str, default='Test Points', help='Name of SVG layer containing pogo pins')
    parser.add_argument('-n', '--name', default='jig', help='Output KiCAD project name')
    args = parser.parse_args()

    if not path.exists(args.output):
        os.mkdir(args.output)

    if not path.isdir(args.output):
        raise SystemError(f'Output path "{args.output}" is not a directory')

    with open(args.svg, 'r') as f:
        doc = xe.fromstring(f.read())
    pogo_circle_ids = [ circle.attrib['id'] for circle in svg_find_elements(doc, f'{{{SVG_NS["svg"]}}}circle', args.svg_layer) ]
    # scale = svg_get_scale(doc)
    page_x, page_y, page_w, page_h = svg_get_viewbox_mm(doc)
    MM_PER_IN = 25.4
    SVG_DEF_DPI = 96
    px_to_mm = lambda px: px/SVG_DEF_DPI * MM_PER_IN
    query = inkscape_query_all(args.svg)
    dims = [ query[id] for id in pogo_circle_ids ]
    assert all( abs(1 - w/h) < 0.001 for _x, _y, w, h in dims )
    print('origin:', page_x, page_y)
    print('dims:', page_w, page_h)
    pins = [ (
                (page_x + px_to_mm(x) + px_to_mm(w)/2,
                 page_y - page_h + px_to_mm(y) + px_to_mm(w)/2),
                px_to_mm(w)) for x, y, w, h in dims ]

    doc = ezdxf.readfile(args.outline)
    outline = []
    for line in doc.modelspace().query('LINE'):
        (x1, y1, _z1), (x2, y2, _z2) = line.dxf.start, line.dxf.end
        outline.append(((x1, -y1), (x2, -y2)))

    with open(path.join(args.output, f'{args.name}.sch'), 'w', encoding='utf8') as sch:
        sch.write(sch_template(f'{args.name} generated schematic (PogoJig v{__version__})', len(pins), yspace=args.yspace))

    with open(path.join(args.output, f'{args.name}.kicad_pcb'), 'w', encoding='utf8') as pcb:
        pcb.write(pcb_template(outline, pins, annular=args.annular))

    with open(path.join(args.output, f'{args.name}.pro'), 'w', encoding='utf8') as f:
        f.write(pkgutil.get_data('pogojig.kicad', 'kicad.pro').decode('utf8'))

    with open(path.join(args.output, f'{args.name}-cache.lib'), 'w', encoding='utf8') as f:
        f.write(pkgutil.get_data('pogojig.kicad', 'kicad-cache.lib').decode('utf8'))