#!/usr/bin/env python3
"""build_ekey.py — generate the M.2 E-key (Wi-Fi) sub-sheet for the CM5 carrier.

The inverse of retire_block.py: places a library connector + passives, attaches a
named label at every USED pin connection point, and a (no_connect) at every unused
pin. Connectivity is by KiCad's by-name net model (a label coincident with a pin
endpoint joins that pin to the net) — no wire routing, which keeps generation exact.

Locked net map (see CM5_Carrier_Pinout_BOM.md §4, design confirmed E-key-only):
  PCIe link-0 -> the 9 freed CM5 pins ; +3V3_RF (4 pins) ; GND (11) ; W_DISABLE x2 pull-ups.
  RX is NOT AC-coupled on the carrier (caps live on the card's TX, per M.2 spec).

Emits carrier/M2_Ekey.kicad_sch. Pin connection point: abs=(px+lx, py-ly) for a
rot0/no-mirror placement (the transform validated in retire_block.py).
"""
import re, sys, os, uuid as uuidlib

ROOT_UUID = "e63e39d7-6ac0-4ffd-8aa3-1841a4541b55"
PROJECT   = "CM5IO"
SHEET_SYMBOL_UUID = "00000000-0000-0000-0000-00005 e000ekey".replace(" ", "")  # placeholder, overwritten below
KSS = "/Applications/KiCad/KiCad.app/Contents/SharedSupport/symbols"
HERE = os.path.dirname(os.path.abspath(__file__))
CARRIER = os.path.join(HERE, "..", "carrier")

def U(label):
    return str(uuidlib.uuid5(uuidlib.NAMESPACE_DNS, "cm5carrier-ekey-" + label))

# deterministic sheet-symbol uuid (used in instance paths + the root sheet block)
SHEET_SYMBOL_UUID = U("sheet-symbol")
SHEET_FILE_UUID   = U("sheet-file")

def paren_end(t, s):
    d = 0
    for i in range(s, len(t)):
        if t[i] == '(': d += 1
        elif t[i] == ')':
            d -= 1
            if d == 0: return i + 1
    raise ValueError("unbalanced")

def get_symbol(text, name):
    m = re.search(r'\(symbol\s+"' + re.escape(name) + r'"', text)
    if not m: return None
    s = m.start(); return text[s:paren_end(text, s)]

def lib_pins(seg):
    """[(number, lx, ly, angle)] for a lib_symbol block (descends into _x_1 units)."""
    out = []
    for pm in re.finditer(r'\(pin\b', seg):
        s = pm.start(); ps = seg[s:paren_end(seg, s)]
        at = re.search(r'\(at\s+([\d.-]+)\s+([\d.-]+)\s+([\d.-]+)\)', ps)
        nu = re.search(r'\(number\s+"([^"]+)"', ps)
        if at and nu:
            out.append((nu.group(1), float(at.group(1)), float(at.group(2)), float(at.group(3))))
    return out

# ---- gather lib_symbol definitions -------------------------------------------------
carrier_hs = open(os.path.join(CARRIER, "CM5_HighSpeed.kicad_sch")).read()
conn_lib   = open(os.path.join(KSS, "Connector.kicad_sym")).read()
power_lib  = open(os.path.join(KSS, "power.kicad_sym")).read()

# reuse R / C / GND verbatim from the carrier (guaranteed-compatible embeds)
DEV_R  = get_symbol(carrier_hs, "Device:R")
DEV_C  = get_symbol(carrier_hs, "Device:C")
PWR_GND = get_symbol(carrier_hs, "power:GND")
# socket + PWR_FLAG from KiCad libs, renamed to their full lib_id
sock = get_symbol(conn_lib, "Bus_M.2_Socket_E")
sock = sock.replace('(symbol "Bus_M.2_Socket_E"', '(symbol "Connector:Bus_M.2_Socket_E"', 1)

def set_pin_type(seg, number, newtype):
    """Rewrite the electrical type of the lib pin whose (number "N") matches.
    The socket is typed from the host's view; for ERC across the connector we model
    the lines CM5 drives as 'passive' and the RX pair (card drives -> CM5) as 'output'."""
    for pm in re.finditer(r'\(pin\s+(\w+)\s+(\w+)', seg):
        s = pm.start(); blk = seg[s:paren_end(seg, s)]
        if re.search(r'\(number\s+"' + re.escape(number) + r'"', blk):
            head = re.match(r'\(pin\s+\w+', blk).group(0)
            return seg[:s] + blk.replace(head, f'(pin {newtype}', 1) + seg[s + len(blk):]
    return seg
