#!/usr/bin/python3
"""One-shot body-centred placement with incremental copper repair, not regeneration.

Writes a candidate only. Review its DRC before promoting it to the native board.
Requires the retained pre-alignment snapshot and no editor locks.
"""
import hashlib
import json
import time
import subprocess
from pathlib import Path
import pcbnew as p
from shapely.geometry import Point, LineString, box
import route_pcb_draft as r

ROOT = Path(__file__).resolve().parent
BASE = ROOT/'validation/pcb-before-group-alignment.kicad_pcb'
MOVES = {'DS1': ((143,136),(143.97,136)),
         **{f'R{i+1}': ((143+11*i,134),(142.16+11*i,132.6685)) for i in range(3)},
         **{f'SW{i+1}': ((143+12*i,168),(141.72+12*i,168)) for i in range(3)}}


def unlocked():
    assert not list(ROOT.glob('*.lck')), 'KiCad editor lock exists'
    assert not list(ROOT.glob('.~lock.*')), 'Document editor lock exists'


def geom(item):
    if isinstance(item,p.PAD):
        if item.GetShape()==p.PAD_SHAPE_CIRCLE:
            return Point(r.pos(item)).buffer(p.ToMM(item.GetSize().x)/2,quad_segs=32)
        bb=item.GetBoundingBox()
        return box(p.ToMM(bb.GetX()),p.ToMM(bb.GetY()),p.ToMM(bb.GetRight()),p.ToMM(bb.GetBottom()))
    if isinstance(item,p.PCB_VIA):
        return Point(r.pos(item)).buffer(.35,quad_segs=32)
    return LineString([xy(item.GetStart()),xy(item.GetEnd())]).buffer(p.ToMM(item.GetWidth())/2,quad_segs=32)


def xy(v):
    return (p.ToMM(v.x),p.ToMM(v.y))


def layers(item):
    return set(r.LAYERS) if isinstance(item,(p.PAD,p.PCB_VIA)) else {item.GetLayer()}


def groups(board, net):
    items=[a for f in board.GetFootprints() for a in f.Pads() if a.GetNetCode()==net]
    items += [t for t in board.GetTracks() if t.GetNetCode()==net]
    shapes=[geom(i) for i in items]
    parents=list(range(len(items)))
    def root(i):
        while parents[i]!=i:
            i=parents[i]
        return i
    for i,a in enumerate(items):
        for j in range(i):
            if layers(a)&layers(items[j]) and shapes[i].distance(shapes[j])<1e-6:
                parents[root(i)]=root(j)
    result={}
    for i,a in enumerate(items):
        result.setdefault(root(i),[]).append(a)
    return list(result.values())


def anchors(group):
    result=[]
    for a in group:
        points=[r.pos(a)] if isinstance(a,(p.PAD,p.PCB_VIA)) else [xy(a.GetStart()),xy(a.GetEnd())]
        result.extend((point,layers(a)) for point in points)
    return result


