#!/usr/bin/env python3
"""
extract_doe_region.py -- spool the DOE-PERSIST-TRIAL binary ring off
disk/artemis.img into a CSV, for the multiuser/multitasking DoE campaign
report (FABRIC-3.md §XXXV.10).

Reads the SAME devblock_from_top-addressed ring doe_region.c writes
(include/starkernel/doe_region.h has the authoritative layout: control
header at devblock_from_top=98, slot devblocks growable from 99, 64-byte
doe_trial_slot_t records, 64 slots/devblock). Two access modes:

  --img PATH        Read the image file directly (default). No privilege
                     needed, works on a plain regular file. This is what
                     "reading the drive" reduces to when the image is
                     just a file, which is the normal QEMU-dev case.

  --loop PATH        Attach PATH via losetup first (needs sudo), read the
                     resulting /dev/loopN node, then detach it again on
                     exit -- the literal "mount the virtual drive on a
                     loop device" step for when this runs against a real
                     block device or a copy pulled off real hardware. The
                     ring is raw binary at a fixed byte offset, not a
                     filesystem -- there is nothing to `mount` in the
                     POSIX-filesystem sense, so this attaches the loop
                     device and reads it exactly like a regular file.

Either mode is read-only and never modifies the image. Safe to run
against a live campaign's artemis.img (the ring's control header CRC
means a torn read mid-write is simply rejected as "ring not yet
initialized" by the same discipline doe_region.c itself uses -- see
doe_region_ctrl_read()); worst case is missing the very latest record
until the next run of this script.
"""

import argparse
import csv
import struct
import subprocess
import sys
from contextlib import contextmanager
from pathlib import Path

DEVBLOCK_SIZE = 4096

# Mirrors include/starkernel/doe_region.h and include/starkernel/log_region.h
ARTEMIS_SIG_DEVBLOCK_FROM_TOP = 64
LOG_REGION_DEVBLOCK_FROM_TOP_BASE = ARTEMIS_SIG_DEVBLOCK_FROM_TOP + 1  # 65
LOG_REGION_MAX_DEVBLOCKS = 32
DOE_REGION_DEVBLOCK_FROM_TOP_BASE = (
    LOG_REGION_DEVBLOCK_FROM_TOP_BASE + 1 + LOG_REGION_MAX_DEVBLOCKS
)  # 98

DOE_SLOT_SIZE = 64
DOE_SLOTS_PER_DEVBLOCK = DEVBLOCK_SIZE // DOE_SLOT_SIZE  # 64
DOE_REGION_MAGIC = 0x44454F44  # 'DOED'
DOE_REGION_VERSION_0 = 0

# doe_region_ctrl_t: magic(u64) devblocks(u32) head_slot(u32) tail_slot(u32)
# record_count(u32) hdr_crc(u64), little-endian, rest is padding to 4096.
CTRL_FMT = "<QIIIIQ"
CTRL_SIZE = struct.calcsize(CTRL_FMT)

# doe_trial_slot_t: fleet_k_q48(u64) switch_count_cumulative(u64)
# run_id(u32) cfg(u32) rep(u32) nw(u32) mode(u32) fail_count(u32)
# fleet_conserved(u32) isa(8s), rest is padding to 64.
SLOT_FMT = "<QQIIIIIII8s"
SLOT_SIZE = struct.calcsize(SLOT_FMT)
assert SLOT_SIZE <= DOE_SLOT_SIZE


def dft_to_offset(devblock_from_top: int, total_devblocks: int) -> int:
    idx = total_devblocks - 1 - devblock_from_top
    if idx < 0:
        raise ValueError(
            f"devblock_from_top={devblock_from_top} exceeds total_devblocks={total_devblocks}"
        )
    return idx * DEVBLOCK_SIZE


@contextmanager
def loop_attach(img_path: Path):
    """Attach img_path via losetup (read-only), yield the /dev/loopN path,
    detach on exit. Requires sudo -- the literal 'mount on a loop device'
    step; nothing here is a filesystem mount."""
    result = subprocess.run(
        ["sudo", "losetup", "--find", "--show", "--read-only", str(img_path)],
        capture_output=True, text=True, check=True,
    )
    loop_dev = result.stdout.strip()
    print(f"  attached {img_path} -> {loop_dev}", file=sys.stderr)
    try:
        yield loop_dev
    finally:
        subprocess.run(["sudo", "losetup", "--detach", loop_dev], check=False)
        print(f"  detached {loop_dev}", file=sys.stderr)


