import struct,sys,os
src=sys.argv[1]; dst=sys.argv[2]; opts=set(sys.argv[3:])
d=bytearray(open(src,'rb').read())
def W(*ws): return b''.join(struct.pack('>H',w&0xffff) for w in ws)
def L(v): return struct.pack('>I',v)
def asm(base,items):
    # items: bytes | ('lbl',name) | ('bra',opword,name)  (16-bit displacement branches)
    pos=base; labels={}; out=[]
    for it in items:
        if isinstance(it,tuple) and it[0]=='lbl': labels[it[1]]=pos
        elif isinstance(it,tuple): pos+=4
        else: pos+=len(it)
    code=bytearray(); pos=base
    for it in items:
        if isinstance(it,tuple) and it[0]=='lbl': continue
        if isinstance(it,tuple):
            disp=labels[it[2]]-(pos+2); code+=W(it[1],disp); pos+=4
        else: code+=it; pos+=len(it)
    return bytes(code)
def put(addr,code,expect=None):
    if expect is not None: assert d[addr:addr+len(expect)]==expect,(hex(addr),d[addr:addr+len(expect)].hex())
    d[addr:addr+len(code)]=code
HOOK1=0x2FF000; HOOK2=0x2FF080
assert all(b==0xff for b in d[HOOK1:HOOK1+0x200]), "free space not free"
# ---- Fix 1: pass target selection prefers the other human teammate ----
if 'nosel' not in opts:
    put(HOOK1, asm(HOOK1,[
        W(0xC3C1),            # muls.w d1,d1
        W(0xC5C2),            # muls.w d2,d2
        W(0xD481),            # add.l d1,d2
        W(0xD480),            # add.l d0,d2       (orig cost = dist^2 + angle penalty)
        W(0x082B,3,0x62),     # btst #3,$62(a3)   passer human-controlled?
        ('bra',0x6700,'cmp'),
        W(0x0829,3,0x62),     # btst #3,$62(a1)   candidate human-controlled?
        ('bra',0x6700,'cmp'),
        W(0x7400),            # moveq #0,d2       -> best possible cost
        ('lbl','cmp'),
        W(0xB484),            # cmp.l d4,d2
        ('bra',0x6200,'done'),# bhi
        W(0x2802),            # move.l d2,d4
        W(0x2049),            # movea.l a1,a0
        ('lbl','done'),
        W(0x4E75)]))          # rts
    orig=bytes.fromhex('c3c1c5c2d481d480b484620000062802 2049'.replace(' ',''))
    put(0xBD44, W(0x4EB9)+L(HOOK1)+W(0x6000,0xBD56-0xBD4C)+W(0x4E71,0x4E71,0x4E71,0x4E71), orig)
# ---- Fix 2: bigger stick pickup reach for the human who is the pass target ----
if 'norecv' not in opts:
    R=int(next((o.split('=')[1] for o in opts if o.startswith('r=')),'24'))
    put(HOOK2, asm(HOOK2,[
        W(0x343C,0x000E),           # move.w #14,d2
        W(0x263C)+L(0xC4),          # move.l #196,d3
        W(0x082A,3,0x62),           # btst #3,$62(a2)   human-controlled?
        ('bra',0x6700,'std'),
        W(0x382A,0x0052),           # move.w $52(a2),d4
        W(0xB878,0xBEE0),           # cmp.w $bee0.w,d4  is he the pass target?
        ('bra',0x6600,'std'),
        W(0x343C,R),                # move.w #R,d2
        W(0x263C)+L(R*R),           # move.l #R*R,d3
        ('lbl','std'),
        W(0xB042), ('bra',0x6E00,'fail'),   # cmp.w d2,d0 ; bgt
        W(0x4442),                          # neg.w d2
        W(0xB042), ('bra',0x6D00,'fail'),   # cmp.w d2,d0 ; blt
        W(0xD26A,0x0014), W(0x926B,0x0014), # add.w $14(a2),d1 ; sub.w $14(a3),d1
        W(0xB242), ('bra',0x6D00,'fail'),   # cmp.w d2,d1 ; blt  (d2 negative)
        W(0x4442),                          # neg.w d2
        W(0xB242), ('bra',0x6E00,'fail'),   # cmp.w d2,d1 ; bgt
        W(0xC1C0), W(0xC3C1), W(0xD081),    # d0 = dx^2+dy^2
        W(0x2203),                          # move.l d3,d1
        W(0x4A2B,0x005E),                   # tst.b $5e(a3)
        ('bra',0x6F00,'cmp2'),              # ble
        W(0xE489),                          # lsr.l #2,d1
        ('lbl','cmp2'),
        W(0xB081), ('bra',0x6200,'fail'),   # cmp.l d1,d0 ; bhi
        W(0x4EF9)+L(0x14EC8),               # jmp pickup
        ('lbl','fail'),
        W(0x4EF9)+L(0x14ED0)]))             # jmp no-pickup path
    put(0x14E84, W(0x4EF9)+L(HOOK2), bytes.fromhex('b07c000e6e000046'))