def main():
    unlocked()
    before=r.PCB.read_bytes()
    assert before==BASE.read_bytes(), 'Input differs from preserved baseline'
    protected={name:(ROOT/name).read_bytes() for name in ('serial-carrier.kicad_pro','serial-carrier.kicad_prl','serial-carrier.kicad_sch')}
    board=p.LoadBoard(str(r.PCB))
    fps={f.GetReference():f for f in board.GetFootprints()}
    for ref,(old,new) in MOVES.items():
        assert r.pos(fps[ref])==old
        fps[ref].SetPosition(r.point(new))
        if ref.startswith('R'):
            fps[ref].SetValue('2k 1%')
    pads=[a for f in board.GetFootprints() for a in f.Pads()]
    removed=[]
    # Only remove copper made illegal by a moved pad/hole. Everything else is
    # initially retained, including all unrelated nets and manual routing.
    moved_pads=[a for ref in MOVES for a in fps[ref].Pads()]
    for t in list(board.GetTracks()):
        if any((not a.GetNetCode() or a.GetNetCode()!=t.GetNetCode()) and
               geom(t).distance(geom(a))<.25-1e-6 for a in moved_pads):
            removed.append(t.m_Uuid.AsString())
            board.Delete(t)

    repaired=[]
    deadline=time.monotonic()+600
    for ref,(old,new) in MOVES.items():
        for pad in fps[ref].Pads():
            if not pad.GetNumber(): continue
            target=r.pos(pad)
            start=(round(target[0]+old[0]-new[0],6),round(target[1]+old[1]-new[1],6))
            if geom(pad).covers(Point(start)): continue
            attached=[t for t in board.GetTracks() if t.GetNetCode()==pad.GetNetCode() and start in (xy(t.GetStart()),xy(t.GetEnd()))]
            if not attached: continue
            width=.5 if pad.GetNetname() in ('/+3V3','/GND') else .25
            blocked,via=r.masks(board,pads,pad.GetNetCode(),width)
            allowed=set.union(*(layers(t) for t in attached))
            for z,layer in enumerate(r.LAYERS):
                if layer not in allowed: blocked[r.node(start,z)]=True
            path,_=r.astar(blocked,via,start,target,deadline,500000)
            assert path, ('Cannot extend moved pad',ref,pad.GetNumber())
            r.add_path(board,path,start,target,pad.GetNetCode(),width)
            repaired.append(dict(net=pad.GetNetname(),start=start,end=target))
    nets=sorted({a.GetNetCode() for a in pads if a.GetNumber() and a.GetNetCode()}, key=lambda n: (board.FindNet(n).GetNetname() in ('/+3V3','/GND'), not board.FindNet(n).GetNetname().startswith('/OLED'), n))
    for net in nets:
        while len(parts:=groups(board,net))>1:
            options=[]
            for i,aa in enumerate(parts):
                for bb in parts[i+1:]:
                    for a,la in anchors(aa):
                        for b,lb in anchors(bb):
                            options.append(((a[0]-b[0])**2+(a[1]-b[1])**2,a,la,b,lb))
            options.sort(key=lambda v:v[0])
            name=parts[0][0].GetNetname()
            width=.5 if name in ('/+3V3','/GND') else .25
            blocked,via=r.masks(board,pads,net,width)
            success=False
            for _,a,la,b,lb in options[:400]:
                mask=blocked.copy()
                for z,layer in enumerate(r.LAYERS):
                    if layer not in la: mask[r.node(a,z)]=True
                    if layer not in lb: mask[r.node(b,z)]=True
                path,_=r.astar(mask,via,a,b,deadline,250000)
                if path:
                    r.add_path(board,path,a,b,net,width)
                    repaired.append(dict(net=name,start=a,end=b))
                    print(repaired[-1],flush=True)
                    success=True
                    break
                assert time.monotonic()<deadline, 'Bounded routing deadline reached'
            assert success, ('Cannot reconnect retained copper',name)
    unlocked()
    assert r.PCB.read_bytes()==before
    assert all((ROOT/n).read_bytes()==data for n,data in protected.items())
    out=ROOT/'validation/pcb-group-candidate.kicad_pcb'
    # Let KiCad identify true dangling stubs after repair. Do not prune valid
    # retained routes before reconnecting the moved pads.
    from add_device_silkscreen import add_outlines
    drc=ROOT/'validation/pcb-group-candidate-drc.json'
    for attempt in range(20):
        add_outlines(board)
        assert p.SaveBoard(str(out),board,True)
        subprocess.run(['kicad-cli','pcb','drc','--format','json','--all-track-errors','--severity-all','-o',str(drc),str(out)],check=True,timeout=60)
        results=json.loads(drc.read_text())
        ids={i['uuid'] for v in results['violations'] if v['type']=='track_dangling' for i in v['items']}
        if not ids:
            assert not results['unconnected_items'], results['unconnected_items']
            assert all(v['type']=='lib_footprint_issues' for v in results['violations']), results['violations']
            break
        for t in list(board.GetTracks()):
            if t.m_Uuid.AsString() in ids:
                removed.append(t.m_Uuid.AsString())
                board.Delete(t)
    else:
        raise AssertionError('Dangling-stub cleanup did not converge')
    report=dict(baseline_sha256=hashlib.sha256(before).hexdigest(),moves_absolute_mm=MOVES,
                removed_track_uuids=removed,repaired_connections=repaired,
                retained_tracks=sum(t.m_Uuid.AsString() not in removed for t in p.LoadBoard(str(BASE)).GetTracks()),
                protected_sha256={n:hashlib.sha256(data).hexdigest() for n,data in protected.items()})
    (ROOT/'validation/pcb-group-migration.json').write_text(json.dumps(report,indent=2)+'\n')

if __name__=='__main__':
    main()