#!/usr/bin/env python3
"""Capture streaming ETM instruction trace unattended via J-Trace + Ozone.

Generates a throwaway Ozone project (never touches the committed
hw/bsp/**/ozone/*.jdebug), launches Ozone on a virtual display (xvfb-run),
drives the whole session over Ozone's automation TCP socket (UM08025 §6.7):
connect -> flash -> run for --duration-ms under streaming trace -> halt ->
export. Outputs in --out:
  code_profile.txt   hot functions (run/fetch counts) + code coverage
  itrace.csv         raw instruction history (only with --trace-csv)
  ozone_console.log, jlink.log, ozone_gui.log   session evidence

Requires firmware built with -DTRACE_ETM=1 and the board wired to a J-Trace.
Analyze results with etm_profile.py in this directory.
"""

import argparse
import glob
import os
import re
import shutil
import signal
import socket
import string
import subprocess
import sys
import tempfile
import time

REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), *[".."] * 4))

# Throwaway Ozone project. SP/PC-from-vector-table hooks match the committed
# reference projects (Cortex-M generic). $$(InstallDir) -> literal $(InstallDir).
PROJECT_TEMPLATE = string.Template("""\
/* Auto-generated by etm_capture.py (etm-trace skill) - throwaway, do not commit. */
void OnProjectLoad (void) {
  Project.SetDevice ("$device");
  Project.SetHostIF ("USB", "$probe");
  Project.SetTargetIF ("SWD");
  Project.SetTIFSpeed ("$tif_speed");
  Project.SetTraceSource ("Trace Pins");
  Project.SetTracePortWidth ($port_width);
$timing_line$core_clock_line$timestamps_line$hss_line$power_lines  Edit.SysVar (VAR_TRACE_MAX_INST_CNT, $max_inst);
  Edit.Preference (PREF_TIMESTAMP_FORMAT, TIMESTAMP_FORMAT_TIME);
  Project.AddSvdFile ("$$(InstallDir)/Config/Peripherals/ARMv7M.svd");
  Project.SetConsoleLogFile ("$outdir/ozone_console.log");
  Project.SetJLinkLogFile ("$outdir/jlink.log");
$os_plugin_line  File.Open ("$elf");
}
$jlink_script_hook

$reset_hook

$download_hook
$user_funcs""")

# Cortex-M generic SP/PC-from-vector-table init. A committed board reference
# overrides these verbatim (e.g. RT1176 apps in FlexSPI NOR need the ROM
# bootloader to do SP/PC init, plus a JTAG_nTRST pad fix).
DEFAULT_HOOK = """\
void %s (void) {
  unsigned int SP;
  unsigned int PC;
  unsigned int VectorTableAddr;
  VectorTableAddr = Elf.GetBaseAddr();
  if (VectorTableAddr == 0xFFFFFFFF) {
    Util.Log("etm_capture: failed to get program base");
  } else {
    SP = Target.ReadU32(VectorTableAddr);
    Target.SetReg("SP", SP);
    PC = Target.ReadU32(VectorTableAddr + 4);
    Target.SetReg("PC", PC);
  }
}"""

# Symbolic constants (TP_OP_*, EXPORT_*) only evaluate in project-script
# context, never over the automation socket - so these live in generated user
# functions invoked via Script.Exec (UM08025 SS6.7, SS7.9.9.1).
TRACEPOINT_FUNC = """\

void SetupTracepoints (void) {
%s}
"""

LINES_CSV_FUNC = """\

void ExportLinesCsv (void) {
  Export.CodeProfile ("%s", EXPORT_AS_CSV | EXPORT_CSV_LINES | EXPORT_FILE_PATHS, "");
}
"""

INSTS_CSV_FUNC = """\

void ExportInstsCsv (void) {
  Export.CodeProfile ("%s", EXPORT_AS_CSV | EXPORT_CSV_INSTS | EXPORT_FILE_PATHS, "");
}
"""


JLINK_SCRIPT_HOOK = """\

void BeforeTargetConnect (void) {
  Project.SetJLinkScript ("%s");
}
"""