# ---- Fix 3: goalie passes: never blind, softer short passes ----
HOOK4=0x2FF140
if 'nogk' not in opts:
    GMIN=int(next((o.split('=')[1] for o in opts if o.startswith('gmin=')),'3'))
    GADD=int(next((o.split('=')[1] for o in opts if o.startswith('gadd=')),'1'))
    # Goalie aims like a skater (stick direction, +-45 deg cone, other human preferred inside it).
    # Only if nobody is in that direction: instead of a blind pass straight out, pick the other
    # human (co-op) or else the nearest teammate, regardless of direction.
    HOOK5=0x2FF180
    put(HOOK5, asm(HOOK5,[
        W(0x4A84), ('bra',0x6A00,'found'),          # tst.l d4 ; bpl found
        W(0x4A6B,0x0034), ('bra',0x6600,'blind'),   # skater -> original blind pass
        W(0x7605),                                   # moveq #5,d3
        W(0x327C,0xB04A),                            # movea.w #$b04a,a1
        W(0x0C6B,0x0006,0x0052), ('bra',0x6D00,'loop'),  # cmpi.w #6,$52(a3) ; blt
        W(0xD2FC,0x0300),                            # adda.w #$300,a1 (team 2)
        ('lbl','loop'),
        W(0xB7C9), ('bra',0x6700,'next'),            # cmpa.l a1,a3 ; beq
        W(0x4A69,0x0034), ('bra',0x6700,'next'),     # goalie / not a skater
        W(0x0829,2,0x63), ('bra',0x6600,'next'),     # unavailable
        W(0x3211), W(0x9278,0xB74A),                 # d1 = x - puck x
        W(0x3429,0x0014), W(0x9478,0xB75E),          # d2 = y - puck y
        W(0xC3C1), W(0xC5C2), W(0xD481),             # d2 = dx^2+dy^2
        W(0x082B,3,0x62), ('bra',0x6700,'cmp'),      # passer human?
        W(0x0829,3,0x62), ('bra',0x6700,'cmp'),      # candidate human?
        W(0x7400),                                   # -> other human wins
        ('lbl','cmp'),
        W(0xB484), ('bra',0x6200,'next'),            # cmp.l d4,d2 ; bhi
        W(0x2802), W(0x2049),                        # d4=d2 ; a0=a1
        ('lbl','next'),
        W(0xD2FC,0x0080),
        ('bra',0x51CB,'loop'),                       # dbra d3,loop
        W(0x4A84), ('bra',0x6B00,'blind'),
        ('lbl','found'),
        W(0x4EF9)+L(0xBD64),
        ('lbl','blind'),
        W(0x4EF9)+L(0xBD6C)]))
    put(0xBD5E, W(0x4EF9)+L(HOOK5), bytes.fromhex('4a846b00000a'))
    # flight time d2: game truncates it, so short passes fly up to 2x speed. Goalie: round up + minimum.
    put(HOOK4, asm(HOOK4,[
        W(0xB47C,0x0018),            # cmp.w #$18,d2
        ('bra',0x6300,'ok'),         # bls
        W(0x7418),                   # moveq #$18,d2
        ('lbl','ok'),
        W(0x4A6B,0x0034),            # tst.w $34(a3)  passer goalie?
        ('bra',0x6600,'done'),
        W(0x0642,GADD),              # addi.w #GADD,d2
        W(0xB47C,GMIN),              # cmp.w #GMIN,d2
        ('bra',0x6C00,'done'),       # bge
        W(0x343C,GMIN),              # move.w #GMIN,d2
        ('lbl','done'),
        W(0x4E75)]))
    put(0xBF54, W(0x4EB9)+L(HOOK4)+W(0x4E71,0x4E71), bytes.fromhex('b47c0018630000047418'))
