"""Generate the KiCad schematic from design.py (label-style: every pin gets a stub + net label)."""
import copy
import os
import sys

sys.path.insert(0, os.path.dirname(__file__))
import design as D
from board import uid, fmt
from kilib import load_symbol, sym_pins
from sexpr import QStr, dumps, find, findall

OUT = os.path.expanduser('~/Documents/projects/sandTable-local/sandtable-controller')
NAME = 'sandtable-controller'
ROOT = uid('rootsheet')
STUB = 2.54
LIBNICK = 'SandTable'

# custom symbols kept in the project library
CUSTOM = {
    'Driver_Motor:Pololu_Breakout_A4988': ('DRV8825_StepStick', {
        '2': ('VDD/~{FAULT}', 'passive'), '3': ('A2', None), '4': ('A1', None), '5': ('B1', None), '6': ('B2', None),
        '9': ('~{EN}', None), '10': ('M0', None), '11': ('M1', None), '12': ('M2', None),
        '13': ('~{RST}', None), '14': ('~{SLP}', None)}),
    'Interface_USB:CH340C': ('CH340C_3V3', {'4': ('V3', 'passive')}),
}

GROUPS = [
    ('POWER INPUT  12V -> 5V', (12.7, 25.4, 287.02, 104.14), ['J1', 'F1', 'D1', 'C3', 'PS1', 'C4', 'D2', '#FLG1', '#FLG2', '#FLG3']),
    ('3.3V', (294.64, 25.4, 398.78, 104.14), ['U3', 'C5', 'C6']),
    ('SENSORS / LED / EXPANSION', (406.4, 25.4, 579.12, 238.76), ['J5', 'R9', 'R10', 'C13', 'D6', 'J6', 'R11', 'R12', 'C14', 'D7', 'U4', 'C15', 'R13', 'J7', 'J8']),
    ('USB-C + USB-SERIAL + AUTO-RESET', (12.7, 116.84, 190.5, 264.16), ['J2', 'U5', 'R1', 'R2', 'D3', 'U2', 'C7', 'Q1', 'R3', 'Q2', 'R4', '#FLG4']),
    ('ESP32', (198.12, 116.84, 398.78, 264.16), ['U1', 'C9', 'C10', 'R5', 'C8', 'SW1', 'R6', 'SW2', 'R8', 'R14', 'D4', 'R15', 'D5']),
    ('STEPPER DRIVERS (DRV8825, 1/32 step: M0-M2, RST, SLP tied high; JP1 bridged only for A4988)', (12.7, 276.86, 398.78, 406.4), ['A1', 'C1', 'C16', 'J3', 'A2', 'C2', 'C17', 'J4', 'JP1']),
    ('MOUNTING (Arduino Uno pattern)', (406.4, 251.46, 579.12, 299.72), ['H1', 'H2', 'H3', 'H4']),
]
FLAGS = {'#FLG1': 'GND', '#FLG2': '+12V', '#FLG3': '+5V', '#FLG4': 'VBUS'}


def snap(v, g=2.54):
    return round(v / g) * g


def symbol_def(lib_id):
    """(lib_id used in the schematic, symbol definition list)."""
    sym = copy.deepcopy(load_symbol(lib_id))
    lib_name = str(sym[1])
    if lib_id in CUSTOM:
        new, renames = CUSTOM[lib_id]
        for unit in findall(sym, 'symbol'):
            unit[1] = QStr(str(unit[1]).replace(lib_name, new, 1))
            for p in findall(unit, 'pin'):
                num = str(find(p, 'number')[1])
                if num in renames:
                    nm, typ = renames[num]
                    find(p, 'name')[1] = QStr(nm)
                    if typ:
                        p[1] = typ
        full = f'{LIBNICK}:{new}'
        sym[1] = QStr(full)
        return full, sym, new
    sym[1] = QStr(lib_id)
    return lib_id, sym, None