def resolve_board(board):
    """Board config from its committed ozone reference project, else board.cmake/mk."""
    cfg = {"device": None, "tif_speed": "4 MHz", "timing": None, "port_width": 4,
           "core_clock": None, "ref": None}
    jdebugs = sorted(glob.glob(f"{REPO_ROOT}/hw/bsp/*/boards/{board}/ozone/*.jdebug"))
    if jdebugs:
        cfg["ref"] = jdebugs[0]
        text = open(jdebugs[0]).read()
        # config regexes must not match //-commented lines
        cfgtext = re.sub(r"^\s*//.*$", "", text, flags=re.M)
        # inherit every user function verbatim except OnProjectLoad, which the
        # template owns (e.g. STM32H5's AfterTargetConnect must clock the trace
        # CoreSight domain; RT1176 replaces the SP/PC reset/download hooks for
        # ROM-bootloader boot; Nordic hooks call a _SetupTarget helper that
        # must ride along or hook execution fails at runtime)
        extra = []
        for m in re.finditer(r"^void (\w+)\s*\(void\)\s*\{.*?^\}", text,
                             re.M | re.S):
            name, block = m.group(1), m.group(0)
            if name == "OnProjectLoad":
                continue
            elif name == "BeforeTargetConnect":
                # the generated project synthesizes its own BeforeTargetConnect
                # (JLINK_SCRIPT_HOOK) from the SetJLinkScript regex below;
                # inheriting the reference's copy too would emit a duplicate
                # function definition
                continue
            elif name == "AfterTargetReset":
                cfg["reset_hook"] = block
            elif name == "AfterTargetDownload":
                cfg["download_hook"] = block
            else:
                extra.append(block)
        if extra:
            cfg["connect_hook"] = "\n\n".join(extra) + "\n"
        for key, pat in (("device", r'Project\.SetDevice\s*\(\s*"([^"]+)"'),
                         ("tif_speed", r'Project\.SetTIFSpeed\s*\(\s*"([^"]+)"'),
                         ("timing", r'Project\.SetTraceTiming\s*\(([-\d\s,]+)\)'),
                         ("port_width", r'Project\.SetTracePortWidth\s*\(\s*(\d+)'),
                         ("core_clock", r'VAR_TRACE_CORE_CLOCK\s*,\s*(\d+)')):
            m = re.search(pat, cfgtext)
            if m:
                cfg[key] = m.group(1).strip()
        # inherit a J-Link script (e.g. RT1176 must declare its off-ROM-table
        # TPIU/funnel); relative paths resolve against the reference's dir
        m = re.search(r'Project\.SetJLinkScript\s*\(\s*"([^"]+)"', cfgtext)
        if m:
            # $(ProjectDir) = the reference's own directory
            rel = m.group(1).replace("$(ProjectDir)", ".")
            cfg["jlink_script"] = os.path.normpath(os.path.join(
                os.path.dirname(jdebugs[0]), rel))
    else:
        for path in glob.glob(f"{REPO_ROOT}/hw/bsp/*/boards/{board}/board.cmake"):
            m = re.search(r'JLINK_DEVICE\s+([^\s)]+)\s*\)', open(path).read())
            if m:
                if "${" in m.group(1):
                    sys.exit(f"error: {path} defines JLINK_DEVICE via an "
                             f"unexpanded CMake variable ({m.group(1)}) - pass "
                             f"--device explicitly for this board")
                cfg["device"] = m.group(1)
                cfg["ref"] = path
                break
    if not cfg["device"]:
        sys.exit(f"error: cannot resolve J-Link device for board '{board}' "
                 f"(no hw/bsp/*/boards/{board}/ozone/*.jdebug or board.cmake)")
    return cfg


