AI Flow

"""project-common: the tool-neutral half of an EDA project workflow (adom-aiflow <tool> <step>).

Each backend's workflow script (tools/<tool>-project.py) is run with this file in front of it, so these
helpers are shared: the run's desktop and thread, one adom-bridge verb, the project state under
<run>/<tool>/, the flow's proof files, wiki component pages, distributor numbers checked against the
exact MPN, and a reference KiCad board re-framed into a board frame (lower-left corner at 0,0, y up).
A script sets TOOL before calling anything that touches state (init(tool)).
"""
import argparse, csv, json, math, os, re, shutil, subprocess, sys, time
from pathlib import Path

RUN = Path(os.environ.get('AIFLOW_RUN', 'aiflow-run'))
TARGET = os.environ.get('AIFLOW_TARGET', '')
if not TARGET:
    # no run yet and no --target: the bridge's primary machine (the desktop this workspace belongs to)
    try:
        _t = json.loads(subprocess.run(['adom-bridge', 'targets'], capture_output=True, text=True, timeout=30).stdout)
        TARGET = _t.get('primaryMachine') or ''
    except Exception:
        TARGET = ''
THREAD = os.environ.get('AIFLOW_THREAD', 'adom-aiflow')
HOME_WIN = os.environ.get('AIFLOW_BOX_DIR', '')  # folder on the box; default under the user's Documents
TOOL = 'project'
STATE = RUN / TOOL


def init(tool):
    global TOOL, STATE
    TOOL = tool
    STATE = RUN / tool


def ok(msg, hints=()):
    print('OK: ' + msg)
    for h in hints:
        if h: print('Hint: ' + h)
    sys.exit(0)


def err(msg, hints=()):
    print('ERROR: ' + msg)
    for h in hints:
        if h: print('Hint: ' + h)
    sys.exit(1)


def clip(text, n=600):
    """Error text within n characters, keeping its beginning (the cause) and its end (issue #76)."""
    text = text or ""
    if len(text) <= n:
        return text
    head = int(n * 0.7)
    return text[:head] + " ... " + text[len(text) - (n - head - 5):]


def say(msg):
    print(msg, flush=True)


def run(cmd, timeout=600, check=False):
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    if check and p.returncode != 0:
        err('%s failed: %s' % (cmd[0], clip((p.stdout + p.stderr).strip(), 400)))
    return p


def ab(verb, args, timeout=600):
    """One adom-bridge verb on the run's desktop; the reply with a JSON-string `output` unwrapped."""
    if not TARGET: err('the run has no desktop target', ['Start the run with --target <box> (adom-aiflow start ... --target), or set AIFLOW_TARGET.'])
    a = dict(args); a.setdefault('reason', 'adom-aiflow %s' % TOOL)
    p = run(['adom-bridge', '--target', TARGET, '--ai-thread', THREAD, verb, json.dumps(a)], timeout=timeout)
    try: d = json.loads(p.stdout)
    except Exception: return {'success': False, 'error': clip(p.stdout + p.stderr, 400)}
    o = d.get('output')
    if isinstance(o, str) and o.startswith('{'):
        try: d = json.loads(o)
        except Exception: pass
    # a tool's error box (Altium paints its text, so the picture is the only copy): keep it where the
    # run's AI can look at it, and say where
    shot = d.get('errorShot') if isinstance(d, dict) else None
    if isinstance(shot, str) and len(shot) > 100:
        import base64
        STATE.mkdir(parents=True, exist_ok=True)
        f = STATE / ('errorShot-%s-%d.png' % (verb, int(time.time())))
        f.write_bytes(base64.b64decode(shot.split(',')[-1]))
        d['errorShot'] = str(f)
        d['_hint'] = ((d.get('_hint') or '') + ' The error box is saved at %s: look at it.' % f).strip()
    return d


def load_state():
    f = STATE / 'project.json'
    return json.loads(f.read_text()) if f.is_file() else {'parts': {}}


def save_state(st):
    STATE.mkdir(parents=True, exist_ok=True)
    (STATE / 'project.json').write_text(json.dumps(st, indent=1))


