#!/usr/bin/env python3
"""gbshot.py <sock> <out.bmp> — reconstruct the current GB screen from VRAM
(BG + window + sprites, both CGB banks) and save a 24-bit BMP. Grayscale from
the 2bpp tile values (enough to read menus/sprites). Uses the sl0pboy control socket."""
import socket, sys, struct

def main():
    sock, out = sys.argv[1], sys.argv[2]
    c = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM); c.settimeout(6); 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):
        out=[]
        while n>0:
            k=min(4096,n); h=cmd(f"read 0x{a:04X} {k}").split()[-1]
            out+=[int(h[i*2:i*2+2],16) for i in range(k)]; a+=k; n-=k
        return out
    reg=lambda a: rd(a,1)[0]
    LCDC=reg(0xFF40);SCY=reg(0xFF42);SCX=reg(0xFF43);WY=reg(0xFF4A);WX=reg(0xFF4B)
    cmd("write 0xFF4F 00"); v0=rd(0x8000,0x2000)
    cmd("write 0xFF4F 01"); v1=rd(0x8000,0x2000)
    cmd("write 0xFF4F 00"); oam=rd(0xFE00,0xA0)
    vb=lambda bank,addr:(v0 if bank==0 else v1)[addr-0x8000]
    bgmap=0x9C00 if (LCDC>>3)&1 else 0x9800
    wmap =0x9C00 if (LCDC>>6)&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(mapbase,mx,my):
        ent=mapbase+(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)]
    winon=(LCDC>>5)&1
    for sy in range(144):
        for sx in range(160):
            if winon and sy>=WY and sx>=WX-7:
                fb[sy][sx]=fetch(wmap,sx-(WX-7),sy-WY)
            elif LCDC&1:
                fb[sy][sx]=fetch(bgmap,(sx+SCX)&0xFF,(sy+SCY)&0xFF)
    h=16 if (LCDC>>2)&1 else 8
    if (LCDC>>1)&1:
        for i in range(40):
            y=oam[i*4]-16;x=oam[i*4+1]-8;t=oam[i*4+2];a=oam[i*4+3]
            bank=(a>>3)&1;xf=(a>>5)&1;yf=(a>>6)&1
            if h==16:t&=0xFE
            for part in range(h//8):
                tp=t+(part if not yf else (h//8-1-part))
                base=0x8000+tp*16
                for ty in range(8):
                    ry=7-ty if yf else ty
                    b1=vb(bank,base+ry*2);b2=vb(bank,base+ry*2+1)
                    yy=y+part*8+ty
                    if 0<=yy<144:
                        for tx in range(8):
                            rx=7-tx if xf else tx
                            val=((b1>>(7-rx))&1)|(((b2>>(7-rx))&1)<<1)
                            if val==0: continue
                            xx=x+tx
                            if 0<=xx<160: fb[yy][xx]=shade[val]
    W,H=160,144;rowpad=(-W*3)%4;data=bytearray()
    for ry in range(H-1,-1,-1):
        for x in range(W):
            g=fb[ry][x];data+=bytes((g,g,g))
        data+=b"\x00"*rowpad
    hdr=b"BM"+struct.pack("<IHHI",54+len(data),0,0,54)+struct.pack("<IiiHHIIiiII",40,W,H,1,24,0,len(data),2835,2835,0,0)
    open(out,"wb").write(hdr+data)
    print(f"saved {out}  (map=0x{reg(0xD35D):02X} pos {reg(0xD360)},{reg(0xD361)})")

if __name__=="__main__": main()