#!/usr/bin/env python3
"""stitch_overworld.py <sock> <out.png> — drive EXPLORE mode, DFS the map
connection graph, screenshot every reachable overworld map, and composite them
into one big overworld image positioned by connection direction.
Assumes EXPLORE is ALREADY open (on the origin map) when this starts."""
import socket, sys, time
from PIL import Image, ImageDraw

sock, out = sys.argv[1], sys.argv[2]
c = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM); c.settimeout(8); c.connect(sock); c.recv(200)
def cmd(l):
    c.sendall((l+"\n").encode()); b=b""
    while b"\n" not in b: b+=c.recv(65536)
    return b.split(b"\n",1)[0].decode()
def rd(a,n):
    o=[]
    while n>0:
        k=min(4096,n); h=cmd(f"read 0x{a:04X} {k}").split()[-1]
        o+=[int(h[i*2:i*2+2],16) for i in range(k)]; a+=k; n-=k
    return o
reg=lambda a: rd(a,1)[0]

def capture():
    LCDC=reg(0xFF40);SCY=reg(0xFF42);SCX=reg(0xFF43)
    cmd("write 0xFF4F 00"); v0=rd(0x8000,0x2000)
    cmd("write 0xFF4F 01"); v1=rd(0x8000,0x2000)
    cmd("write 0xFF4F 00")
    vb=lambda bank,addr:(v0 if bank==0 else v1)[addr-0x8000]
    bgmap=0x9C00 if (LCDC>>3)&1 else 0x9800
    uns=(LCDC>>4)&1
    ta=lambda idx:0x8000+idx*16 if uns else 0x9000+((idx^0x80)-0x80)*16
    shade=[255,170,85,0]
    fb=[[255]*160 for _ in range(144)]
    def fetch(mx,my):
        ent=bgmap+(my//8)*32+(mx//8)
        idx=vb(0,ent);attr=vb(1,ent)
        bank=(attr>>3)&1;xf=(attr>>5)&1;yf=(attr>>6)&1
        a=ta(idx); ry=(7-(my%8)) if yf else (my%8); rx=(7-(mx%8)) if xf else (mx%8)
        b1=vb(bank,a+ry*2);b2=vb(bank,a+ry*2+1)
        return shade[((b1>>(7-rx))&1)|(((b2>>(7-rx))&1)<<1)]
    for sy in range(144):
        for sx in range(160):
            if LCDC&1: fb[sy][sx]=fetch((sx+SCX)&0xFF,(sy+SCY)&0xFF)
    img=Image.new("L",(160,144))
    img.putdata([fb[y][x] for y in range(144) for x in range(160)])
    return img.convert("RGB")

def curmap(): return reg(0xD35D)
def conns():
    cn=reg(0xD36F)
    return {'N':reg(0xD370) if cn&8 else None,
            'S':reg(0xD37B) if cn&4 else None,
            'W':reg(0xD386) if cn&2 else None,
            'E':reg(0xD391) if cn&1 else None}
def press(k, settle=1.4):
    cmd("+"+k); time.sleep(0.26); cmd("-"+k); time.sleep(settle)

def move(k, src, target):
    """press k to go src->target, with retries; only re-press while still on src."""
    for _ in range(4):
        press(k)
        m=curmap()
        if m==target: return True
        if m!=src: return False   # ended somewhere unexpected; don't hammer
    return False

OFF={'N':(0,-1),'S':(0,1),'W':(-1,0),'E':(1,0)}
KEY={'N':'up','S':'down','W':'left','E':'right'}
OPP={'N':'S','S':'N','W':'E','E':'W'}

cmd("release")
shots={}   # map_id -> RGB image
grid={}    # map_id -> (col,row)
cells={}   # (col,row) -> map_id  (collision detect)
order=[]

def visit(mid,col,row):
    shots[mid]=capture(); grid[mid]=(col,row); order.append(mid)
    if (col,row) in cells and cells[(col,row)]!=mid:
        print(f"  COLLISION at ({col},{row}): 0x{cells[(col,row)]:02X} vs 0x{mid:02X}")
    cells[(col,row)]=mid
    cn=conns()
    print(f"map 0x{mid:02X} @({col},{row}) conn "+" ".join(f"{d}=0x{v:02X}" for d,v in cn.items() if v is not None))
    for d in ('N','S','W','E'):
        nb=cn[d]
        if nb is None or nb in grid: continue
        if not move(KEY[d], mid, nb):
            print(f"  !! {d} 0x{mid:02X}->0x{nb:02X} failed (on 0x{curmap():02X}); resync")
            # try to get back to mid before continuing
            if curmap()!=mid: move(KEY[OPP[d]], curmap(), mid)
            continue
        dc,dr=OFF[d]; visit(nb,col+dc,row+dr)
        if not move(KEY[OPP[d]], nb, mid):
            print(f"  !! backtrack 0x{nb:02X}->0x{mid:02X} failed (on 0x{curmap():02X})")

start=curmap()
print(f"start map=0x{start:02X}")
visit(start,0,0)

cols=[p[0] for p in grid.values()]; rows=[p[1] for p in grid.values()]
minc,maxc=min(cols),max(cols); minr,maxr=min(rows),max(rows)
W=(maxc-minc+1)*160; H=(maxr-minr+1)*144
canvas=Image.new("RGB",(W,H),(40,40,40))
dr=ImageDraw.Draw(canvas)
for mid,(col,row) in grid.items():
    x=(col-minc)*160; y=(row-minr)*144
    canvas.paste(shots[mid],(x,y))
    dr.rectangle([x,y,x+159,y+143],outline=(255,0,0))
    dr.text((x+2,y+2),f"{mid:02X}",fill=(255,0,0))
canvas.save(out)
print(f"\nstitched {len(shots)} maps -> {out}  ({W}x{H})")