def bbox(sym):
    xs, ys = [], []
    for unit in findall(sym, 'symbol'):
        for g in unit:
            if not isinstance(g, list):
                continue
            if g[0] == 'rectangle':
                s, e = find(g, 'start'), find(g, 'end')
                xs += [float(s[1]), float(e[1])]; ys += [float(s[2]), float(e[2])]
            elif g[0] in ('polyline',):
                for xy in findall(find(g, 'pts'), 'xy'):
                    xs.append(float(xy[1])); ys.append(float(xy[2]))
            elif g[0] == 'circle':
                c = find(g, 'center'); r = float(find(g, 'radius')[1])
                xs += [float(c[1]) - r, float(c[1]) + r]; ys += [float(c[2]) - r, float(c[2]) + r]
    for p in sym_pins(sym):
        xs.append(p['x']); ys.append(p['y'])
    if not xs:
        return (-2.54, -2.54, 2.54, 2.54)
    return (min(xs), min(ys), max(xs), max(ys))


def pin_net_map():
    m = {}
    for net, pins in D.NETS.items():
        for p in pins:
            m[p] = net
    return m


def outward(angle):
    """Unit vector (schematic coords, y down) pointing away from the symbol body."""
    a = int(round(angle)) % 360
    return {0: (-1, 0), 180: (1, 0), 90: (0, 1), 270: (0, -1)}[a]


def label_angle(o):
    return {(1, 0): 0, (-1, 0): 180, (0, -1): 90, (0, 1): 270}[o]


def font(size=1.27, hide=False, justify=None, bold=False):
    f = ['font', ['size', fmt(size), fmt(size)]] + ([['bold', 'yes']] if bold else [])
    e = ['effects', f]
    if justify:
        e.append(['justify'] + justify)
    if hide:
        e.append(['hide', 'yes'])
    return e