# Lines CM5 drives (TX 35/37, REFCLK 47/49, PERST 52): socket side passive.
for n in ['35', '37', '47', '49', '52']:
    sock = set_pin_type(sock, n, 'passive')
# Lines the CARD drives into CM5 inputs -> model socket side as output:
#   RX pair (41/43) + CLKREQ# (53) + WAKE# (55) are all CM5 inputs (open-drain for
#   CLKREQ/WAKE; the card pulls them). 'output' gives each net a driver for ERC.
for n in ['41', '43', '53', '55']:
    sock = set_pin_type(sock, n, 'output')
pflag = get_symbol(power_lib, "PWR_FLAG")
pflag = pflag.replace('(symbol "PWR_FLAG"', '(symbol "power:PWR_FLAG"', 1)
for nm, blk in [("Device:R", DEV_R), ("Device:C", DEV_C), ("power:GND", PWR_GND),
                ("Connector:Bus_M.2_Socket_E", sock), ("power:PWR_FLAG", pflag)]:
    if not blk: sys.exit("missing lib_symbol: " + nm)

SOCK_PINS = {n: (lx, ly, a) for n, lx, ly, a in lib_pins(sock)}   # 67 pins

# ---- the net map -------------------------------------------------------------------
# pin -> ('hier', net) cross-sheet PCIe ; ('gnd',) ; ('v33',) ; ('loc', net) local
HIER = {
    '35': 'PCIE_TX_P', '37': 'PCIE_TX_N', '41': 'PCIE_RX_P', '43': 'PCIE_RX_N',
    '47': 'PCIE_CLK_P', '49': 'PCIE_CLK_N', '52': 'PCIE_nRST', '53': 'PCIE_nCLKREQ',
    '55': 'PCIE_nWAKE',
}
V33  = ['2', '4', '72', '74']
GND  = ['1', '7', '18', '33', '39', '45', '51', '57', '63', '69', '75']
LOCAL = {'56': 'WIFI_DISABLE#', '54': 'WIFI_DISABLE2#'}
used = set(HIER) | set(V33) | set(GND) | set(LOCAL)
NC = sorted((p for p in SOCK_PINS if p not in used), key=lambda x: int(x))

assert len(used) + len(NC) == len(SOCK_PINS), (len(used), len(NC), len(SOCK_PINS))

# ---- placement ---------------------------------------------------------------------
def g(v):  # snap to the 1.27mm (50mil) schematic grid so labels bind to pins
    return round(round(v / 1.27) * 1.27, 2)
PX, PY = g(170.0), g(105.0)
def pin_abs(num):
    lx, ly, a = SOCK_PINS[num]
    return (round(PX + lx, 2), round(PY - ly, 2), a)

def lbl_angle(num):
    lx, ly, a = SOCK_PINS[num]
    if abs(lx) >= abs(ly): return 0 if lx > 0 else 180
    return 90 if ly > 0 else 270

def hier_label(name, x, y, ang, shape="bidirectional"):
    return (f'\t(hierarchical_label "{name}"\n\t\t(shape {shape})\n\t\t(at {x} {y} {ang})\n'
            f'\t\t(effects (font (size 1.27 1.27)) (justify left))\n\t\t(uuid "{U("hl-"+name+"-"+str(x)+"-"+str(y))}")\n\t)\n')

def glob_label(name, x, y, ang):
    return (f'\t(global_label "{name}"\n\t\t(shape input)\n\t\t(at {x} {y} {ang})\n'
            f'\t\t(effects (font (size 1.27 1.27)) (justify left))\n\t\t(uuid "{U("gl-"+name+"-"+str(x)+"-"+str(y))}")\n\t)\n')

def loc_label(name, x, y, ang):
    return (f'\t(label "{name}"\n\t\t(at {x} {y} {ang})\n'
            f'\t\t(effects (font (size 1.27 1.27)) (justify left))\n\t\t(uuid "{U("ll-"+name+"-"+str(x)+"-"+str(y))}")\n\t)\n')

def no_connect(x, y, tag):
    tag = tag if tag[-1].isdigit() else tag + "1"  # refs must end in a digit (KiCad annotation)
    return f'\t(no_connect (at {x} {y}) (uuid "{U("nc-"+tag)}"))\n'