def trace_only_points(elf, syms_arg):
    """Tracepoint lines for --trace-only: start trace at each symbol's entry,
    stop at each return instruction inside it (pop ...pc / bx lr, via objdump).
    Hardware comparators are scarce (ETM-M7) - keep the symbol list short."""
    lines = ""
    nm = subprocess.run(["arm-none-eabi-nm", "-S", "--defined-only", elf],
                        capture_output=True, text=True).stdout
    for want in [s.strip() for s in syms_arg.split(",") if s.strip()]:
        m = re.search(rf"^([0-9a-f]+) ([0-9a-f]+) [TtWw] {re.escape(want)}$",
                      nm, re.M)
        if not m:
            sys.exit(f"error: --trace-only symbol '{want}' not in ELF")
        lo, sz = int(m.group(1), 16) & ~1, int(m.group(2), 16)
        lines += f'  Trace.SetPoint (TP_OP_START_TRACE, "{want}");\n'
        dis = subprocess.run(
            ["arm-none-eabi-objdump", "-d", f"--start-address={lo:#x}",
             f"--stop-address={lo + sz:#x}", elf],
            capture_output=True, text=True).stdout
        exits = re.findall(
            r"^\s*([0-9a-f]+):.*?(?:(?:pop|ldmia[.\w]*\s+sp!,)[^\n]*\bpc\b|bx\s+lr)",
            dis, re.M | re.I)
        if not exits:
            sys.exit(f"error: no return instruction found in '{want}'")
        for addr in exits:
            lines += f'  Trace.SetPoint (TP_OP_STOP_TRACE, "0x{int(addr, 16):08X}");\n'
    return lines


def resolve_probe(probe):
    """Ozone's SetHostIF needs a serial - with several probes connected a
    nickname makes it block on a selection dialog. JLinkExe DOES resolve
    nicknames, so borrow its banner to map nickname -> serial. The serial only
    ever lands in the throwaway project file, never in committed files."""
    if not probe or probe.isdigit():
        return probe
    r = subprocess.run(["JLinkExe", "-USB", probe, "-nogui", "1"],
                       input="qc\n", capture_output=True, text=True, timeout=30)
    m = re.search(r"S/N:\s*(\d+)", r.stdout)
    if not m:
        sys.exit(f"error: cannot resolve probe nickname '{probe}' to a serial "
                 f"(JLinkExe -USB {probe} found no emulator)")
    return m.group(1)