def main():
    pin_net = pin_net_map()
    nc = set(D.NO_CONNECT)
    items = []          # schematic body items
    lib_syms = {}
    custom_defs = {}
    instances = {}      # ref -> (lib_id, sym, value, footprint, lcsc, descr)
    for ref, (sym_id, fp_id, value, *_rest) in D.PARTS.items():
        lcsc, descr = _rest[-2], _rest[-1]
        full, sym, custom = symbol_def(sym_id)
        lib_syms[full] = sym
        if custom:
            custom_defs[custom] = sym
        instances[ref] = (full, sym, value, fp_id, lcsc, descr)
    flag_full, flag_sym, _ = symbol_def('power:PWR_FLAG')
    lib_syms[flag_full] = flag_sym
    for f in FLAGS:
        instances[f] = (flag_full, flag_sym, 'PWR_FLAG', '', '', '')

    # sanity: every pin of every part is either in a net or no-connect
    missing = []
    for ref, (full, sym, *_r) in instances.items():
        if ref.startswith('#'):
            continue
        for p in sym_pins(sym):
            key = f'{ref}.{p["num"]}'
            if key not in pin_net and key not in nc:
                missing.append(key)
    if missing:
        print('UNASSIGNED PINS:', missing)

    texts = []
    for title, (x0, y0, x1, y1), refs in GROUPS:
        texts.append(['text', QStr(title), ['exclude_from_sim', 'no'], ['at', fmt(x0), fmt(y0 - 3.81), '0'],
                      font(2.0, bold=True, justify=['left', 'bottom']), ['uuid', QStr(uid('grp', title))]])
        texts.append(['rectangle', ['start', fmt(x0 - 2.54), fmt(y0 - 2.54)], ['end', fmt(x1), fmt(y1)],
                      ['stroke', ['width', '0.2'], ['type', 'dash']], ['fill', ['type', 'none']],
                      ['uuid', QStr(uid('grprect', title))]])
        # pack symbols left->right, wrap
        cx, cy, row_h = x0 + 2.54, y0 + 2.54, 0
        for ref in refs:
            full, sym, value, fp_id, lcsc, descr = instances[ref]
            bx0, by0, bx1, by1 = bbox(sym)
            pins = sym_pins(sym)
            lab = max([len(pin_net.get(f'{ref}.{p["num"]}', '') or ('X' if ref.startswith('#') else '')) for p in pins] + [0])
            lab_w = STUB + 1.0 * 1.27 * lab + 1.27
            left = any(outward(p['angle'])[0] < 0 for p in pins)
            right = any(outward(p['angle'])[0] > 0 for p in pins)
            up = any(outward(p['angle'])[1] < 0 for p in pins)
            dn = any(outward(p['angle'])[1] > 0 for p in pins)
            w = (bx1 - bx0) + (lab_w if left else 2.54) + (lab_w if right else 2.54) + 5.08
            h = (by1 - by0) + (lab_w if up else 5.08) + (lab_w if dn else 5.08) + 2.54
            if cx + w > x1 and cx > x0 + 2.54:
                cx = x0 + 2.54
                cy += row_h
                row_h = 0
            # symbol origin so that its bbox sits inside the slot
            ox = snap(cx + (lab_w if left else 2.54) - bx0, 1.27)
            oy = snap(cy + (lab_w if up else 5.08) + by1, 1.27)
            instances[ref] = (full, sym, value, fp_id, lcsc, descr, ox, oy, (by1, by0))
            cx += w
            row_h = max(row_h, h)
        if cy + row_h > y1 + 0.1:
            print(f'  ! group "{title}" overflows its box by {cy + row_h - y1:.1f} mm')

    sym_nodes = []
    flag_count = 0
    for ref, inst in instances.items():
        full, sym, value, fp_id, lcsc, descr, ox, oy, (top, bot) = inst
        is_flag = ref.startswith('#')
        node = ['symbol', ['lib_id', QStr(full)], ['at', fmt(ox), fmt(oy), '0'], ['unit', '1'],
                ['exclude_from_sim', 'no'], ['in_bom', 'no' if (is_flag or ref.startswith(('H', 'JP'))) else 'yes'],
                ['on_board', 'no' if is_flag else 'yes'], ['dnp', 'no'],
                ['uuid', QStr(uid('sym', ref))]]
        props = {str(p[1]): p for p in findall(sym, 'property')}
        ref_y = oy - top - 1.27
        val_y = oy - bot + 2.54
        node.append(['property', QStr('Reference'), QStr(ref), ['at', fmt(ox), fmt(ref_y), '0'],
                     font(hide=is_flag)])
        node.append(['property', QStr('Value'), QStr(value), ['at', fmt(ox), fmt(val_y), '0'], font(hide=is_flag)])
        node.append(['property', QStr('Footprint'), QStr(fp_id), ['at', fmt(ox), fmt(oy), '0'], font(hide=True)])
        ds = props.get('Datasheet')
        node.append(['property', QStr('Datasheet'), QStr(str(ds[2]) if ds else ''), ['at', fmt(ox), fmt(oy), '0'], font(hide=True)])
        node.append(['property', QStr('Description'), QStr(descr), ['at', fmt(ox), fmt(oy), '0'], font(hide=True)])
        if lcsc:
            node.append(['property', QStr('LCSC'), QStr(lcsc), ['at', fmt(ox), fmt(oy), '0'], font(hide=True)])
        pins = sym_pins(sym)
        seen = set()
        for p in pins:
            if p['num'] in seen:
                continue
            seen.add(p['num'])
            node.append(['pin', QStr(p['num']), ['uuid', QStr(uid('pin', ref, p['num']))]])
        node.append(['instances', ['project', QStr(NAME), ['path', QStr('/' + ROOT), ['reference', QStr(ref)], ['unit', '1']]]])
        sym_nodes.append(node)
        # wires + labels / no-connects
        done = set()
        for p in pins:
            key = f'{ref}.{p["num"]}'
            px, py = ox + p['x'], oy - p['y']
            if (round(px, 3), round(py, 3)) in done:
                continue
            done.add((round(px, 3), round(py, 3)))
            net = FLAGS[ref] if is_flag else pin_net.get(key)
            if net is None:
                if key in nc or p['type'] == 'no_connect':
                    items.append(['no_connect', ['at', fmt(px), fmt(py)], ['uuid', QStr(uid('nc', key, px, py))]])
                continue
            o = outward(p['angle'])
            qx, qy = px + STUB * o[0], py + STUB * o[1]
            items.append(['wire', ['pts', ['xy', fmt(px), fmt(py)], ['xy', fmt(qx), fmt(qy)]],
                          ['stroke', ['width', '0'], ['type', 'default']], ['uuid', QStr(uid('w', key, px, py))]])
            ang = label_angle(o)
            just = ['left', 'bottom'] if ang in (0, 90) else ['right', 'bottom']
            items.append(['label', QStr(net), ['at', fmt(qx), fmt(qy), str(ang)], font(justify=just),
                          ['uuid', QStr(uid('lbl', key, px, py))]])

    root = ['kicad_sch', ['version', '20250114'], ['generator', QStr('eeschema')], ['generator_version', QStr('9.0')],
            ['uuid', QStr(ROOT)], ['paper', QStr('A2')],
            ['title_block', ['title', QStr('SandTable Controller')], ['date', QStr('2026-10-22')], ['rev', QStr(D.VERSION)],
             ['company', QStr('lazypugs')],
             ['comment', '1', QStr('ESP32-WROOM-32E (FluidNC) + 2x DRV8825 StepStick, 12V in')],
             ['comment', '2', QStr('Drop-in for the Arduino+CNC shield bay in the sand table (Uno hole pattern)')]],
            ['lib_symbols'] + list(lib_syms.values())]
    root += texts + items + sym_nodes
    root.append(['sheet_instances', ['path', QStr('/'), ['page', QStr('1')]]])
    root.append(['embedded_fonts', 'no'])
    os.makedirs(OUT, exist_ok=True)
    with open(os.path.join(OUT, NAME + '.kicad_sch'), 'w') as f:
        f.write(dumps(root) + '\n')
    # project symbol library with the custom symbols
    lib = ['kicad_symbol_lib', ['version', '20241209'], ['generator', QStr('sandtable-gen')], ['generator_version', QStr('9.0')]]
    for name, s in custom_defs.items():
        s2 = copy.deepcopy(s)
        s2[1] = QStr(name)
        lib.append(s2)
    with open(os.path.join(OUT, 'sandtable.kicad_sym'), 'w') as f:
        f.write(dumps(lib) + '\n')
    symlibs = sorted({sid.split(':')[0] for sid in [v[0] for v in D.PARTS.values()] + ['power:PWR_FLAG']})
    with open(os.path.join(OUT, 'sym-lib-table'), 'w') as f:
        f.write('(sym_lib_table\n  (version 7)\n')
        f.write('  (lib (name "SandTable")(type "KiCad")(uri "${KIPRJMOD}/sandtable.kicad_sym")(options "")(descr "SandTable custom symbols"))\n')
        for lib in symlibs:
            f.write(f'  (lib (name "{lib}")(type "KiCad")(uri "${{KICAD9_SYMBOL_DIR}}/{lib}.kicad_sym")(options "")(descr ""))\n')
        f.write(')\n')
    fplibs = sorted({v[1].split(':')[0] for v in D.PARTS.values()})
    with open(os.path.join(OUT, 'fp-lib-table'), 'w') as f:
        f.write('(fp_lib_table\n  (version 7)\n')
        for lib in fplibs:
            f.write(f'  (lib (name "{lib}")(type "KiCad")(uri "${{KICAD9_FOOTPRINT_DIR}}/{lib}.pretty")(options "")(descr ""))\n')
        f.write(')\n')
    print('wrote schematic,', len(sym_nodes), 'symbols,', sum(1 for i in items if i[0] == 'label'), 'labels,',
          sum(1 for i in items if i[0] == 'no_connect'), 'no-connects')


if __name__ == '__main__':
    main()