def read_ring(dev_path: str):
    """Returns a list of dicts, one per live trial record, in ring order
    (oldest first). Empty list if the ring has never been written."""
    size = Path(dev_path).stat().st_size if not dev_path.startswith("/dev/") else None
    if size is None:
        with open(dev_path, "rb") as f:
            f.seek(0, 2)
            size = f.tell()

    total_devblocks = size // DEVBLOCK_SIZE
    if total_devblocks <= DOE_REGION_DEVBLOCK_FROM_TOP_BASE:
        raise ValueError(
            f"image too small ({size} bytes, {total_devblocks} devblocks) "
            f"for the DOE region fence (needs > {DOE_REGION_DEVBLOCK_FROM_TOP_BASE})"
        )

    ctrl_off = dft_to_offset(DOE_REGION_DEVBLOCK_FROM_TOP_BASE, total_devblocks)
    with open(dev_path, "rb") as f:
        f.seek(ctrl_off)
        ctrl_raw = f.read(DEVBLOCK_SIZE)

    magic, devblocks, head_slot, tail_slot, record_count, hdr_crc = struct.unpack_from(
        CTRL_FMT, ctrl_raw, 0
    )
    magic_lo = magic & 0xFFFFFFFF
    version = (magic >> 32) & 0xFF
    if magic_lo != DOE_REGION_MAGIC or version != DOE_REGION_VERSION_0:
        print("  DOE region not yet initialized (no trials persisted yet)", file=sys.stderr)
        return []

    total_slots = devblocks * DOE_SLOTS_PER_DEVBLOCK
    records = []
    with open(dev_path, "rb") as f:
        idx = head_slot
        for _ in range(record_count):
            slot_devblock_dft = DOE_REGION_DEVBLOCK_FROM_TOP_BASE + 1 + (idx // DOE_SLOTS_PER_DEVBLOCK)
            slot_off_in_devblock = (idx % DOE_SLOTS_PER_DEVBLOCK) * DOE_SLOT_SIZE
            blk_off = dft_to_offset(slot_devblock_dft, total_devblocks)
            f.seek(blk_off + slot_off_in_devblock)
            raw = f.read(SLOT_SIZE)
            (fleet_k_q48, switch_count, run_id, cfg, rep, nw, mode,
             fail_count, fleet_conserved, isa_raw) = struct.unpack(SLOT_FMT, raw)
            isa = isa_raw.split(b"\x00", 1)[0].decode("ascii", errors="replace")
            records.append({
                "isa": isa,
                "run_id": run_id,
                "cfg": cfg,
                "rep": rep,
                "nw": nw,
                "mode": mode,
                "fail_count": fail_count,
                "fleet_k_q48": fleet_k_q48,
                "fleet_conserved": fleet_conserved,
                "switch_count_cumulative": switch_count,
            })
            idx = (idx + 1) % total_slots

    print(f"  ring: devblocks={devblocks} head={head_slot} tail={tail_slot} "
          f"record_count={record_count} -> {len(records)} records read", file=sys.stderr)
    return records


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--img", default="disk/artemis.img", help="Path to the disk image (default: disk/artemis.img)")
    ap.add_argument("--loop", action="store_true", help="Attach --img via losetup first (needs sudo)")
    ap.add_argument("-o", "--out", default=None, help="Output CSV path (default: stdout)")
    args = ap.parse_args()

    img_path = Path(args.img)
    if not img_path.exists():
        print(f"extract_doe_region: image not found: {img_path}", file=sys.stderr)
        sys.exit(1)

    if args.loop:
        with loop_attach(img_path) as loop_dev:
            records = read_ring(loop_dev)
    else:
        records = read_ring(str(img_path))

    fieldnames = ["isa", "run_id", "cfg", "rep", "nw", "mode", "fail_count",
                  "fleet_k_q48", "fleet_conserved", "switch_count_cumulative"]

    out_f = open(args.out, "w", newline="") if args.out else sys.stdout
    try:
        writer = csv.DictWriter(out_f, fieldnames=fieldnames)
        writer.writeheader()
        for rec in records:
            writer.writerow(rec)
    finally:
        if args.out:
            out_f.close()
            print(f"  wrote {len(records)} rows -> {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()