#!/usr/bin/env python3
"""Two-process smoke test for snapshot_link D→P RDMA byte transfer.

Spawns scripts/snapshot_link_receiver.py via subprocess.Popen with stderr
piped to ``<tmpdir>/recv.stderr.log`` for post-mortem if something dies.

Sender (this process):
    1. Spawns receiver child, waits for endpoint.json
    2. Brings up own SnapshotPeer (no recv buffer), registers a send buffer
    3. For each size: fill pattern, batch_transfer_sync_write, signal child,
       wait for child's ack
    4. Reads child's stdout (one JSON event per line) for verification

Pass = every size yields a child "verify" event with ok=true.

Usage:
    bash scripts/setup_env.sh && uv run --no-sync python scripts/smoke_snapshot_link.py

Env (optional):
    SNAPSHOT_LINK_HOST       default 127.0.0.1
    SNAPSHOT_LINK_IB         default mlx5_60
    SNAPSHOT_LINK_RECV_PORT  default 17777
    SNAPSHOT_LINK_SEND_PORT  default 17778
"""

from __future__ import annotations

import argparse
import ctypes
import hashlib
import json
import os
import subprocess
import sys
import tempfile
import time
from pathlib import Path

_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE.parent / "src"))


SIZES_BYTES_DEFAULT = [
    1 << 10,   # 1 KB
    1 << 14,   # 16 KB
    1 << 18,   # 256 KB
    1 << 20,   # 1 MB
    1 << 22,   # 4 MB
    1 << 24,   # 16 MB
    1 << 26,   # 64 MB
]


def _pattern_byte(i: int, seed: int) -> int:
    return (i * 2654435761 + seed) & 0xFF


