#!/usr/bin/env python3
"""render_overworld.py — build the seamless Kanto overworld straight from the
PKY source data (no emulator). Ground truth for how maps connect.

  maps/<Label>.blk            raw W*H block indices
  gfx/blocksets/<t>.bst       16 bytes/block = 4x4 tile indices
  gfx/tilesets/<t>.2bpp       16 bytes/tile  = 8x8 2bpp
  data/maps/headers/*.asm     tileset + connections (dir, target, offset)
  constants/map_constants.asm map_const NAME, W, H   (blocks)

Placement (global block coords, top-left of each map), A --dir(offset o)--> B:
  north: (ax+o, ay-Bh)   south: (ax+o, ay+Ah)
  west : (ax-Bw, ay+o)   east : (ax+Aw, ay+o)
"""
import os, re, sys, glob
from PIL import Image

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def P(*a): return os.path.join(ROOT, *a)

# --- map dimensions (blocks) -------------------------------------------------
dims = {}
for m in re.finditer(r'map_const\s+(\w+),\s*(\d+),\s*(\d+)', open(P('constants/map_constants.asm')).read()):
    dims[m.group(1)] = (int(m.group(2)), int(m.group(3)))

# --- tileset const order  <->  tileset name ----------------------------------
tconsts = re.findall(r'^\s*const\s+(\w+)', open(P('constants/tileset_constants.asm')).read(), re.M)
tnames  = re.findall(r'^\s*tileset\s+(\w+),', open(P('data/tilesets/tileset_headers.asm')).read(), re.M)
const2tileset = dict(zip(tconsts, tnames))    # OVERWORLD -> Overworld

# --- tileset name -> (bst path, 2bpp path) -----------------------------------
tblk, tgfx = {}, {}
pend_blk, pend_gfx = [], []
for line in open(P('gfx/tilesets.asm')):
    mb = re.match(r'\s*(\w+)_Block::(?:\s*INCBIN\s+"([^"]+)")?', line)
    mg = re.match(r'\s*(\w+)_GFX::(?:\s*INCBIN\s+"([^"]+)")?', line)
    if mg:
        pend_gfx.append(mg.group(1))
        if mg.group(2):
            for n in pend_gfx: tgfx[n] = mg.group(2)
            pend_gfx = []
    elif mb:
        pend_blk.append(mb.group(1))
        if mb.group(2):
            for n in pend_blk: tblk[n] = mb.group(2)
            pend_blk = []

# --- parse every map header: label, const, tileset, connections --------------
maps = {}            # const -> dict(label, tileset, conns=[(dir,tgtconst,off)])
const2label = {}
for f in glob.glob(P('data/maps/headers/*.asm')):
    txt = open(f).read()
    h = re.search(r'map_header\s+(\w+),\s*(\w+),\s*(\w+),', txt)
    if not h: continue
    label, const, tileset = h.group(1), h.group(2), h.group(3)
    const2label[const] = label
    conns = [(m.group(1), m.group(3), int(m.group(4)))
             for m in re.finditer(r'connection\s+(\w+),\s*(\w+),\s*(\w+),\s*(-?\d+)', txt)]
    maps[const] = dict(label=label, tileset=tileset, conns=conns)

# --- blockset / tileset caches -----------------------------------------------
SHADE = [(255,255,255),(168,168,168),(80,80,80),(0,0,0)]
_bst, _gfx, _rendered = {}, {}, {}
def blockset(name):
    if name not in _bst: _bst[name] = open(P(tblk[name]),'rb').read()
    return _bst[name]
def tileset_gfx(name):
    if name not in _gfx: _gfx[name] = open(P(tgfx[name]),'rb').read()
    return _gfx[name]
def tile_img(gfx, idx):
    px = []
    o = idx*16
    for row in range(8):
        b1, b2 = gfx[o+row*2], gfx[o+row*2+1]
        for col in range(8):
            v = ((b1>>(7-col))&1) | (((b2>>(7-col))&1)<<1)
            px.append(SHADE[v])
    im = Image.new('RGB',(8,8)); im.putdata(px); return im

def render_map(const):
    if const in _rendered: return _rendered[const]
    label = maps[const]['label']; tname = const2tileset[maps[const]['tileset']]
    w,h = dims[const]
    blk = open(P('maps', label+'.blk'),'rb').read()
    bs, gfx = blockset(tname), tileset_gfx(tname)
    img = Image.new('RGB',(w*32,h*32),(0,0,0))
    for br in range(h):
        for bc in range(w):
            if br*w+bc >= len(blk): continue
            block = blk[br*w+bc]
            for ty in range(4):
                for tx in range(4):
                    ti = bs[block*16 + ty*4 + tx]
                    img.paste(tile_img(gfx,ti), (bc*32+tx*8, br*32+ty*8))
    _rendered[const] = img
    return img

# --- BFS placement from PALLET_TOWN ------------------------------------------
pos = {'PALLET_TOWN': (0,0)}
q = ['PALLET_TOWN']
while q:
    a = q.pop(0)
    if a not in maps: continue
    ax, ay = pos[a]; aw, ah = dims[a]
    for d, b, o in maps[a]['conns']:
        if b in pos or b not in dims: continue
        bw, bh = dims[b]
        if   d == 'north': pos[b] = (ax+o, ay-bh)
        elif d == 'south': pos[b] = (ax+o, ay+ah)
        elif d == 'west':  pos[b] = (ax-bw, ay+o)
        elif d == 'east':  pos[b] = (ax+aw, ay+o)
        else: continue
        q.append(b)

placed = [c for c in pos if c in maps and os.path.exists(P('maps', maps[c]['label']+'.blk'))]
minx = min(pos[c][0] for c in placed); miny = min(pos[c][1] for c in placed)
maxx = max(pos[c][0]+dims[c][0] for c in placed)
maxy = max(pos[c][1]+dims[c][1] for c in placed)
W, H = (maxx-minx)*32, (maxy-miny)*32
print(f"placed {len(placed)} maps; world = {maxx-minx}x{maxy-miny} blocks = {W}x{H}px")

canvas = Image.new('RGB',(W,H),(20,20,30))
for c in sorted(placed, key=lambda c: (pos[c][1],pos[c][0])):
    x = (pos[c][0]-minx)*32; y = (pos[c][1]-miny)*32
    canvas.paste(render_map(c), (x,y))

out = sys.argv[1] if len(sys.argv)>1 else '/tmp/kanto_overworld.png'
canvas.save(out)
# also a downscaled preview for quick viewing
scale = max(1, max(W,H)//1600 + 1)
canvas.resize((W//scale, H//scale)).save(out.replace('.png','_preview.png'))
print(f"saved {out} ({W}x{H})  + preview /{scale}x")