def gnd_sym(x, y, tag):
    tag = tag if tag[-1].isdigit() else tag + "1"  # refs must end in a digit (KiCad annotation)
    return (f'\t(symbol\n\t\t(lib_id "power:GND")\n\t\t(at {x} {y} 0)\n\t\t(unit 1)\n'
            f'\t\t(exclude_from_sim no) (in_bom yes) (on_board yes) (dnp no)\n'
            f'\t\t(uuid "{U("gnd-"+tag)}")\n'
            f'\t\t(property "Reference" "#PWR{tag}" (at {x} {y+3.0} 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
            f'\t\t(property "Value" "GND" (at {x} {y+2.0} 0) (effects (font (size 1.27 1.27))))\n'
            f'\t\t(property "Footprint" "" (at {x} {y} 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
            f'\t\t(property "Datasheet" "" (at {x} {y} 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
            f'\t\t(pin "1" (uuid "{U("gndpin-"+tag)}"))\n'
            f'\t\t(instances (project "{PROJECT}" (path "/{ROOT_UUID}/{SHEET_SYMBOL_UUID}" (reference "#PWR{tag}") (unit 1))))\n\t)\n')

def pwr_flag(x, y, tag):
    tag = tag if tag[-1].isdigit() else tag + "1"  # refs must end in a digit (KiCad annotation)
    return (f'\t(symbol\n\t\t(lib_id "power:PWR_FLAG")\n\t\t(at {x} {y} 0)\n\t\t(unit 1)\n'
            f'\t\t(exclude_from_sim no) (in_bom yes) (on_board yes) (dnp no)\n'
            f'\t\t(uuid "{U("pflag-"+tag)}")\n'
            f'\t\t(property "Reference" "#FLG{tag}" (at {x} {y-3.0} 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
            f'\t\t(property "Value" "PWR_FLAG" (at {x} {y-2.0} 0) (effects (font (size 1.27 1.27))))\n'
            f'\t\t(property "Footprint" "" (at {x} {y} 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
            f'\t\t(property "Datasheet" "" (at {x} {y} 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
            f'\t\t(pin "1" (uuid "{U("pflagpin-"+tag)}"))\n'
            f'\t\t(instances (project "{PROJECT}" (path "/{ROOT_UUID}/{SHEET_SYMBOL_UUID}" (reference "#FLG{tag}") (unit 1))))\n\t)\n')

def passive(lib_id, ref, value, footprint, x, y, tag):
    return (f'\t(symbol\n\t\t(lib_id "{lib_id}")\n\t\t(at {x} {y} 0)\n\t\t(unit 1)\n'
            f'\t\t(exclude_from_sim no) (in_bom yes) (on_board yes) (dnp no)\n'
            f'\t\t(uuid "{U("sym-"+ref)}")\n'
            f'\t\t(property "Reference" "{ref}" (at {x+2.54} {y-1.0} 0) (effects (font (size 1.27 1.27)) (justify left)))\n'
            f'\t\t(property "Value" "{value}" (at {x+2.54} {y+1.0} 0) (effects (font (size 1.27 1.27)) (justify left)))\n'
            f'\t\t(property "Footprint" "{footprint}" (at {x} {y} 90) (effects (font (size 1.27 1.27)) (hide yes)))\n'
            f'\t\t(property "Datasheet" "" (at {x} {y} 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
            f'\t\t(pin "1" (uuid "{U("p1-"+ref)}"))\n'
            f'\t\t(pin "2" (uuid "{U("p2-"+ref)}"))\n'
            f'\t\t(instances (project "{PROJECT}" (path "/{ROOT_UUID}/{SHEET_SYMBOL_UUID}" (reference "{ref}") (unit 1))))\n\t)\n')

def socket_symbol():
    # the placed J_E (multi-unit symbol uses unit 1; both graphic units share lib def)
    props = ('\t\t(property "Reference" "J_E1" (at 170 53 0) (effects (font (size 1.27 1.27))))\n'
             '\t\t(property "Value" "M.2_Ekey_WiFi" (at 170 55 0) (effects (font (size 1.27 1.27))))\n'
             '\t\t(property "Footprint" "Connector_PCBEdge:M.2_2230-xx-E" (at 170 105 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
             '\t\t(property "Datasheet" "" (at 170 105 0) (effects (font (size 1.27 1.27)) (hide yes)))\n'
             '\t\t(property "Description" "M.2 A+E key socket, Wi-Fi (AsiaRF AW7915/MT7915)" (at 170 105 0) (effects (font (size 1.27 1.27)) (hide yes)))\n')
    pins = "".join(f'\t\t(pin "{n}" (uuid "{U("jepin-"+n)}"))\n' for n in sorted(SOCK_PINS, key=int))
    return (f'\t(symbol\n\t\t(lib_id "Connector:Bus_M.2_Socket_E")\n\t\t(at {PX} {PY} 0)\n\t\t(unit 1)\n'
            f'\t\t(exclude_from_sim no) (in_bom yes) (on_board yes) (dnp no)\n'
            f'\t\t(uuid "{U("J_E1")}")\n' + props + pins +
            f'\t\t(instances (project "{PROJECT}" (path "/{ROOT_UUID}/{SHEET_SYMBOL_UUID}" (reference "J_E1") (unit 1))))\n\t)\n')

