#!/usr/bin/env bash
# Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
# See LICENSE.md for license information.
# SPDX-License-Identifier: MIT
#
# profile-iigs.sh - PC-sampling profiler for a JoeyLib program on the IIgs.
#
# Boots a self-booting program volume (make-iigs-game-disk.sh's) under
# headless MAME and samples the 65816's program bank and counter every
# PROFILE_PERIOD_US of EMULATED time, from a Lua coroutine that emu.wait()s
# between samples. A per-frame sample would land at the same point of a
# frame-synced game loop every time; a period that is not a multiple of the
# frame's spreads the samples over the whole frame.
#
# The OMF carries no symbols and the GS/OS Loader puts each bank-aligned
# segment in whatever bank it likes, so the samples are tied back to the
# link map by the segments' BYTES: at the end the script reads the start of
# every bank a sample landed in and scores it against every segment image
# (the Loader's relocations change only a few bytes). A bank that matches
# no segment is GS/OS, a tool, NTP or the ROM, and is reported by bank.
#
# The map comes from building the program with JOEY_LINK_MAP=<dir>, which
# keeps <dir>/NAME.map, NAME.seg<N>.bin and NAME.json (clang-build.sh):
#
#   rm -f build/iigs/bin/STAXI && JOEY_LINK_MAP=/some/dir make iigs
#
# The map lists static functions too (JOEY_LINK_MAP links with link816's
# --map-locals), but one the compiler inlined has no code of its own: its
# time is its caller's row.
#
# Usage:   scripts/profile-iigs.sh <image> <map> [seconds] [skipSeconds]
#            image        the program's boot volume: .po (800KB, the 3.5
#                         drive) or .2mg (anything larger, a CFFA2 card)
#            map          NAME.map, with NAME.seg<N>.bin beside it
#            seconds      emulated seconds to run, default 120
#            skipSeconds  ignore samples before this, default 40 (boot)
# Env:     PROFILE_PERIOD_US  sample period in emulated microseconds,
#                             default 997
#          PROFILE_TOP        rows to print, default 40
#          PROFILE_RAMSIZE    MAME -ramsize, default 4M
#          PROFILE_SAMPLES    a file to keep the raw samples in (S addr
#                             count, F bank bytes, T total), for a look
#                             inside a bank no segment matched
# Requires: ./install.sh install iigs.
set -euo pipefail

img=${1:?usage: profile-iigs.sh <image> <map> [seconds] [skipSeconds]}
map=${2:?usage: profile-iigs.sh <image> <map> [seconds] [skipSeconds]}
seconds=${3:-120}
skip=${4:-40}
periodUs=${PROFILE_PERIOD_US:-997}
top=${PROFILE_TOP:-40}
ramsize=${PROFILE_RAMSIZE:-4M}

repo=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
source "$repo/scripts/lib/joeyEnv.sh"
joeyNeed MAME mame
joeyNeed IIGS_ROMPATH iigsFirmware

[ -f "$img" ] || { echo "profile-iigs: no such image $img" >&2; exit 2; }
[ -f "$map" ] || { echo "profile-iigs: no such map $map" >&2; exit 2; }
[ -f "${map%.map}.seg1.bin" ] || { echo "profile-iigs: ${map%.map}.seg1.bin missing (build with JOEY_LINK_MAP=<dir>)" >&2; exit 2; }

# MAME opens a hard-disk image read-write: run a private copy.
work=$(mktemp -d -t joeylib-iigsprof.XXXXXX)
trap 'rm -rf "$work"' EXIT
name=$(basename "$img")
cp "$img" "$work/$name"
if [ "$(stat -c %s "$img")" -le 900000 ]; then
    media=(-flop3 "$work/$name")
else
    media=(-sl7 cffa2 -hard1 "$work/$name")
fi

# Bytes of a bank compared against a segment's image.
fingerprintBytes=512

cat > "$work/prof.lua" <<LUA
local cpu = manager.machine.devices[":maincpu"]
local mem = cpu.spaces["program"]
local out = io.open("samples.txt", "w")
local counts = {}
local banks = {}
local taken = 0