def _fill_pattern(buf, length: int, seed: int) -> None:
    tile_size = 4096
    tile = bytes(_pattern_byte(i, seed) for i in range(tile_size))
    tile_arr = (ctypes.c_ubyte * tile_size).from_buffer_copy(tile)
    n_full = length // tile_size
    rem = length - n_full * tile_size
    base = ctypes.addressof(buf)
    src_addr = ctypes.addressof(tile_arr)
    for k in range(n_full):
        ctypes.memmove(base + k * tile_size, src_addr, tile_size)
    if rem:
        ctypes.memmove(base + n_full * tile_size, src_addr, rem)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default=os.environ.get("SNAPSHOT_LINK_HOST", "127.0.0.1"))
    ap.add_argument("--ib", default=os.environ.get("SNAPSHOT_LINK_IB", "mlx5_60"))
    ap.add_argument("--recv-port", type=int,
                    default=int(os.environ.get("SNAPSHOT_LINK_RECV_PORT", "17777")))
    ap.add_argument("--send-port", type=int,
                    default=int(os.environ.get("SNAPSHOT_LINK_SEND_PORT", "17778")))
    ap.add_argument("--max-bytes", type=int, default=128 * 1024 * 1024)
    ap.add_argument("--sizes", default=",".join(str(s) for s in SIZES_BYTES_DEFAULT))
    args = ap.parse_args()

    sizes = [int(s) for s in args.sizes.split(",")]
    tmpdir = Path(tempfile.mkdtemp(prefix="snapshot_link_smoke_"))
    control_path = tmpdir / "endpoint.json"
    recv_stderr_log = tmpdir / "recv.stderr.log"

    recv_cmd = [
        sys.executable,
        str(_HERE / "snapshot_link_receiver.py"),
        "--host", args.host,
        "--port", str(args.recv_port),
        "--ib", args.ib,
        "--max-bytes", str(args.max_bytes),
        "--control-path", str(control_path),
        "--sizes", args.sizes,
    ]
    recv_stderr = open(recv_stderr_log, "w")
    print(f"[sender] launching receiver: {' '.join(recv_cmd)}", flush=True)
    print(f"[sender] receiver stderr → {recv_stderr_log}", flush=True)
    recv_proc = subprocess.Popen(
        recv_cmd,
        stdout=subprocess.PIPE,
        stderr=recv_stderr,
        bufsize=1,
        universal_newlines=True,
    )

    try:
        # Wait for endpoint metadata
        deadline = time.time() + 60.0
        while time.time() < deadline:
            if control_path.exists():
                try:
                    meta = json.loads(control_path.read_text())
                    if meta.get("ready"):
                        break
                except Exception:
                    pass
            if recv_proc.poll() is not None:
                _dump_recv_stderr(recv_stderr_log)
                print(f"[sender] FAIL: receiver exited early (rc={recv_proc.returncode})")
                return 1
            time.sleep(0.1)
        else:
            print("[sender] FAIL: timed out waiting for receiver endpoint", flush=True)
            return 1

        print(f"[sender] receiver endpoint: {meta}", flush=True)

        from agentic_pd_hybrid.snapshot_link import SnapshotPeer, SnapshotEndpoint
        endpoint = SnapshotEndpoint(
            session_id=meta["session_id"],
            base_ptr=int(meta["base_ptr"]),
            capacity_bytes=int(meta["capacity_bytes"]),
        )
        peer = SnapshotPeer(
            host=args.host,
            port=args.send_port,
            ib_device=args.ib,
            receive_capacity_bytes=0,
        )
        send_buf = (ctypes.c_byte * args.max_bytes)()
        send_addr = ctypes.addressof(send_buf)
        peer.register_send_buffer(send_addr, args.max_bytes)
        print(f"[sender] own session_id={peer.session_id}, send_buf @ {hex(send_addr)} ({args.max_bytes} B)", flush=True)

        transfers = []
        for size in sizes:
            if size > args.max_bytes:
                continue
            seed = int(time.time() * 1e6) & 0xFFFFFFFF
            _fill_pattern(send_buf, size, seed)
            t0 = time.perf_counter()
            ret = peer.push(endpoint, send_addr, 0, size, remote_offset=0)
            t1 = time.perf_counter()
            dt_ms = (t1 - t0) * 1000.0
            gbps = (size * 8.0 / 1e9) / max(t1 - t0, 1e-9)
            print(f"[sender] push size={size:>10d}  ret={ret}  "
                  f"dur={dt_ms:>9.3f} ms  thru={gbps:>6.3f} Gbps",
                  flush=True)
            signal_path = control_path.with_suffix(f".do{size}")
            ack_path = control_path.with_suffix(f".ack{size}")
            signal_path.write_text(str(seed))
            ack_deadline = time.time() + 60.0
            while time.time() < ack_deadline:
                if ack_path.exists():
                    break
                if recv_proc.poll() is not None:
                    print(f"[sender] FAIL: receiver died after size={size}", flush=True)
                    _dump_recv_stderr(recv_stderr_log)
                    return 1
                time.sleep(0.05)
            transfers.append({
                "size": size, "ret": ret, "dur_ms": round(dt_ms, 3),
                "thru_Gbps": round(gbps, 3),
                "ack": ack_path.exists(),
            })

        peer.close()

        # Drain child stdout — each line is a JSON event
        try:
            recv_proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            recv_proc.terminate()
            recv_proc.wait(timeout=5)

        events = []
        if recv_proc.stdout is not None:
            for raw in recv_proc.stdout:
                raw = raw.strip()
                if not raw:
                    continue
                try:
                    events.append(json.loads(raw))
                except json.JSONDecodeError:
                    events.append({"event": "non-json", "raw": raw})

        print("=" * 78)
        print("[receiver] events:")
        verify_ok = 0
        verify_fail = 0
        for ev in events:
            print(f"  {ev}")
            if ev.get("event") == "verify":
                if ev.get("ok"):
                    verify_ok += 1
                else:
                    verify_fail += 1

        recv_stderr.close()
        _dump_recv_stderr(recv_stderr_log, header="--- receiver stderr ---")

        overall = "PASS" if verify_fail == 0 and verify_ok == len(transfers) else "FAIL"
        print("=" * 78)
        print(f"OVERALL: {overall}  verify_ok={verify_ok}  verify_fail={verify_fail}  "
              f"transfers={len(transfers)}")
        return 0 if overall == "PASS" else 1

    finally:
        try:
            recv_proc.terminate()
            recv_proc.wait(timeout=5)
        except Exception:
            try:
                recv_proc.kill()
            except Exception:
                pass


def _dump_recv_stderr(path: Path, header: str = "--- receiver stderr (last 40) ---") -> None:
    try:
        text = path.read_text()
    except FileNotFoundError:
        return
    print(header, flush=True)
    for line in text.splitlines()[-40:]:
        print(f"  {line}", flush=True)


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