gps-denied-onboard/e2e/fixtures/sitl_replay_builder/build_p01_fixtures.py

"""FT-P-01 fixture builder (AZ-598).

Produces:

* ``outbound_messages_<fc_kind>_<host>.json`` — per-image SUT outbound GPS
  estimates, in image-order. ``null`` entries encode per-image timeouts.
* ``observer_<fc_kind>_<host>.json`` — minimal observer config so
  ``sitl_observer.get_observer`` succeeds when the fixtures are activated.

Strategy: drive the production ``gps-denied-replay`` CLI against a 1 fps
MP4 encoded from the FT-P-01 still-image set and a synthetic stationary
tlog, then read the resulting FDR JSONL for per-frame outbound estimates.
Compared with the rejected "live SITL docker capture" path this:

* Adds no new SUT-side frame-ingestion code (reuses
  ``ReplayInputAdapter`` + ``VideoFileFrameSource``).
* Bypasses the SITL container entirely (FT-P-01 tests upstream
  geo-estimate accuracy; the FC is just a delivery channel).
* Runs as a single subprocess instead of a multi-container compose.

The helpers below are intentionally dependency-injectable so the unit
tests can mock OpenCV / pymavlink / subprocess / filesystem without
touching real hardware or libraries.
"""

from __future__ import annotations

import argparse
import json
import logging
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Iterable, Sequence

from e2e.fixtures.sitl_replay_builder._common import (
    DEFAULT_CLI_BIN,
    run_gps_denied_replay,
    write_observer_fixture,
)

_LOG = logging.getLogger(__name__)

DEFAULT_FPS = 1.0
DEFAULT_TLOG_DURATION_S = 120
DEFAULT_TLOG_HZ = 200
DEFAULT_FDR_KIND = "outbound_position_estimate"


@dataclass(frozen=True)
class BuilderConfig:
    """Per-invocation builder configuration."""

    input_dir: Path
    output_dir: Path
    fc_kind: str
    host: str
    fps: float = DEFAULT_FPS
    tlog_duration_s: int = DEFAULT_TLOG_DURATION_S
    tlog_hz: int = DEFAULT_TLOG_HZ
    fdr_kind: str = DEFAULT_FDR_KIND
    cli_bin: str = DEFAULT_CLI_BIN


# Step 1 — encode the still images into a 1 fps MP4


def encode_stills_to_mp4(
    image_paths: Sequence[Path],
    output_mp4: Path,
    *,
    fps: float = DEFAULT_FPS,
    _video_writer_factory: Callable | None = None,
    _imread: Callable | None = None,
) -> int:
    """Encode `image_paths` (in order) as an MP4 at `fps`. Returns frame count.

    Raises ``FileNotFoundError`` when no image paths are supplied or when
    any input image cannot be read.

    The OpenCV dependencies are injected via the underscore-prefixed
    parameters so unit tests can run without OpenCV being available.
    """
    if not image_paths:
        raise FileNotFoundError(
            "encode_stills_to_mp4: image_paths is empty; nothing to encode"
        )

    if _video_writer_factory is None or _imread is None:
        import cv2

        _imread = _imread or (lambda path: cv2.imread(str(path), cv2.IMREAD_COLOR))
        if _video_writer_factory is None:
            _fourcc = cv2.VideoWriter_fourcc(*"mp4v")

            def _video_writer_factory(out: Path, width: int, height: int):
                return cv2.VideoWriter(str(out), _fourcc, fps, (width, height))

    first_frame = _imread(image_paths[0])
    if first_frame is None:
        raise FileNotFoundError(
            f"encode_stills_to_mp4: failed to read {image_paths[0]}"
        )
    height, width = first_frame.shape[:2]
    output_mp4.parent.mkdir(parents=True, exist_ok=True)

    writer = _video_writer_factory(output_mp4, width, height)
    try:
        writer.write(first_frame)
        for path in image_paths[1:]:
            frame = _imread(path)
            if frame is None:
                raise FileNotFoundError(
                    f"encode_stills_to_mp4: failed to read {path}"
                )
            writer.write(frame)
    finally:
        writer.release()

    return len(image_paths)


# Step 2 — generate a synthetic stationary tlog


