# build_season.py in.bin out.bin  -- adds the season-mode code (season.s, assembled with ../m68kasm.sh) to a patched ROM.
import sys, struct, subprocess, os
here = os.path.dirname(os.path.abspath(__file__))
src, dst = sys.argv[1], sys.argv[2]
d = bytearray(open(src, 'rb').read())
ORG = 0x2F8000
# team strength 0..255 from the ROM's ratings (top 12 skaters + best goalie), for the simulated league games
def tptr(t): return struct.unpack('>I', d[0x30E + 4 * t:0x312 + 4 * t])[0]
score = []
for t in range(32):
    a = tptr(t); h = struct.unpack('>6H', d[a:a + 12]); i = a + h[0]; sk = []; gl = []
    while i < a + h[2] - 2:
        n = struct.unpack('>H', d[i:i + 2])[0]; r = [int(c, 16) for c in d[i + n + 1:i + n + 8].hex()]; i += n + 8
        if r[6] == 0 and r[8] == 0 and r[9] == 0: gl.append(r[3] + r[4] + r[1])
        else: sk.append(r[1] + r[2] + r[3] + r[4] + r[5] + r[8] + r[9] + r[12])
    sk.sort(reverse=True)
    score.append(sum(sk[:12]) / 12 + 1.5 * (max(gl) if gl else 9))
lo, hi = min(score), max(score)
st = [int(255 * (x - lo) / (hi - lo)) for x in score]
open(os.path.join(here, 'strength.inc'), 'w').write('| generated by build_season.py: team strength 0..255 (ANA..WPG)\nstrength:\n' +
    ''.join('        .byte   %s\n' % ','.join(str(v) for v in st[k:k + 8]) for k in range(0, 32, 8)) + '        .even\n')