def gen_project(cfg, args, outdir):
    timing = ""
    if args.trace_timing is not None:
        d = [int(v) for v in str(args.trace_timing).split(",")]
        if len(d) not in (1, 4):
            sys.exit("error: --trace-timing takes one value or d0,d1,d2,d3")
        d = d * 4 if len(d) == 1 else d
        timing = ("  Project.SetTraceTiming "
                  f"({d[0]}, {d[1]}, {d[2]}, {d[3]});\n")
    elif cfg["timing"]:
        timing = f"  Project.SetTraceTiming ({cfg['timing']});\n"
    core_clock = args.core_clock or cfg["core_clock"]
    clk_line = f"  Edit.SysVar (VAR_TRACE_CORE_CLOCK, {core_clock});\n" if core_clock else ""
    ts_line = ("  Edit.SysVar (VAR_TRACE_TIMESTAMPS_ENABLED, 0);\n"
               if args.no_timestamps else "")
    user_funcs = ""
    if cfg.get("connect_hook"):
        user_funcs += "\n" + cfg["connect_hook"]
    if args.trace_only:
        user_funcs += TRACEPOINT_FUNC % trace_only_points(
            os.path.abspath(args.elf), args.trace_only)
    if args.profile_lines_csv:
        user_funcs += LINES_CSV_FUNC % os.path.join(outdir, "profile_lines.csv")
    if args.profile_insts_csv:
        user_funcs += INSTS_CSV_FUNC % os.path.join(outdir, "profile_insts.csv")
    if args.os_plugin and not glob.glob(
            f"/opt/SEGGER/Ozone*/Plugins/OS/{args.os_plugin}.js"):
        sys.exit(f"error: RTOS plugin '{args.os_plugin}' not in "
                 f"/opt/SEGGER/Ozone*/Plugins/OS (e.g. FreeRTOSPlugin_CM7)")
    os_plugin = (f'  Project.SetOSPlugin ("{args.os_plugin}");\n'
                 if args.os_plugin else "")
    if args.attach:
        # attach to the running target: no download, no reset - trace a window
        # mid-run (firmware must match the ELF and have trace pins enabled)
        os_plugin += "  Debug.SetConnectMode (CM_ATTACH_HALT);\n"
    # HSS sampling rate belongs in OnProjectLoad (persistent) per UM08025 4.6.2;
    # setting it mid-session over the socket is rejected.
    hss = (f"  Edit.SysVar (VAR_HSS_SPEED, {args.sample_hz});\n"
           if args.sample else "")
    power = ("  Edit.SysVar (VAR_TARGET_POWER_ON, 1);\n"
             f"  Edit.SysVar (VAR_POWER_SAMPLING_SPEED, {args.power_hz});\n"
             if args.power else "")
    jlink_script = args.jlink_script or cfg.get("jlink_script")
    jls_hook = (JLINK_SCRIPT_HOOK % os.path.abspath(jlink_script)
                if jlink_script else "")
    proj = os.path.join(outdir, "etm_capture.jdebug")
    if args.trace_width:
        cfg["port_width"] = args.trace_width
    with open(proj, "w") as f:
        f.write(PROJECT_TEMPLATE.substitute(
            device=cfg["device"], probe=resolve_probe(args.probe),
            tif_speed=cfg["tif_speed"],
            port_width=cfg["port_width"], timing_line=timing, core_clock_line=clk_line,
            timestamps_line=ts_line, max_inst=args.max_inst, outdir=outdir,
            elf=os.path.abspath(args.elf), user_funcs=user_funcs,
            os_plugin_line=os_plugin, hss_line=hss, power_lines=power,
            jlink_script_hook=jls_hook,
            reset_hook=cfg.get("reset_hook", DEFAULT_HOOK % "AfterTargetReset"),
            download_hook=cfg.get("download_hook",
                                  DEFAULT_HOOK % "AfterTargetDownload")))
    return proj