def generate_stationary_tlog(
    output_tlog: Path,
    *,
    duration_s: int = DEFAULT_TLOG_DURATION_S,
    hz: int = DEFAULT_TLOG_HZ,
    _mavlink_writer_factory: Callable | None = None,
) -> int:
    """Write a tlog with `duration_s * hz` stationary RAW_IMU + ATTITUDE pairs.

    The output is the minimum tlog content ``ReplayInputAdapter`` requires:
    monotonic-timestamp RAW_IMU + ATTITUDE messages so the AZ-405 tlog
    pre-validator (`AC-13`) doesn't reject the input.

    The samples encode zero accel/gyro/attitude — auto-sync will refuse to
    find a take-off, so callers MUST drive ``gps-denied-replay`` with an
    explicit ``--time-offset-ms 0`` to bypass auto-sync.

    Returns the number of message PAIRS written.
    """
    if duration_s <= 0:
        raise ValueError(f"duration_s must be positive; got {duration_s}")
    if hz <= 0:
        raise ValueError(f"hz must be positive; got {hz}")

    if _mavlink_writer_factory is None:
        from pymavlink import mavutil

        def _mavlink_writer_factory(out: Path):
            return mavutil.mavlogfile(str(out), write=True)

    output_tlog.parent.mkdir(parents=True, exist_ok=True)
    pairs = 0
    writer = _mavlink_writer_factory(output_tlog)
    try:
        period_us = int(1_000_000 / hz)
        total_pairs = duration_s * hz
        for i in range(total_pairs):
            time_us = i * period_us
            writer.write(_pack_raw_imu_zero(time_us))
            writer.write(_pack_attitude_zero(time_us // 1000))
            pairs += 1
    finally:
        close = getattr(writer, "close", None)
        if callable(close):
            close()

    return pairs


def _pack_raw_imu_zero(time_usec: int) -> bytes:
    """Pack a zero-motion RAW_IMU MAVLink frame (msg id 27).

    Constructed with pymavlink's MAVLink2 packer so the produced bytes are
    a wire-compatible MAVLink frame including header + CRC. Stationary
    semantics: all accel/gyro/mag fields are zero except the Z accel which
    carries one g (gravity, ~9.81 m/s² × 1000 in mg).
    """
    from pymavlink.dialects.v20 import ardupilotmega as mavlink

    packer = mavlink.MAVLink(file=None, srcSystem=1, srcComponent=1)
    msg = mavlink.MAVLink_raw_imu_message(
        time_usec=time_usec,
        xacc=0,
        yacc=0,
        zacc=-9810,
        xgyro=0,
        ygyro=0,
        zgyro=0,
        xmag=0,
        ymag=0,
        zmag=0,
        id=0,
        temperature=0,
    )
    return msg.pack(packer)


def _pack_attitude_zero(time_boot_ms: int) -> bytes:
    """Pack a zero-motion ATTITUDE MAVLink frame (msg id 30)."""
    from pymavlink.dialects.v20 import ardupilotmega as mavlink

    packer = mavlink.MAVLink(file=None, srcSystem=1, srcComponent=1)
    msg = mavlink.MAVLink_attitude_message(
        time_boot_ms=time_boot_ms,
        roll=0.0,
        pitch=0.0,
        yaw=0.0,
        rollspeed=0.0,
        pitchspeed=0.0,
        yawspeed=0.0,
    )
    return msg.pack(packer)


# Step 3 — drive `gps-denied-replay` against the generated video+tlog
# (`run_gps_denied_replay` is re-exported from `_common.py` so b78 + b79 share one impl.)


# Step 4 — extract per-frame outbound estimates from the FDR JSONL


def parse_fdr_for_outbound_estimates(
    fdr_path: Path,
    *,
    fdr_kind: str = DEFAULT_FDR_KIND,
    lat_key: str = "lat_deg",
    lon_key: str = "lon_deg",
) -> list[dict]:
    """Walk `fdr_path` (JSONL) and return outbound-estimate payloads in order.

    A record contributes one entry when its ``kind`` matches `fdr_kind` AND
    its payload carries both `lat_key` and `lon_key`. Other records are
    silently skipped (the FDR carries many record types per the
    `_docs/02_document/contracts/fdr/` schema). Malformed JSON lines raise
    ``ValueError`` with the line number.
    """
    if not fdr_path.is_file():
        raise FileNotFoundError(f"FDR JSONL not found: {fdr_path}")

    out: list[dict] = []
    with fdr_path.open("r", encoding="utf-8") as fp:
        for line_no, line in enumerate(fp, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(
                    f"malformed FDR JSON at {fdr_path}:{line_no}: {exc.msg}"
                ) from exc
            if record.get("kind") != fdr_kind:
                continue
            payload = record.get("payload", {})
            if not isinstance(payload, dict):
                continue
            if lat_key not in payload or lon_key not in payload:
                continue
            out.append(
                {
                    "lat_deg": float(payload[lat_key]),
                    "lon_deg": float(payload[lon_key]),
                }
            )
    return out


# Step 5 — write the two fixture files in the b75/b78 schema


def write_outbound_messages_fixture(
    output_path: Path,
    image_ids: Sequence[str],
    estimates: Sequence[dict | None],
) -> None:
    """Write `outbound_messages_<fc_kind>_<host>.json`.

    `image_ids` and `estimates` must have the same length. `None` entries
    in `estimates` are persisted as JSON `null` (timeout markers); other
    entries must carry `lat_deg`/`lon_deg`.
    """
    if len(image_ids) != len(estimates):
        raise ValueError(
            f"length mismatch: {len(image_ids)} image_ids vs "
            f"{len(estimates)} estimates"
        )
    messages: list[dict | None] = []
    for image_id, estimate in zip(image_ids, estimates):
        if estimate is None:
            messages.append(None)
            continue
        messages.append(
            {
                "image_id": image_id,
                "lat_deg": float(estimate["lat_deg"]),
                "lon_deg": float(estimate["lon_deg"]),
            }
        )
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps({"messages": messages}, indent=2))


# `write_observer_fixture` is re-exported from `_common.py` (used by both b78 + b79).


# Orchestration


def _resolve_p01_image_paths(input_dir: Path) -> list[Path]:
    """Return the AD0000NN.jpg images under `input_dir`, sorted by name."""
    if not input_dir.is_dir():
        raise FileNotFoundError(f"input dir not found: {input_dir}")
    return sorted(input_dir.glob("AD??????.jpg"))


def build_p01_fixtures(
    cfg: BuilderConfig,
    *,
    _runner: Callable[[Sequence[str]], subprocess.CompletedProcess] | None = None,
    _video_writer_factory: Callable | None = None,
    _imread: Callable | None = None,
    _mavlink_writer_factory: Callable | None = None,
) -> Path:
    """End-to-end FT-P-01 fixture build. Returns the output directory.

    Steps (matches the module docstring):

    1. Resolve the 60 AD0000NN.jpg images under ``cfg.input_dir``.
    2. Encode them at ``cfg.fps`` into ``stills.mp4`` under ``cfg.output_dir``.
    3. Generate a stationary ``stationary.tlog`` under ``cfg.output_dir``.
    4. Run ``gps-denied-replay`` against the pair; write FDR JSONL.
    5. Project FDR outbound-estimate records into the two fixture files.

    Per-frame timeout handling: if the FDR yields fewer estimates than
    images, the trailing image_ids get `null` (timeout) entries. If the
    FDR yields MORE estimates than images (multiple emissions per frame),
    only the first ``len(image_paths)`` estimates are kept and a WARN is
    logged so the operator notices the schema mismatch.
    """
    image_paths = _resolve_p01_image_paths(cfg.input_dir)
    if not image_paths:
        raise FileNotFoundError(
            f"no AD??????.jpg images found under {cfg.input_dir}"
        )

    cfg.output_dir.mkdir(parents=True, exist_ok=True)
    stills_mp4 = cfg.output_dir / "stills.mp4"
    stationary_tlog = cfg.output_dir / "stationary.tlog"
    fdr_jsonl = cfg.output_dir / "fdr.jsonl"

    encode_stills_to_mp4(
        image_paths, stills_mp4, fps=cfg.fps,
        _video_writer_factory=_video_writer_factory, _imread=_imread,
    )
    generate_stationary_tlog(
        stationary_tlog,
        duration_s=cfg.tlog_duration_s,
        hz=cfg.tlog_hz,
        _mavlink_writer_factory=_mavlink_writer_factory,
    )
    run_gps_denied_replay(
        stills_mp4, stationary_tlog, fdr_jsonl,
        cli_bin=cfg.cli_bin, _runner=_runner,
    )

    raw_estimates = parse_fdr_for_outbound_estimates(fdr_jsonl, fdr_kind=cfg.fdr_kind)
    estimates: list[dict | None] = list(raw_estimates[: len(image_paths)])
    if len(raw_estimates) > len(image_paths):
        _LOG.warning(
            "FDR carried %d outbound estimates but only %d images were pushed; "
            "truncating to the per-frame count", len(raw_estimates), len(image_paths)
        )
    while len(estimates) < len(image_paths):
        estimates.append(None)

    outbound_path = cfg.output_dir / f"outbound_messages_{cfg.fc_kind}_{cfg.host}.json"
    observer_path = cfg.output_dir / f"observer_{cfg.fc_kind}_{cfg.host}.json"
    write_outbound_messages_fixture(
        outbound_path,
        image_ids=[p.name for p in image_paths],
        estimates=estimates,
    )
    write_observer_fixture(observer_path)
    return cfg.output_dir


def _main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        prog="build_p01_fixtures",
        description="Build FT-P-01 SITL replay fixtures via gps-denied-replay.",
    )
    parser.add_argument("--input-dir", type=Path, required=True,
                        help="Directory containing AD000001..AD000060.jpg")
    parser.add_argument("--output-dir", type=Path, required=True,
                        help="Output dir for stills.mp4 + stationary.tlog + fixtures")
    parser.add_argument("--fc-kind", choices=("ardupilot", "inav"), default="ardupilot")
    parser.add_argument("--host", default="sitl-host")
    parser.add_argument("--fps", type=float, default=DEFAULT_FPS)
    parser.add_argument("--cli-bin", default=DEFAULT_CLI_BIN)
    args = parser.parse_args(argv)

    logging.basicConfig(level=logging.INFO)
    cfg = BuilderConfig(
        input_dir=args.input_dir,
        output_dir=args.output_dir,
        fc_kind=args.fc_kind,
        host=args.host,
        fps=args.fps,
        cli_bin=args.cli_bin,
    )
    build_p01_fixtures(cfg)
    return 0


if __name__ == "__main__":  # pragma: no cover
    sys.exit(_main())