"""Run inside IDA (via ida-mcp execute_python) to dump address-invariant
per-function signatures for cross-version binary diffing.

Writes psp_out/sigs/<input-stem>.json
"""
import idaapi, idc, idautils, ida_funcs, ida_bytes, ida_auto, ida_segment
import ida_gdl, hashlib, json, os


def _prep():
    seg = ida_segment.getnseg(0)
    # SOS carries an ARM prologue/vector region at 0x0-0x3c0; rest is Thumb.
    if seg is not None and seg.start_ea == 0:
        try:
            idc.split_sreg_range(0x0, 'T', 0, idc.SR_user)
            idc.split_sreg_range(0x3C0, 'T', 1, idc.SR_user)
        except Exception as e:
            print('sreg:', e)
    ida_auto.auto_wait()


def _sig(fn):
    toks, ncall = [], 0
    ea = fn.start_ea
    while ea < fn.end_ea:
        if ida_bytes.is_code(ida_bytes.get_flags(ea)):
            m = idc.print_insn_mnem(ea)
            if m:
                parts = [m]
                for oi in range(3):
                    t = idc.get_operand_type(ea, oi)
                    if t == 0:
                        break
                    if t == 1:
                        parts.append('r%d' % idc.get_operand_value(ea, oi))
                    elif t == 5:
                        v = idc.get_operand_value(ea, oi)
                        parts.append('i%x' % v if abs(v) < 0x10000 else 'IMMBIG')
                    elif t in (2, 3, 4):
                        parts.append('MEM')
                    elif t in (6, 7):
                        parts.append('ADDR')
                    else:
                        parts.append('t%d' % t)
                toks.append('|'.join(parts))
                if m.lower().startswith(('bl', 'blx')):
                    ncall += 1
        ea = idc.next_head(ea, fn.end_ea)
    try:
        nbb = len(list(ida_gdl.FlowChart(fn)))
    except Exception:
        nbb = 0
    return toks, nbb, ncall


def run():
    _prep()
    out = {}
    for f in idautils.Functions():
        fn = ida_funcs.get_func(f)
        toks, nbb, ncall = _sig(fn)
        out['%x' % fn.start_ea] = {
            'size': fn.end_ea - fn.start_ea,
            'n': len(toks),
            'bb': nbb,
            'calls': ncall,
            'h': hashlib.sha256('\n'.join(toks).encode()).hexdigest()[:16],
        }
    stem = os.path.basename(idc.get_input_file_path()).rsplit('.', 1)[0]
    d = '/Users/user/dev/amd-psp-re/psp_out/sigs'
    os.makedirs(d, exist_ok=True)
    p = os.path.join(d, stem + '.json')
    json.dump(out, open(p, 'w'))
    print('%s: %d funcs -> %s' % (stem, len(out), p))
    return len(out)


run()