print('strength', st)
# scorer / assist weights per roster slot (goalies 0), for the simulated games
pw = {'g': [], 'a': []}
for t in range(32):
    a = tptr(t); h = struct.unpack('>6H', d[a:a + 12]); i = a + h[0]; wg = []; wa = []
    while i < a + h[2] - 2:
        n = struct.unpack('>H', d[i:i + 2])[0]; r = [int(c, 16) for c in d[i + n + 1:i + n + 8].hex()]; i += n + 8
        if r[6] == 0 and r[8] == 0 and r[9] == 0: wg.append(0); wa.append(0); continue
        g = r[3] + r[9] + r[5] + r[8]           # off. awareness, accuracy, shot power, stickhandling
        a_ = r[12] + r[3] + r[8] + r[1]          # passing, off. awareness, stickhandling, agility
        wg.append(min(255, g * g // 6)); wa.append(min(255, a_ * a_ // 6))
    while len(wg) < 26: wg.append(0); wa.append(0)
    pw['g'].append(wg[:26]); pw['a'].append(wa[:26])
def table(lbl, rows):
    out = lbl + ':\n'
    for r in rows: out += '        .word   %d\n        .byte   %s\n' % (sum(r), ','.join(map(str, r)))
    return out
open(os.path.join(here, 'pweights.inc'), 'w').write('| generated by build_season.py: per team word total + 26 slot weights\n' +
    table('pw_goal', pw['g']) + table('pw_ast', pw['a']) + '        .even\n')
# goalies per team: starter (line 1 G), backup (first other goalie), number of goalies (they come first in the roster)
gk = []
for t in range(32):
    a = tptr(t); h = struct.unpack('>6H', d[a:a + 12]); i = a + h[0]; n_g = 0
    while i < a + h[2] - 2:
        n = struct.unpack('>H', d[i:i + 2])[0]; r = [int(c, 16) for c in d[i + n + 1:i + n + 8].hex()]; i += n + 8
        if r[6] == 0 and r[8] == 0 and r[9] == 0: n_g += 1
        else: break
    st_ = d[a + h[3]] - 1
    assert 0 <= st_ < n_g <= 3, (t, st_, n_g)
    gk.append((st_, [g for g in range(n_g) if g != st_][0], n_g))
open(os.path.join(here, 'gk.inc'), 'w').write('| generated by build_season.py: per team starter, backup (roster index), number of goalies, 0\ngktab:\n' +
    ''.join('        .byte   %d,%d,%d,0\n' % g for g in gk))
print('goalies', gk)
subprocess.run([os.path.join(here, '..', 'm68kasm.sh'), os.path.join(here, 'season.s'), '%X' % ORG, os.path.join(here, 'season.bin')], check=True)
code = open(os.path.join(here, 'season.bin'), 'rb').read()
assert all(b == 0xFF for b in d[ORG:ORG + len(code)]), 'free space not free'
d[ORG:ORG + len(code)] = code
HOOK = 0x7C2E                                  # jsr $FF88E on the game-over path
assert d[HOOK:HOOK + 6] == bytes.fromhex('4eb9000ff88e'), d[HOOK:HOOK + 6].hex()
d[HOOK:HOOK + 6] = bytes.fromhex('4eb9') + struct.pack('>I', ORG)
# schedule: a fixed shuffle of opponent offsets 1..31 (opponent = (your team + offset) mod 32)
import random
rng = random.Random(1994); off = list(range(1, 32)); rng.shuffle(off)
open(os.path.join(here, 'sched.inc'), 'w').write('| generated by build_season.py: opponent offsets for game k mod 31\nsched:\n' +
    ''.join('        .byte   %s\n' % ','.join(str(v) for v in off[k:k + 8]) for k in range(0, 31, 8)) + '        .even\n')
# season menu (locks Home/Away/Players to the schedule)
MORG = 0x2F9800
subprocess.run([os.path.join(here, '..', 'm68kasm.sh'), os.path.join(here, 'seasonmenu.s'), '%X' % MORG, os.path.join(here, 'seasonmenu.bin')], check=True)
mc = open(os.path.join(here, 'seasonmenu.bin'), 'rb').read()
assert all(b == 0xFF for b in d[MORG:MORG + len(mc)]) and MORG + len(mc) <= 0x2FA000, 'season menu: no room'
d[MORG:MORG + len(mc)] = mc
assert d[0xF7BAA:0xF7BB0] == bytes.fromhex('4eb9002fe000'), 'fix 8 missing in the base ROM'
d[0xF7BAA:0xF7BB0] = bytes.fromhex('4eb9') + struct.pack('>I', MORG)
assert d[0xF7ABA:0xF7AC0] == bytes.fromhex('4eb9002095d6')
d[0xF7ABA:0xF7AC0] = bytes.fromhex('4eb9') + struct.pack('>I', MORG + 8)   # writehook = _start + 8 (bsr.w + bra.w)
# more menu hooks: team rows and labels (addresses of the labels from the symbol list)
sym = subprocess.run(['docker', 'run', '--rm', '-v', here + ':/w', '-w', '/w', '--entrypoint', 'sh', 'ghcr.io/stephane-d/sgdk', '-c',
    'm68k-elf-as -m68000 -o /tmp/m.o seasonmenu.s && m68k-elf-ld -Ttext=0x%X -e _start -o /tmp/m.elf /tmp/m.o && m68k-elf-nm /tmp/m.elf' % MORG],
    capture_output=True, text=True, check=True).stdout
SYM = {l.split()[2]: int(l.split()[0], 16) for l in sym.splitlines() if len(l.split()) == 3}
assert SYM['writehook'] == MORG + 8, hex(SYM['writehook'])
assert d[0xF7C72:0xF7C78] == bytes.fromhex('227cffffd048')
d[0xF7C72:0xF7C78] = bytes.fromhex('4eb9') + struct.pack('>I', SYM['teamrow'])
assert d[0xF80BC:0xF80C2] == bytes.fromhex('4eb900011a48')
d[0xF80BC:0xF80C2] = bytes.fromhex('4eb9') + struct.pack('>I', SYM['labelhook'])
# (tested and dropped: opening the stats screens from the main menu row itself — they need the scoreboard screen setup)
assert d[0x1732C:0x17332] == bytes.fromhex('4eb9000fcc76')
d[0x1732C:0x17332] = bytes.fromhex('4eb9') + struct.pack('>I', SYM['skipmatch'])
# standings screen, replacing "Other Scores" in the pause / post-game menus
SORG = 0x2F9000
subprocess.run([os.path.join(here, '..', 'm68kasm.sh'), os.path.join(here, 'standings.s'), '%X' % SORG, os.path.join(here, 'standings.bin')], check=True)
sc = open(os.path.join(here, 'standings.bin'), 'rb').read()
assert SORG >= ORG + len(code) and all(b == 0xFF for b in d[SORG:SORG + len(sc)]), 'free space not free'
d[SORG:SORG + len(sc)] = sc
old = b'\x00\x14   Other Scores   \x00\x00\x80\xd4'
new = b'\x00\x14 Season Standings ' + struct.pack('>I', SORG)
n = d.count(old); assert n >= 1, 'menu entry not found'
d[:] = d.replace(old, new); print('menu entries replaced:', n)
# scoring leaders, replacing "Record Holders"
LORG = 0x2FA000
subprocess.run([os.path.join(here, '..', 'm68kasm.sh'), os.path.join(here, 'leaders.s'), '%X' % LORG, os.path.join(here, 'leaders.bin')], check=True)
lc = open(os.path.join(here, 'leaders.bin'), 'rb').read()
assert all(b == 0xFF for b in d[LORG:LORG + len(lc)]) and LORG + len(lc) <= 0x2FB000, 'leaders: no room'
d[LORG:LORG + len(lc)] = lc
old = b'\x00\x14  Record Holders  \x00\x0f\xbc\x14'
new = b'\x00\x14 Scoring Leaders  ' + struct.pack('>I', LORG)
n = d.count(old); assert n >= 1, 'Record Holders entry not found'
d[:] = d.replace(old, new); print('leaders entries replaced:', n)
# one team's players (opened from the standings)
TORG = 0x2FB000
subprocess.run([os.path.join(here, '..', 'm68kasm.sh'), os.path.join(here, 'teamstats.s'), '%X' % TORG, os.path.join(here, 'teamstats.bin')], check=True)
tc = open(os.path.join(here, 'teamstats.bin'), 'rb').read()
assert all(b == 0xFF for b in d[TORG:TORG + len(tc)]) and TORG + len(tc) <= 0x2FC000, 'team stats: no room'
d[TORG:TORG + len(tc)] = tc
# league table with the playoff tree's team banners, after every season game (hook $17336)
LTORG = 0x2FC000
subprocess.run([os.path.join(here, '..', 'm68kasm.sh'), os.path.join(here, 'table.s'), '%X' % LTORG, os.path.join(here, 'table.bin')], check=True)
lt = open(os.path.join(here, 'table.bin'), 'rb').read()
assert all(b == 0xFF for b in d[LTORG:LTORG + len(lt)]) and LTORG + len(lt) <= 0x2FD000, 'league table: no room'
d[LTORG:LTORG + len(lt)] = lt
assert d[0x17336:0x1733E] == bytes.fromhex('4a78d0486700e128'), d[0x17336:0x1733E].hex()
d[0x17336:0x1733E] = bytes.fromhex('4ef9') + struct.pack('>I', LTORG) + bytes.fromhex('4e71')
# "New Season" in the pre-game menu: a copy of the Regular Season pre-game list with one more entry,
# and the list pointer (table at $131CC, entry 0) moved to the copy
NORG = 0x2FA800
subprocess.run([os.path.join(here, '..', 'm68kasm.sh'), os.path.join(here, 'newseason.s'), '%X' % NORG, os.path.join(here, 'newseason.bin')], check=True)
nc = open(os.path.join(here, 'newseason.bin'), 'rb').read()
assert all(b == 0xFF for b in d[NORG:NORG + len(nc)]) and NORG + len(nc) <= 0x2FAC00, 'new season: no room'
d[NORG:NORG + len(nc)] = nc
LIST0 = struct.unpack('>I', d[0x131CC:0x131D0])[0]
assert LIST0 == 0x19A84, hex(LIST0)
i = LIST0 + 8; items = []
while d[i:i + 2] == b'\x00\x14': items.append(bytes(d[i:i + 24])); i += 24
assert d[i:i + 4] == bytes.fromhex('0004ff00')
items.append(b'\x00\x14' + b'    New Season    ' + struct.pack('>I', NORG))
NL = 0x2FAC00
blob = bytes(d[LIST0:LIST0 + 8]) + b''.join(items) + bytes(d[i:i + 4])
assert all(b == 0xFF for b in d[NL:NL + len(blob)])
d[NL:NL + len(blob)] = blob
d[0x131CC:0x131D0] = struct.pack('>I', NL)
print('pre-game menu:', len(items), 'entries')
s = sum(struct.unpack('>%dH' % ((len(d) - 0x200) // 2), d[0x200:])) & 0xffff
d[0x18E:0x190] = struct.pack('>H', s)
open(dst, 'wb').write(d); print('season ROM', dst, len(code), 'bytes of code at', hex(ORG))