def box_dir():
    """The project's folder on the desktop, under the user's Documents (one per run)."""
    if HOME_WIN: return HOME_WIN.rstrip('/')
    st = load_state()
    if st.get('boxDir'): return st['boxDir']
    r = ab('shell_execute', {'command': 'echo %USERPROFILE%'})
    out = [l.strip() for l in str(r.get('stdout') or r.get('output') or '').splitlines() if l.strip() and '%' not in l]
    home = out[-1] if out else 'C:/Users/Public'
    # one folder per run, never shared: named after the project folder and the run (every run is
    # "aiflow-run", so the run's name alone sent a second project into the first one's folder and
    # its Altium files), and claimed with a marker naming this run; a folder another run claimed
    # gets a numbered sibling
    base = home.replace('\\', '/') + '/Documents/AdomAiFlow/'
    me = str(RUN.resolve())
    stem = re.sub(r'[^A-Za-z0-9._-]+', '-', '%s-%s' % (RUN.resolve().parent.name, RUN.resolve().name)).strip('-') or 'run'
    d = None
    for k in range(1, 50):
        cand = base + stem + ('' if k == 1 else '-%d' % k)
        w = cand.replace('/', '\\')
        r = ab('shell_execute', {'command': 'if exist "%s\\.aiflow-run" (type "%s\\.aiflow-run") else if exist "%s" (echo AIFLOW-TAKEN) else (echo AIFLOW-FREE)' % (w, w, w)})
        o = str(r.get('stdout') or r.get('output') or '').strip()
        if 'AIFLOW-FREE' in o or me in [l.strip() for l in o.splitlines()]:
            d = cand; break
    if d is None: err('no free project folder under %s for %s' % (base, stem), ['Move old AI Flow project folders out of Documents/AdomAiFlow on the desktop.'])
    w = d.replace('/', '\\')
    ab('shell_execute', {'command': 'mkdir "%s" 2>nul & (echo %s)> "%s\\.aiflow-run"' % (w, me, w)})
    st['boxDir'] = d; save_state(st)
    return d


def win(p):
    return p.replace('/', '\\')


def send(paths, dest):
    r = ab('send_files', {'filePaths': [str(Path(p).resolve()) for p in paths], 'dest': dest, 'destIsFolder': True}, timeout=1800)
    if r.get('success') is False: err('could not copy files to the desktop: %s' % str(r.get('error'))[:300], ['Check `adom-bridge --target %s ping` and retry.' % TARGET])


def prove(key, data=None):
    """Record that a step's own work ran (the flow's proof, read by `step` and `finish`): <run>/proofs/<key>.json."""
    d = RUN / 'proofs'; d.mkdir(parents=True, exist_ok=True)
    (d / (key + '.json')).write_text(json.dumps({'at': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()), 'tool': TOOL, 'data': data or {}}, indent=1))


def lookup_numbers(mpn, page_links, cache):
    """Mouser, DigiKey and LCSC numbers for one MPN, each checked to belong to that exact MPN.
    The page's own links count; the distributor CLIs fill the gaps; a dead backend is said, never guessed."""
    got = {k: v for k, v in (page_links or {}).items() if k in ('mouser', 'digikey', 'lcsc') and v}
    notes = []
    if mpn in cache:
        for k, v in cache[mpn].items():
            if k in ('mouser', 'digikey', 'lcsc') and v and not got.get(k): got[k] = v
    def want(k): return not got.get(k)
    norm = lambda s: re.sub(r'[^A-Z0-9]', '', str(s).upper())
    for k, tool in (('mouser', 'adom-mouser'), ('digikey', 'adom-digikey')):
        if not want(k) or not shutil.which(tool): continue
        p = run([tool, 'part', mpn], timeout=90)
        if p.returncode != 0 or not p.stdout.strip().startswith(('{', '[')):
            txt = (p.stdout + p.stderr)
            m = re.search(r'HTTP (\d{3})', txt)
            notes.append('%s service %s' % (k, ('down (HTTP %s)' % m.group(1)) if m else 'unreachable')); continue
        try: d = json.loads(p.stdout)
        except Exception: continue
        items = d if isinstance(d, list) else (d.get('parts') or d.get('results') or [d])
        for it in items:
            if not isinstance(it, dict): continue
            m = it.get('mpn') or it.get('ManufacturerPartNumber') or it.get('manufacturerPartNumber')
            pn = it.get('mouserPartNumber') or it.get('MouserPartNumber') or it.get('digikeyPartNumber') or it.get('DigiKeyPartNumber') or it.get('partNumber')
            if m and pn and norm(m) == norm(mpn): got[k] = pn; break
    if want('lcsc') and shutil.which('adom-jlcpcb'):
        p = run(['adom-jlcpcb', 'search', mpn], timeout=90)
        try:
            d = json.loads(p.stdout); items = d if isinstance(d, list) else (d.get('results') or d.get('parts') or [])
            for it in items:
                m = it.get('mpn') or it.get('mfrPart') or it.get('manufacturerPartNumber')
                c = it.get('lcsc') or it.get('lcscPart') or it.get('code')
                if m and c and norm(m) == norm(mpn): got['lcsc'] = c; break
        except Exception:
            notes.append('lcsc service unreachable')
    return got, notes