# ---- assemble the sheet ------------------------------------------------------------
body = []
body.append(socket_symbol())

# connector pin labels / NC / GND
for num, net in HIER.items():
    # global labels: connect to the matching CM5-side labels across sheets without
    # parent sheet-pin plumbing (the 9 freed CM5 PCIe pins get the same global names).
    x, y, _ = pin_abs(num); body.append(glob_label(net, x, y, lbl_angle(num)))
for num in V33:
    x, y, _ = pin_abs(num); body.append(glob_label("+3V3_RF", x, y, lbl_angle(num)))
for num, net in LOCAL.items():
    x, y, _ = pin_abs(num); body.append(loc_label(net, x, y, lbl_angle(num)))
for num in GND:
    x, y, _ = pin_abs(num); body.append(gnd_sym(x, y, num))
for num in NC:
    x, y, _ = pin_abs(num); body.append(no_connect(x, y, num))

# pull-up resistors: pin1(top) -> signal label ; pin2(bottom) -> +3V3_RF
def place_rc(lib_id, ref, value, fp, x, y, top_label_fn, top_name, bot_name):
    body.append(passive(lib_id, ref, value, fp, x, y, ref))
    body.append(top_label_fn(top_name, x, round(y - 3.81, 2), 90))          # pin1
    body.append(glob_label(bot_name, x, round(y + 3.81, 2), 270))           # pin2

place_rc("Device:R", "R_E1", "10k", "Resistor_SMD:R_0402_1005Metric", g(215.0), g(80.0), loc_label, "WIFI_DISABLE#", "+3V3_RF")
place_rc("Device:R", "R_E2", "10k", "Resistor_SMD:R_0402_1005Metric", g(225.0), g(80.0), loc_label, "WIFI_DISABLE2#", "+3V3_RF")

# decoupling caps: pin1 -> +3V3_RF ; pin2 -> GND
for i, (ref, val) in enumerate([("C_E1", "10u"), ("C_E2", "100n"), ("C_E3", "100n")]):
    cx = g(215.0 + i * 7.62); cy = g(110.0)
    body.append(passive("Device:C", ref, val, "Capacitor_SMD:C_0402_1005Metric", cx, cy, ref))
    body.append(glob_label("+3V3_RF", cx, round(cy - 3.81, 2), 90))         # pin1
    body.append(gnd_sym(cx, round(cy + 3.81, 2), ref + "g"))               # pin2 -> GND symbol

# +3V3_RF is driven by the Power sheet's U_BUCK33R (PWR_FLAG lives there).

# ---- file scaffold -----------------------------------------------------------------
libsyms = "\n".join([DEV_R, DEV_C, PWR_GND, pflag, sock])
sheet = (
    '(kicad_sch\n'
    '\t(version 20231120)\n\t(generator "eeschema")\n\t(generator_version "8.0")\n'
    f'\t(uuid "{SHEET_FILE_UUID}")\n\t(paper "A4")\n'
    '\t(title_block\n\t\t(title "CM5 Carrier - M.2 E-key (Wi-Fi)")\n\t\t(rev "A")\n'
    '\t\t(comment 1 "AW7915 / MT7915, PCIe x1 Gen2, USB NC")\n\t)\n'
    '\t(lib_symbols\n' + libsyms + '\n\t)\n'
    + "".join(body) +
    '\t(sheet_instances\n\t\t(path "/"\n\t\t\t(page "4")\n\t\t)\n\t)\n'
    ')\n'
)

assert sheet.count('(') == sheet.count(')'), f"UNBALANCED ({sheet.count('(')} vs {sheet.count(')')})"
out = os.path.join(CARRIER, "M2_Ekey.kicad_sch")
print(f"pins: {len(SOCK_PINS)}  used: {len(used)} (hier {len(HIER)}, v33 {len(V33)}, gnd {len(GND)}, local {len(LOCAL)})  NC: {len(NC)}")
print(f"sheet-symbol uuid: {SHEET_SYMBOL_UUID}")
print(f"parens balanced: {sheet.count('(')}")
if '--apply' in sys.argv:
    open(out, 'w').write(sheet); print("WROTE " + out)
else:
    print("DRY RUN (pass --apply to write " + out + ")")