#!/usr/bin/env python3
# traceCode.py - recursive-descent trace of the game from known entry points over the runtime images;
# writes code coverage per image and reports indirect jumps / illegal opcodes / unreached areas.
import sys, json
from m6502 import trace, decode

LOADED = {
  "boot":    [(0x0400,0x8800),(0x8C00,0x9000),(0x9300,0xC400),(0xFBB8,0xFFD2)],
  "runtime": [(0x0200,0x8A00),(0x8C00,0x9000),(0x9300,0xD000),(0xE000,0xF700),(0xFBB8,0xFFD2)],
}
def validFor(ranges):
    return lambda a: any(lo <= a < hi for lo, hi in ranges)

def entriesFor(mem, extra):
    e = {0x0800}
    e |= set(extra)
    return e

def run(imageName, ranges, extra, noReturn):
    mem = open(imageName, "rb").read()
    code, calls, jumps, indirects, illegal = trace(mem, entriesFor(mem, extra), validFor(ranges), noReturn=noReturn)
    return mem, code, calls, jumps, indirects, illegal

def report(tag, mem, code, calls, jumps, indirects, illegal, ranges):
    print(f"== {tag}: code bytes {len(code)}, subroutines {len(calls)}, jump targets {len(jumps)}, indirect jmps {len(indirects)}, illegal {len(illegal)}")
    for a in indirects:
        i = decode(mem, a)
        print(f"   jmp (${i.operand:04X}) at ${a:04X}")
    for a in illegal[:20]:
        print(f"   illegal opcode ${mem[a]:02X} at ${a:04X}")
    # unreached runs inside loaded ranges
    total = sum(hi-lo for lo,hi in ranges)
    print(f"   coverage {len(code)}/{total} = {100*len(code)/total:.1f}%")

if __name__ == "__main__":
    extra = [int(x,16) for x in sys.argv[1:]]
    results = {}
    for tag, img, rng in [("boot","../extracted/bootImage.bin",LOADED["boot"]),
                          ("A","../extracted/runtimeImage_A.bin",LOADED["runtime"]),
                          ("B","../extracted/runtimeImage_B.bin",LOADED["runtime"])]:
        mem, code, calls, jumps, ind, ill = run(img, rng, extra, frozenset())
        report(tag, mem, code, calls, jumps, ind, ill, rng)
        results[tag] = sorted(code)
    json.dump({k:[f"{a:04X}" for a in v] for k,v in results.items()}, open("../extracted/traceCode.json","w"))