# ---- Fix 4: long shots (e.g. from the blue line) score a bit more often ----
# Goalie stops everything below height 15 within 12-15 px, and the crossbar is at 13, so a set goalie
# was basically unbeatable from distance (about 0.3% of blue-line shots went in).
# - Shot hook ($C408): distance to the net is stored in free RAM $FFF000. For long shots
#   (distance >= bld) the game's own RNG ($11086) decides once per shot, with chance blp %,
#   that the goalie misses it ($FFF002 set).
# - Goalie reach hook ($15032): while $FFF002 is set the goalie does not touch the puck.
#   Puck still has to be under the bar and between the posts.
# - Lob hook ($C51A): long shots get vz -= vz>>blz, so fewer of them sail over the bar.
# - Pickup hook ($150F6): any player getting the puck clears both flags.
if 'nobl' not in opts:
    DST=int(next((o.split('=')[1] for o in opts if o.startswith('bld=')),'160'))
    P=int(next((o.split('=')[1] for o in opts if o.startswith('blp=')),'15'))
    Z=int(next((o.split('=')[1] for o in opts if o.startswith('blz=')),'3'))
    T=int(next((o.split('=')[1] for o in opts if o.startswith('blt=')),'40'))
    HS=0x2FF300; HR=0x2FF340; HP=0x2FF380; HZ=0x2FF3C0
    assert all(b==0xff for b in d[0x2FF300:0x2FF400]), "free space not free"
    put(HS, asm(HS,[
        W(0x3600),                    # move.w d0,d3          (original)
        W(0x31C3,0xF000),             # move.w d3,$f000.w     shot distance
        W(0x4238,0xF002),             # clr.b $f002.w
        W(0xB67C,DST),                # cmp.w #DST,d3
        ('bra',0x6D00,'out'),         # blt
        W(0x303C,100),                # move.w #100,d0
        W(0x4EB9)+L(0x11086),         # jsr rnd(0..99)
        W(0xB07C,P),                  # cmp.w #P,d0
        ('bra',0x6C00,'out'),         # bge
        W(0x11FC,T,0xF002),           # move.b #T,$f002.w     goalie will miss this one (countdown)
        ('lbl','out'),
        W(0x3003),                    # move.w d3,d0
        W(0x0838,4,0xC2EA),           # btst #4,$c2ea.w       (original; flags used after return)
        W(0x4E75)]))
    put(0xC408, W(0x4EB9)+L(HS)+W(0x4E71), bytes.fromhex('360008380004c2ea'))
    put(HR, asm(HR,[
        W(0x205F),                    # movea.l (a7)+,a0      (original)
        W(0x4A38,0xF002),             # tst.b $f002.w
        ('bra',0x6700,'norm'),
        W(0x5338,0xF002),             # subq.b #1,$f002.w     (counts while the puck is near a goalie)
        W(0x4EF9)+L(0x150C2),         # jmp: puck not in reach
        ('lbl','norm'),
        W(0x0C78,0x000F,0xB762),      # cmpi.w #$f,$b762.w    (original)
        W(0x4EF9)+L(0x1503A)]))       # jmp back (flags kept)
    put(0x15032, W(0x4EF9)+L(HR)+W(0x4E71), bytes.fromhex('205f0c78000fb762'))
    put(HP, asm(HP,[
        W(0x4A6A,0x0034),             # tst.w $34(a2)         goalie?
        ('bra',0x6600,'take'),
        W(0x4A38,0xF002),             # tst.b $f002.w
        ('bra',0x6700,'take'),
        W(0x588F),                    # addq.l #4,a7
        W(0x4EF9)+L(0x15464),         # goalie can't grab it: return like the original abort
        ('lbl','take'),
        W(0x42B8,0xF000),             # clr.l $f000.w
        W(0x08AA,0,0x62),             # bclr #0,$62(a2)       (original)
        W(0x4E75)]))
    put(0x150F6, W(0x4EB9)+L(HP), bytes.fromhex('08aa00000062'))
    if Z>0:
        put(HZ, asm(HZ,[
            W(0x0C78,DST,0xF000),     # cmpi.w #DST,$f000.w
            ('bra',0x6D00,'st'),      # blt
            W(0x3F01),                # move.w d1,-(a7)
            W(0x3203),                # move.w d3,d1
            W(0xE049|((Z&7)<<9)),     # lsr.w #Z,d1
            W(0x9641),                # sub.w d1,d3
            W(0x321F),                # move.w (a7)+,d1
            ('lbl','st'),
            W(0x31C3,0xB776),         # move.w d3,$b776.w     (original)
            W(0x08B8,4,0xC2F6),       # bclr #4,$c2f6.w       (original; Z flag used after return)
            W(0x4E75)]))
        put(0xC51A, W(0x4EB9)+L(HZ)+W(0x4E71,0x4E71), bytes.fromhex('31c3b77608b80004c2f6'))