def wiki_page(slug, dest):
    """Clone a wiki component page into dest (fresh every time: its numbers may have been added since)."""
    shutil.rmtree(dest, ignore_errors=True)
    r = run(['adom-wiki', 'repo', 'clone', slug, '--dir', str(dest)], timeout=900)
    if r.returncode != 0 or not (dest / 'page.json').is_file():
        return None
    return json.loads((dest / 'page.json').read_text())


def pick_files(d, comp):
    parts = comp.get('parts', {})
    def first(*names):
        for n in names:
            if n and (d / n).is_file(): return d / n
        return None
    sym = first(parts.get('symbol')) or next(iter(sorted(x for x in d.glob('*.kicad_sym') if 'outline' not in x.name)), None)
    outline = next(iter(sorted(d.glob('*-outline.kicad_sym'))), None)
    fp = first(parts.get('footprint')) or next(iter(sorted(d.glob('*.kicad_mod'))), None)
    step = first(parts.get('model_3d_step'), parts.get('step')) or next(iter(sorted(x for x in d.glob('*.step') if 'unetched' not in x.name and 'uncoloured' not in x.name)), None)
    return sym, outline, fp, step


def kicad_footprints(text):
    out = {}
    for m in re.finditer(r'\(footprint "([^"]+)".*?\(at ([-\d.]+) ([-\d.]+)(?: ([-\d.]+))?\).*?\(property "Reference" "([^"]+)"', text, re.S):
        out[m.group(5)] = (float(m.group(2)), float(m.group(3)), float(m.group(4) or 0) % 360)
    return out


def kicad_outline(text):
    xs, ys = [], []
    for m in re.finditer(r'\(gr_(?:line|rect|arc)[^()]*\(start ([-\d.]+) ([-\d.]+)\)\s*(?:\(mid [-\d.]+ [-\d.]+\)\s*)?\(end ([-\d.]+) ([-\d.]+)\).*?\(layer "Edge.Cuts"\)', text, re.S):
        xs += [float(m.group(1)), float(m.group(3))]; ys += [float(m.group(2)), float(m.group(4))]
    if not xs: err('the reference board has no Edge.Cuts outline')
    return min(xs), min(ys), max(xs), max(ys)


def reference(kicad, spec_path, out_dir, frame_name):
    """A reference KiCad board's placement and the spec's coordinates in the board frame a y-up tool uses
    (the outline's lower-left corner at 0,0). Writes moves.json (the run's plan frame, for place check/land),
    fixed.json (the tool's frame: outline size and fixed parts) and spec.json."""
    text = Path(kicad).read_text(); x0, y0, x1, y1 = kicad_outline(text)
    fps = kicad_footprints(text); spec = json.loads(Path(spec_path).read_text())
    def to_run(x, y): return round(x - x0, 3), round(y - y1, 3)
    fixed = set(spec.get('fixedRefs') or [])
    moves = {r: {'x': to_run(x, y)[0], 'y': to_run(x, y)[1], 'rotation': rot} for r, (x, y, rot) in fps.items() if r not in fixed}
    out = Path(out_dir); out.mkdir(parents=True, exist_ok=True)
    (out / 'moves.json').write_text(json.dumps(moves, indent=1))
    fixed_pos = {r: {'x': round(x - x0, 3), 'y': round(y1 - y, 3), 'rot': rot} for r, (x, y, rot) in fps.items() if r in fixed}
    allpos = {r: {'x': round(x - x0, 3), 'y': round(y1 - y, 3), 'rot': rot} for r, (x, y, rot) in fps.items()}
    (out / 'fixed.json').write_text(json.dumps({'outlineMm': [round(x1 - x0, 3), round(y1 - y0, 3)], 'fixed': fixed_pos, 'all': allpos}, indent=1))
    moved = 0
    def tr(p): return [round(p[0] - x0, 3), round(p[1] - y1, 3)]
    for key in ('pours', 'planeUnder'):
        for z in spec.get(key, []) or []:
            if 'polygon' in z and z['polygon'] and any(abs(v) > 40 for pt in z['polygon'] for v in pt): z['polygon'] = [tr(p) for p in z['polygon']]; moved += 1
    for v in spec.get('stitchVias', []) or []:
        if abs(v.get('x', 0)) > 40 or abs(v.get('y', 0)) > 40: v['x'], v['y'] = tr([v['x'], v['y']]); moved += 1
    spec['frameNote'] = "Coordinates are in the %s board's frame (reference board shifted by (%g, %g))." % (frame_name, -x0, -y1)
    (out / 'spec.json').write_text(json.dumps(spec, indent=1))
    dru = Path(kicad).with_suffix('.kicad_dru')
    if dru.is_file(): shutil.copy(dru, out / 'board.kicad_dru')  # the fab's rules travel with the design
    return moves, fixed_pos, (x1 - x0, y1 - y0), moved