class OzoneSession:
    """Drive Ozone via its automation TCP socket (one connection at a time)."""

    def __init__(self, port, logf):
        self.port = port
        self.logf = logf
        self.sock = None

    def log(self, msg):
        line = f"[{time.strftime('%H:%M:%S')}] {msg}"
        print(line, flush=True)
        self.logf.write(line + "\n")
        self.logf.flush()

    def connect(self, timeout_s):
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            try:
                self.sock = socket.create_connection(("127.0.0.1", self.port), timeout=5)
                self.sock.settimeout(0.5)
                self.log(f"connected to Ozone automation socket :{self.port}")
                return
            except OSError:
                time.sleep(1)
        raise TimeoutError(f"Ozone automation socket :{self.port} not reachable "
                           f"after {timeout_s}s (see ozone_gui.log)")

    def drain(self, wait_s=1.0):
        buf = b""
        end = time.time() + wait_s
        while time.time() < end:
            try:
                chunk = self.sock.recv(65536)
                if not chunk:
                    break
                buf += chunk
                end = time.time() + 0.5  # keep reading while data flows
            except socket.timeout:
                pass
        text = buf.decode(errors="replace")
        for ln in text.splitlines():
            self.log(f"  ozone> {ln}")
        return text

    def send(self, cmd, wait_s=1.0):
        self.log(f"cmd: {cmd}")
        self.sock.sendall((cmd + "\n").encode())
        return self.drain(wait_s)

    def wait_echo(self, cmd, timeout_s):
        """Send cmd; wait until Ozone echoes its execution (echo comes after the
        command completed, e.g. a large Export). Returns all received text."""
        name = cmd.split("(")[0].strip()
        text = self.send(cmd, 1.0)
        deadline = time.time() + timeout_s
        while name + " (" not in text and name + "(" not in text:
            if time.time() > deadline:
                raise TimeoutError(f"no echo for '{name}' after {timeout_s}s")
            text += self.drain(1.0)
        return text

    def is_halted(self, timeout_s):
        """Poll Debug.IsHalted until it returns 1; parse the '// returns 0xN' echo."""
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            text = self.send("Debug.IsHalted", 1.0)
            end2 = time.time() + 4.0
            while "Debug.IsHalted" not in text and time.time() < end2:
                text += self.drain(1.0)
            m = re.findall(r"Debug\.IsHalted\s*\(\);\s*//\s*returns\s*0x(\d+)", text)
            if m and m[-1] == "1":
                return True
            time.sleep(1)
        return False


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--board", help="TinyUSB board name, e.g. stm32h743eval")
    p.add_argument("--device", help="J-Link device name for non-TinyUSB targets "
                   "(e.g. STM32F407VE for a SEGGER trace reference board); "
                   "requires --elf, bypasses hw/bsp board resolution")
    p.add_argument("--tif-speed", default="4 MHz",
                   help="SWD speed for --device targets (default '4 MHz')")
    p.add_argument("--jlink-script",
                   help="J-Link script file (.pex/.JLinkScript) for trace-pin "
                        "init when the firmware doesn't do it (SEGGER per-MCU "
                        "examples); wired into BeforeTargetConnect")
    p.add_argument("--power", action="store_true",
                   help="power the target from the probe and record a power "
                        "profile (power.csv); probe power is switched off after "
                        "the session. Target must be wired for probe power!")
    p.add_argument("--power-hz", type=int, default=10000,
                   help="power sampling frequency in Hz (default 10000)")
    p.add_argument("--elf", help="firmware ELF built with -DTRACE_ETM=1 "
                   "(default: examples/cmake-build-<board>/device/cdc_msc/cdc_msc.elf)")
    p.add_argument("--duration-ms", type=int, default=10000, help="traced run time")
    p.add_argument("--out", help="output dir (default: mkdtemp under /tmp)")
    p.add_argument("--port", type=int, default=19201,
                   help="automation socket port (19200 = interactive Ozone default; keep 19201)")
    p.add_argument("--probe", default="",
                   help="J-Link USB nickname or serial ('' = sole connected probe)")
    p.add_argument("--trace-csv", action="store_true",
                   help="also export raw instruction history (itrace.csv, can be >100 MB)")
    p.add_argument("--max-inst", type=int, default=10000000,
                   help="VAR_TRACE_MAX_INST_CNT: instruction-trace window/export depth")
    p.add_argument("--core-clock", type=int,
                   help="CPU Hz for timestamp conversion (default: board reference value)")
    p.add_argument("--trace-timing",
                   help="trace sample delay in ps (-5000..5000, overrides the "
                        "board reference; sweep this when the stream dies with "
                        "unknown-packet errors). One value for all 4 pins, or "
                        "'d0,d1,d2,d3' to de-skew individual lines (boards can "
                        "have per-line RC delays, e.g. strap pulls on muxed pads)")
    p.add_argument("--trace-width", type=int, choices=(1, 2, 4),
                   help="trace port width override (fewer pins = tolerant of a "
                        "single bad line, at reduced bandwidth)")
    p.add_argument("--no-timestamps", action="store_true",
                   help="disable trace timestamps (less trace bandwidth -> fewer "
                        "overflows/decode errors; itrace.csv loses its time column)")
    p.add_argument("--trace-only",
                   help="comma-separated symbols: trace ONLY these functions via "
                        "hardware tracepoints (e.g. an ISR + SysTick_Handler for "
                        "calibration); needs few symbols (scarce comparators)")
    p.add_argument("--profile-lines-csv", action="store_true",
                   help="also export per-source-line profile/coverage counters "
                        "(profile_lines.csv)")
    p.add_argument("--profile-insts-csv", action="store_true",
                   help="also export per-instruction counters (profile_insts.csv, "
                        "enables branch-bias analysis in etm_profile.py)")
    p.add_argument("--attach", action="store_true",
                   help="attach to the RUNNING target instead of flash+reset: "
                        "capture a window mid-run (narrowing debug). The flashed "
                        "firmware must match --elf and have been built with "
                        "TRACE_ETM=1")
    p.add_argument("--sample",
                   help="comma-separated C expressions to sample periodically "
                        "during the run (samples.csv, e.g. 'system_ticks')")
    p.add_argument("--sample-hz", type=int, default=1000,
                   help="data sampling frequency in Hz (default 1000)")
    p.add_argument("--os-plugin",
                   help="Ozone RTOS-awareness plugin for task/ISR-attributed "
                        "timeline, e.g. FreeRTOSPlugin_CM7 (see "
                        "/opt/SEGGER/Ozone*/Plugins/OS)")
    args = p.parse_args()

    if not args.board and not args.device:
        sys.exit("error: need --board (TinyUSB) or --device + --elf (other targets)")
    if args.device and not args.elf:
        sys.exit("error: --device requires --elf")
    if args.max_inst > 10000000:
        # >10M yielded a silently EMPTY Export.Trace on Ozone V3.50
        print("warning: --max-inst clamped to 10000000 (larger values produce "
              "an empty instruction-trace export)", file=sys.stderr)
        args.max_inst = 10000000
    if not args.elf:
        args.elf = (f"{REPO_ROOT}/examples/cmake-build-{args.board}"
                    f"/device/cdc_msc/cdc_msc.elf")
    if not os.path.isfile(args.elf):
        sys.exit(f"error: ELF not found: {args.elf}\n"
                 f"build it with -DTRACE_ETM=1 (see the etm-trace skill) or pass --elf")

    if args.trace_only:
        print("warning: --trace-only is EXPERIMENTAL - on some targets windows "
              "for low-rate handlers fail to record, and timestamps are invalid "
              "across trace gaps (instruction counts remain exact). For ISR "
              "timing prefer a full --trace-csv capture + etm_profile.py --isr.",
              file=sys.stderr)
    if args.device:
        cfg = {"device": args.device, "tif_speed": args.tif_speed, "timing": None,
               "port_width": 4, "core_clock": None, "ref": "--device"}
    else:
        cfg = resolve_board(args.board)
    outdir = os.path.abspath(args.out) if args.out else tempfile.mkdtemp(
        prefix=f"etm-{args.board or args.device}-")
    os.makedirs(outdir, exist_ok=True)
    proj = gen_project(cfg, args, outdir)

    ozone_bin = shutil.which("ozone") or shutil.which("Ozone")
    if not ozone_bin:
        sys.exit("error: ozone not on PATH (install SEGGER Ozone)")
    cmd = [ozone_bin, "-project", proj, "-port", str(args.port)]
    if shutil.which("xvfb-run"):
        cmd = ["xvfb-run", "-a"] + cmd
    elif os.environ.get("DISPLAY"):
        print("warning: xvfb-run not found - Ozone window will appear on "
              f"DISPLAY={os.environ['DISPLAY']} and may steal keyboard focus",
              file=sys.stderr)
    else:
        sys.exit("error: no DISPLAY and no xvfb-run; install xvfb")

    gui_log = open(os.path.join(outdir, "ozone_gui.log"), "w")
    ses_logf = open(os.path.join(outdir, "session.log"), "w")
    ses = OzoneSession(args.port, ses_logf)
    ses.log(f"board={args.board} device={cfg['device']} ref={cfg['ref']}")
    ses.log(f"elf={args.elf}")
    ses.log(f"out={outdir}")
    proc = subprocess.Popen(cmd, stdout=gui_log, stderr=gui_log,
                            start_new_session=True)
    profile_out = os.path.join(outdir, "code_profile.txt")
    itrace_out = os.path.join(outdir, "itrace.csv")
    try:
        ses.connect(20)
        ses.drain(3)  # version banner
        ses.send("Debug.Start", 5)
        if not ses.is_halted(90):
            raise TimeoutError("Debug.Start did not reach the startup completion "
                               "point (connect/flash failed? see jlink.log)")
        ses.log("startup complete (halted at main)")
        if args.trace_only:
            ses.wait_echo('Script.Exec ("SetupTracepoints")', 15)
        ses.send('Window.Show ("Code Profile")', 2)
        if args.trace_csv:
            ses.send('Window.Show ("Instruction Trace")', 2)
        if args.sample:
            ses.send('Window.Show ("Data Sampling")', 2)
            for expr in args.sample.split(","):
                ses.send(f'Window.Add ("Data Sampling", "{expr.strip()}")', 2)
        ses.send("Coverage.ExcludeNOPs()", 2)

        ses.log(f"=== traced run: {args.duration_ms} ms ===")
        ses.send("Debug.Continue", 1)
        time.sleep(args.duration_ms / 1000.0)
        ses.send("Debug.Halt", 3)
        if not ses.is_halted(30):
            raise TimeoutError("target did not halt")
        ses.send("Window.WaitForUpdateComplete(120000)", 5)

        ses.wait_echo(f'Export.CodeProfile ("{profile_out}", 0, "")', 60)
        if args.profile_lines_csv:
            ses.wait_echo('Script.Exec ("ExportLinesCsv")', 60)
        if args.profile_insts_csv:
            ses.wait_echo('Script.Exec ("ExportInstsCsv")', 120)
        if args.trace_csv:
            ses.wait_echo(f'Export.Trace ("{itrace_out}", 0)', 300)
        if args.sample:
            ses.wait_echo(f'Export.DataGraphs ("{outdir}/samples.csv")', 60)
        if args.power:
            ses.wait_echo(f'Export.PowerGraphs ("{outdir}/power.csv")', 60)
        ses.send("Debug.Stop", 5)
        ses.send("File.Exit", 2)
    finally:
        for _ in range(15):
            if proc.poll() is not None:
                break
            time.sleep(1)
        if proc.poll() is None:
            ses.log("killing leftover Ozone process group")
            os.killpg(proc.pid, signal.SIGKILL)
        if args.power:
            # guarantee probe power is off, whatever happened above
            sel = ["-USB", args.probe] if args.probe else []
            r = subprocess.run(["JLinkExe", *sel, "-nogui", "1"],
                               input="power off\nqc\n", capture_output=True,
                               text=True, timeout=30)
            ses.log("probe power off " +
                    ("issued" if r.returncode == 0 else f"FAILED rc={r.returncode}"))

    if not (os.path.isfile(profile_out) and os.path.getsize(profile_out) > 0
            and "Code Profile Report" in open(profile_out, errors="replace").read(200)):
        sys.exit(f"error: capture ran but {profile_out} is missing/empty - "
                 f"check {outdir}/session.log and ozone_console.log")
    if args.trace_csv and not (os.path.isfile(itrace_out)
                               and os.path.getsize(itrace_out) > 0):
        sys.exit(f"error: --trace-csv requested but {itrace_out} is missing/empty")
    if args.power and not os.path.getsize(os.path.join(outdir, "power.csv")):
        sys.exit("error: --power requested but power.csv is missing/empty")
    if "Trace collection stopped!" in open(os.path.join(outdir, "session.log"),
                                           errors="replace").read():
        sys.exit("error: trace stream died mid-run (unknown trace data packet) - "
                 "the profile only covers up to that point. Retry with "
                 "--no-timestamps, or reduce the core clock (see SKILL.md).")
    prof = open(profile_out, errors="replace").read()
    m = re.search(r"^\s*Total\s*\|[\d ]*\|\s*([\d ]+)$", prof, re.M)
    if not m or int(m.group(1).replace(" ", "") or 0) == 0:
        sys.exit("error: session completed but NO trace data was collected "
                 "(profile totals are zero) - trace signal not reaching the "
                 "probe: check wiring/connector, trace pinmux, sample timing.")
    print(f"\ncapture OK: {outdir}")
    print(f"  code_profile.txt  ({os.path.getsize(profile_out)} bytes)")
    if args.trace_csv:
        print(f"  itrace.csv        ({os.path.getsize(itrace_out)} bytes)")
    print(f"analyze: python3 {os.path.dirname(os.path.abspath(__file__))}"
          f"/etm_profile.py {outdir} --elf {args.elf}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