# ---- Fix 5: off-screen arrow (e.g. goalie in shootouts) drawn 32 px too far left ----
# $164D6 places the arrow at (obj x - camera x) + $100, i.e. sprite x 128 + 128 = centre of the original
# 256 px wide screen. This hack runs 320 px (H40), centre 160, so the arrow needs + $120.
if 'noarrow' not in opts:
    put(0x16590, W(0x0640,0x0120), bytes.fromhex('06400100'))
# ---- Fix 6: Fighting setting, menu Off/On only, stored where nothing overwrites it ----
# The hack keeps the Fighting option in $FFDF00, but that word lies inside its own event log
# ($DEF6 + $DEF4, 14 bytes per entry, written at $1EFC12). After the first entry the setting holds
# player data (e.g. $0D12) and fights are on or off regardless of the menu. Value 0 = no fights
# ($2001DE), 2 = arcade (lower thresholds). Menu showed 0 as "Also On".
# - Every $DF00.w reference -> $F004.w (free RAM). The menu reaches it via (a2,d7)+$EA6 with
#   a2=$D048, d7=$12 -> new constant $F004-$D048-$12 = $1FAA.
# - Option count 3 -> 2 (both count tables) and value 0 labelled "Off": menu is Off / On.
if 'nofight' not in opts:
    for a,old in [(0x2001DE,'4a78df00'),(0x200AF0,'31c0df00')]:
        put(a, bytes.fromhex(old[:4])+W(0xF004), bytes.fromhex(old))
    for a in (0x2002C2,0x20043E,0x2004B4,0x200802):
        put(a, W(0x0C78,2,0xF004), bytes.fromhex('0c780002df00'))
    for a in (0x2095E0,0x209600,0x209624):
        put(a, W(0x0647,0x1FAA), bytes.fromhex('06470ea6'))
    for a in (0x2095C0,0x2095D4):
        put(a, W(2), bytes.fromhex('0003'))
    put(0x20956E, b'Off'+b' '*17, b'Also On'+b' '*13)
# ---- Fix 8: Fighting shows "N/A" in the menu when Penalties is Off (fights need penalties on, see $2001DE/$D056) ----
# The menu's value printer ($F7BAA: jsr $20961A) already prints "N/A" ($F7CEE) for Per. Length in shootouts.
# Hook that call: after it, option $12 (Fighting) with Penalties ($D056) = 0 → print "N/A" the same way.
HOOK8=0x2FE000   # $2FE000-: small menu hooks ($2FF400- is used by the season project)
if 'nofightna' not in opts and 'nofight' not in opts:
    assert all(b==0xff for b in d[HOOK8:HOOK8+0x40])
    put(HOOK8, asm(HOOK8,[
        W(0x4EB9)+L(0x20961A),        # jsr $20961A          (original)
        W(0x0C47,0x0012),             # cmpi.w #$12,d7       Fighting row?
        ('bra',0x6600,'out'),         # bne out
        W(0x4A78,0xD056),             # tst.w $D056.w        Penalties
        ('bra',0x6600,'out'),         # bne out
        W(0x588F),                    # addq.l #4,a7         drop our return address
        W(0x227C)+L(0xF7CEE),         # movea.l #$F7CEE,a1   "N/A"
        W(0x4EF9)+L(0x11BA4),         # jmp $11BA4           print it (returns to the menu loop)
        ('lbl','out'),
        W(0x4E75)]))                  # rts
    put(0xF7BAA, W(0x4EB9)+L(HOOK8), bytes.fromhex('4eb90020961a'))