# ---------------------------------------------------------------- manufacturing files (shared by the lanes)

TP_NAME = re.compile(r'test[ _.-]?(point|pad)|^tp[_-]', re.I)
MACHINE_NAME = re.compile(r'machine[ _-]?(contact|pin)|press[ _-]?fit', re.I)


def table(path):
    """A CSV as (header, rows): a tool's preamble (Altium writes a title block above the table) and blank
    rows skipped; the header is the first row naming a designator column."""
    rows = [r for r in csv.reader(open(path, newline='', encoding='utf-8-sig', errors='replace')) if any(c.strip() for c in r)]
    for i, r in enumerate(rows):
        if any(c.strip().lower() in ('designator', 'designators', 'refdes', 'references', 'reference', 'parts', 'part') for c in r):
            return r, rows[i + 1:]
    return (rows[0], rows[1:]) if rows else ([], [])


def fab_rewrite(path, roles, drop_machine=False):
    """Rewrite a BOM or pick-and-place CSV in place: test points out, machine contacts/pins out of the
    CPL (and the fab BOM) or marked DNP in the BOM. A BOM row also names its parts by package and maker, so
    a test point the board source missed is still caught. Returns the refs it changed."""
    head, body = table(path)
    if not head: return set()
    rows = [head] + body
    h = [c.strip().lower() for c in rows[0]]
    di = next((h.index(c) for c in ('designator', 'designators', 'refdes', 'references', 'reference', 'parts', 'part', 'name') if c in h), None)
    if di is None: return set()
    qi = h.index('quantity') if 'quantity' in h else None
    bom = qi is not None or any(c in h for c in ('mpn', 'lcsc', 'lcsc part #'))
    if bom and not drop_machine and 'dnp' not in h: rows[0].append('DNP'); h.append('dnp')
    dn = h.index('dnp') if 'dnp' in h else None
    out, changed = [rows[0]], set()
    for row in rows[1:]:
        if di >= len(row): out.append(row); continue
        refs = [x for x in re.split(r'[ ,;]+', row[di]) if x.strip()]
        sep = ', ' if ',' in row[di] else ' '
        info = ' '.join(row[i] for i, c in enumerate(h) if c in ('package', 'footprint', 'mpn', 'value', 'comment', 'libref') and i < len(row))
        maker = ' '.join(row[i] for i, c in enumerate(h) if c in ('manufacturer', 'mpn') and i < len(row))
        if bom and (TP_NAME.search(info) or re.search(r'pcb copper|board only', maker, re.I)):
            for r in refs: roles.setdefault(r, 'tp')
        elif bom and MACHINE_NAME.search(info):
            for r in refs: roles.setdefault(r, 'machine')
        keep = [r for r in refs if roles.get(r) not in ('tp', 'machine')]
        mach = [r for r in refs if roles.get(r) == 'machine']
        changed |= set(refs) - set(keep)
        def part(rs, dnp):
            nr = list(row) + [''] * (len(h) - len(row)); nr[di] = sep.join(rs)
            if qi is not None: nr[qi] = str(len(rs))
            if dn is not None: nr[dn] = 'DNP' if dnp else ''
            return nr
        if keep: out.append(part(keep, False))
        if mach and bom and not drop_machine: out.append(part(mach, True))
    with open(path, 'w', newline='', encoding='utf-8') as f: csv.writer(f).writerows(out)
    return changed


def refs_in(path, cols=('designator', 'designators', 'refdes', 'references', 'reference', 'parts', 'part'), dnp_too=True):
    head, body = table(path)
    h = [c.strip().lower() for c in head]
    i = next((h.index(c) for c in cols if c in h), None)
    d = h.index('dnp') if 'dnp' in h else None
    out = set()
    for row in body:
        if not dnp_too and d is not None and d < len(row) and row[d].strip(): continue
        if i is not None and i < len(row): out |= {x.strip() for x in re.split(r'[ ,;]+', row[i]) if x.strip()}
    return out