local function finish()
    for addr, n in pairs(counts) do
        out:write(string.format("S %06X %d\n", addr, n))
    end
    for bank, _ in pairs(banks) do
        local bytes = {}
        for i = 0, $fingerprintBytes - 1 do
            bytes[#bytes + 1] = string.format("%02X", mem:read_u8((bank << 16) + i))
        end
        out:write(string.format("F %02X %s\n", bank, table.concat(bytes)))
    end
    out:write(string.format("T %d\n", taken))
    out:close()
    manager.machine:exit()
end

local sampler = coroutine.create(function()
    while manager.machine.time:as_double() < $seconds do
        emu.wait($periodUs / 1000000.0)
        if manager.machine.time:as_double() >= $skip then
            local pb = cpu.state["PB"].value
            local addr = (pb << 16) | cpu.state["PC"].value
            counts[addr] = (counts[addr] or 0) + 1
            banks[pb] = true
            taken = taken + 1
        end
    end
    finish()
end)

local started = false
emu.register_frame_done(function()
    if not started then
        started = true
        local ok, err = coroutine.resume(sampler)
        if not ok then
            io.write("PROFILE-IIGS lua error: " .. tostring(err) .. "\n")
            manager.machine:exit()
        end
    end
end)
LUA

echo "profile-iigs: $name, ${seconds}s emulated (samples from ${skip}s, every ${periodUs} us)..." >&2
(cd "$work" && QT_QPA_PLATFORM=offscreen SDL_VIDEODRIVER=dummy SDL_AUDIODRIVER=dummy \
    timeout 3600 "$MAME" apple2gs \
        -ramsize "$ramsize" -rompath "$IIGS_ROMPATH" "${media[@]}" \
        -video none -sound none -nothrottle -skip_gameinfo \
        -autoboot_script "$work/prof.lua" </dev/null >"$work/mame.log" 2>&1) || true
if [ ! -s "$work/samples.txt" ]; then
    echo "profile-iigs: FAIL -- no samples (MAME output follows)" >&2
    tail -15 "$work/mame.log" >&2
    exit 1
fi
grep 'PROFILE-IIGS' "$work/mame.log" >&2 || true
if [ -n "${PROFILE_SAMPLES:-}" ]; then
    cp "$work/samples.txt" "$PROFILE_SAMPLES"
fi

python3 - "$work/samples.txt" "$map" "$top" <<'PY'
import bisect
import json
import os
import re
import sys

samplesPath, mapPath, top = sys.argv[1], sys.argv[2], int(sys.argv[3])
mapBase = mapPath[:-4] if mapPath.endswith('.map') else mapPath

# The segments: number -> (link bank, image bytes).
segments = {}
manifest = mapBase + '.json'
if os.path.exists(manifest):
    for seg in json.load(open(manifest))['segments']:
        segments[seg['num']] = int(seg['base'], 16) >> 16
else:
    segments[1] = 0
images = {num: open('%s.seg%d.bin' % (mapBase, num), 'rb').read() for num in segments}

# The symbols, by address: the globals, and the locals (static functions)
# when the map has them; a global wins an address both name.
byAddr = {}
section = None
for line in open(mapPath):
    if line.startswith('#'):
        section = 'global' if line.startswith('# global symbols') else ('local' if line.startswith('# local symbols') else None)
        continue
    m = re.match(r'^0x([0-9a-fA-F]{6})\s+(\S+)$', line.strip())
    if m and section is not None:
        addr = int(m.group(1), 16)
        if section == 'global' or addr not in byAddr:
            byAddr[addr] = m.group(2)
syms = sorted(byAddr.items())
symAddrs = [a for a, _ in syms]

counts = {}
fingerprints = {}
total = 0
for line in open(samplesPath):
    parts = line.split()
    if parts[0] == 'S':
        counts[int(parts[1], 16)] = int(parts[2])
    elif parts[0] == 'F':
        fingerprints[int(parts[1], 16)] = bytes.fromhex(parts[2])
    elif parts[0] == 'T':
        total = int(parts[1])

# Runtime bank -> link bank, by the best-scoring segment image.
bankMap = {}
for bank, fp in fingerprints.items():
    best = None
    for num, image in images.items():
        n = min(len(fp), len(image))
        if n < 16:
            continue
        same = sum(1 for i in range(n) if fp[i] == image[i])
        score = same / n
        if best is None or score > best[0]:
            best = (score, num)
    if best is not None and best[0] >= 0.85:
        bankMap[bank] = (segments[best[1]], best[1], best[0])

def where(addr):
    bank = addr >> 16
    if bank in bankMap:
        link = (bankMap[bank][0] << 16) | (addr & 0xFFFF)
        i = bisect.bisect_right(symAddrs, link) - 1
        if i >= 0 and (symAddrs[i] >> 16) == (link >> 16):
            return syms[i][1]
        return 'seg%d+%04X' % (bankMap[bank][1], addr & 0xFFFF)
    if bank >= 0xFC:
        return 'ROM $%02X' % bank
    if bank in (0xE0, 0xE1):
        return 'bank $%02X (firmware/vectors)' % bank
    return 'bank $%02X (GS/OS, tools, NTP)' % bank

byName = {}
for addr, n in counts.items():
    key = where(addr)
    byName[key] = byName.get(key, 0) + n

print('profile-iigs: %d samples' % total)
for bank in sorted(bankMap):
    link, num, score = bankMap[bank]
    print('  bank $%02X = segment %d (link bank $%02X, %.0f%% of %d bytes match)' % (bank, num, link, score * 100, len(fingerprints[bank])))
print('%7s %6s  %s' % ('samples', 'pct', 'where'))
for key, n in sorted(byName.items(), key=lambda kv: -kv[1])[:top]:
    print('%7d %5.1f%%  %s' % (n, 100.0 * n / max(total, 1), key))
PY