# ---- Fix 9: Y on a 6-button pad takes your goalie (Goalie Ctrl Manual, your team without the puck) ----
# The hack's 6-button reader turns Y into B ($1FFC02). ygoalie.s (assembled at $2FE040 by build.sh) keeps that when
# your team has the puck (Y = pass) and otherwise asks the input handler ($B258) to switch to the goalie the way
# the game's own manual-goalie switch does ($B3E4: $B86A, then $C0BC/$C0DA).
YG=0x2FE040
if 'noygoalie' not in opts:
    yb=open(os.path.join(os.path.dirname(os.path.abspath(__file__)),'ygoalie.bin'),'rb').read()
    assert all(b==0xff for b in d[YG:YG+len(yb)]) and len(yb)<0x1C0
    d[YG:YG+len(yb)]=yb
    put(0x1FFC02, W(0x4EF9)+L(YG)+W(0x4E71), bytes.fromhex('080200016600ffe8'))
    put(0xB258, W(0x4EF9)+L(YG+0xC0), bytes.fromhex('4ef90020011a'))
# ---- Fix 7: rosters (rosters/<ABBR>.json) ----
# Team table: 32 longs at $30E -> team record. Record header 6 words: +0 offset to players ($92),
# +2 palette ($0C), +4 offset to name strings (players end + 2, preceded by word $0002), +6 lines ($52),
# +8/+10 small fields. Lines: 8 x 8 bytes {G, LD, RD, LW, C, RW, extra, 0}, 1-based player indices;
# order = ?, Sc1, Sc2, Chk, PP1, PP2, PK1, PK2 (verified in the in-game Line Editor).
# Player: word len (2 + name, padded even), name, BCD jersey, 7 bytes = 14 rating nibbles.
# Players are goalies, forwards, defense; keep the team's G/F/D counts (bitmask at +$4C not decoded).
# The rebuilt record goes to free ROM and the team pointer is moved there.
import json
ROSTER_BASE=0x2F0000
if 'noroster' not in opts:
    rdir=os.path.join(os.path.dirname(os.path.abspath(__file__)),'rosters')
    ORIGP=[struct.unpack('>I',d[0x30E+4*t:0x312+4*t])[0] for t in range(32)]
    tptr=lambda t: ORIGP[t]
    def team_players(t):
        a=tptr(t); h=struct.unpack('>6H',d[a:a+12]); i=a+h[0]; out=[]
        while i<a+h[2]-2:
            n=struct.unpack('>H',d[i:i+2])[0]; out.append((d[i+2:i+n].rstrip(b'\x00').decode('latin1'),bytes(d[i+n:i+n+8]))); i+=n+8
        return out
    allp={}
    for t in range(32):
        for nm,b in team_players(t): allp.setdefault(nm,b)
    pos=ROSTER_BASE
    for fn in sorted(os.listdir(rdir)) if os.path.isdir(rdir) else []:
        R=json.load(open(os.path.join(rdir,fn)))
        if 'team' not in R: continue          # e.g. overrides.json (read by gen.py)
        t=R['team']; a=tptr(t)
        h=list(struct.unpack('>6H',d[a:a+12])); old=team_players(t)
        cnt=lambda pl,g: sum(1 for _,b in pl if b[4:6]==b'\x00\x00' and (b[3]&0x0f)==0) # unused, informational
        plist=R['goalies']+R['forwards']+R['defense']
        assert len(plist)==len(old),(fn,len(plist),len(old))
        body=b''; idx={}
        for k,(nm,jer,src) in enumerate(plist):
            idx[nm]=k+1
            rat=allp[nm][1:8] if src=='rom' else bytes.fromhex(src)
            assert len(rat)==7,(nm,src)
            nb=nm.encode('latin1');
            if len(nb)%2: nb+=b'\x00'
            body+=W(2+len(nb))+nb+bytes([int(jer,16)])+rat
        lines=b''.join(bytes([idx[n] for n in ln]+[0]) for ln in R['lines'])
        assert len(lines)==64
        names_old=a+h[2]
        nxt=min([ORIGP[x] for x in range(32) if ORIGP[x]>a]+[names_old+0x100])
        tail=bytes(d[names_old-2:nxt])            # $0002 + names + rest of record
        head=bytearray(d[a:a+h[0]]); head[h[3]:h[3]+64]=lines
        rec=bytes(head)+body+tail
        h2=h[0]+len(body)+2; rec=rec[:4]+W(h2)+rec[6:]
        assert all(b==0xff for b in d[pos:pos+len(rec)+2]),'roster space not free'
        d[pos:pos+len(rec)]=rec
        d[0x30E+4*t:0x312+4*t]=L(pos)
        print('roster',fn,'team',t,'->',hex(pos),len(rec),'bytes')
        pos=(pos+len(rec)+0xF)&~0xF
# ---- fix header checksum ----
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('ok checksum',hex(s))