#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# bridge_selftest.py - offline self-test of cuda_bridge --bridge (CUBB2/CBD2 v2)
# using plain files in place of the vgsp.ko debugfs files.
#
# Simulates the module side: packs 256-byte CUBB2 job records at the documented
# natural-alignment offsets, prepares the in blob, waits for the CBD2 done
# record, and checks the out blob:
#   (a) mail  n=39        -> done ok, out bytes exact
#   (b) vadd  n=1<<20     -> out c correct, bad=0, checksum 392344090.500000
#   (c) stale seq         -> ignored (no re-execution)
#   (d) invalid magic     -> ignored (idle)
#   (e) bad opcode        -> done written with non-zero status
#   (f) oversized n       -> done written with non-zero status (bound)
#
#   python3 bridge_selftest.py [--dir D] [--bin PATH]
#
# Never touches /dev/shm/vgpu-guest-ram, the guest, QEMU, mdev or vgsp.ko.

import argparse
import array
import os
import struct
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
MAIL = b"the little gpu man says you've got mail"
VADD_CHECKSUM = 392344090.500000
CHECKSUM_TOL = 1e-4

JOB_MAGIC = 0x32424243      # 'CBB2' LE
DONE_MAGIC = 0x32444243     # 'CBD2' LE
VERSION = 2
JOB_SIZE = 256
DONE_SIZE = 256

ST_OK, ST_BADJOB, ST_CUDA, ST_RAM = 0, 1, 2, 3


def pack_job(seq, opcode, n, in_bytes=0, out_bytes=0, grid=(1, 1, 1),
             block=(64, 1, 1), shared=0, sem_va=0, sem_payload=0,
             sem_flags=0, prog_va=0, param_count=2, magic=JOB_MAGIC,
             version=VERSION):
    """Pack a CUBB2 job record field-by-field at the frozen v2 offsets."""
    b = bytearray(JOB_SIZE)
    struct.pack_into("<IIQ", b, 0x00, magic, version, seq)
    struct.pack_into("<I", b, 0x10, opcode)
    struct.pack_into("<3I", b, 0x14, *grid)
    struct.pack_into("<3I", b, 0x20, *block)
    struct.pack_into("<I", b, 0x2c, shared)
    struct.pack_into("<I", b, 0x30, n)
    struct.pack_into("<Q", b, 0x38, in_bytes)
    struct.pack_into("<Q", b, 0x40, out_bytes)
    struct.pack_into("<Q", b, 0x48, sem_va)
    struct.pack_into("<Q", b, 0x50, sem_payload)
    struct.pack_into("<I", b, 0x58, sem_flags)
    struct.pack_into("<Q", b, 0x60, prog_va)
    struct.pack_into("<I", b, 0x68, param_count)
    assert len(b) == JOB_SIZE
    return bytes(b)


def parse_done(data):
    """Parse a CBD2 done record at the frozen v2 offsets."""
    assert len(data) == DONE_SIZE, "done record must be 256 bytes"
    magic, version, seq = struct.unpack_from("<IIQ", data, 0x00)
    status, bad = struct.unpack_from("<II", data, 0x10)
    (checksum_bits,) = struct.unpack_from("<Q", data, 0x18)
    (kernel_ns,) = struct.unpack_from("<Q", data, 0x20)
    out_bytes_written, err_code = struct.unpack_from("<II", data, 0x28)
    name = data[0x30:0x70].split(b"\0", 1)[0].decode(errors="replace")
    (host_ns,) = struct.unpack_from("<Q", data, 0x70)
    checksum = struct.unpack("<d", struct.pack("<Q", checksum_bits))[0]
    return {
        "magic": magic,
        "version": version,
        "seq": seq,
        "status": status,
        "bad": bad,
        "checksum": checksum,
        "kernel_ns": kernel_ns,
        "out_bytes_written": out_bytes_written,
        "err_code": err_code,
        "device": name,
        "host_ns": host_ns,
    }


class Harness:
    def __init__(self, bin_path, workdir, poll_ms=10):
        self.bin = bin_path
        self.dir = workdir
        self.job = os.path.join(workdir, "job.bin")
        self.inp = os.path.join(workdir, "in.bin")
        self.out = os.path.join(workdir, "out.bin")
        self.done = os.path.join(workdir, "done.bin")
        self.log = os.path.join(workdir, "daemon.log")
        self.poll_ms = poll_ms
        self.proc = None
        self.fails = 0

    def start(self):
        os.makedirs(self.dir, exist_ok=True)
        for p in (self.job, self.inp, self.out, self.done):
            try:
                os.unlink(p)
            except FileNotFoundError:
                pass
        self.logfh = open(self.log, "wb")
        self.proc = subprocess.Popen(
            [self.bin, "--bridge", "--job", self.job, "--in", self.inp,
             "--out", self.out, "--done", self.done, "--poll-ms",
             str(self.poll_ms)],
            stdout=self.logfh, stderr=subprocess.STDOUT)
        deadline = time.time() + 60.0
        while time.time() < deadline:
            if self.proc.poll() is not None:
                print("daemon exited early rc=%s" % self.proc.returncode)
                return False
            try:
                with open(self.log, "r", errors="replace") as fh:
                    text = fh.read()
            except OSError:
                text = ""
            if "BRIDGE v2" in text and "DEVICE " in text:
                return True
            time.sleep(0.1)
        print("daemon did not become ready in 60 s")
        return False

    def stop(self):
        if self.proc and self.proc.poll() is None:
            self.proc.terminate()
            try:
                self.proc.wait(timeout=10)
            except subprocess.TimeoutExpired:
                self.proc.kill()
                self.proc.wait()
        try:
            self.logfh.close()
        except Exception:
            pass

    def log_text(self):
        try:
            with open(self.log, "r", errors="replace") as fh:
                return fh.read()
        except OSError:
            return ""

    def write_job(self, rec):
        tmp = self.job + ".tmp"
        with open(tmp, "wb") as fh:
            fh.write(rec)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(tmp, self.job)

    def write_in(self, data):
        tmp = self.inp + ".tmp"
        with open(tmp, "wb") as fh:
            fh.write(data)
            fh.flush()
            os.fsync(fh.fileno())
        os.replace(tmp, self.inp)

    def wait_done(self, seq, timeout=60.0):
        deadline = time.time() + timeout
        while time.time() < deadline:
            try:
                with open(self.done, "rb") as fh:
                    data = fh.read(DONE_SIZE)
            except FileNotFoundError:
                data = b""
            if len(data) == DONE_SIZE and data[:4] == struct.pack(
                    "<I", DONE_MAGIC) and struct.unpack_from(
                        "<I", data, 4)[0] == VERSION and struct.unpack_from(
                            "<Q", data, 8)[0] == seq:
                return parse_done(data)
            time.sleep(0.01)
        return None

    def check(self, name, ok, detail):
        if not ok:
            self.fails += 1
        print("BRIDGE-SELFTEST %s %s %s" % (name, "PASS" if ok else "FAIL",
                                            detail))
        sys.stdout.flush()

    def run(self):
        if not self.start():
            print("BRIDGE-SELFTEST setup FAIL (daemon)")
            return 1
        print("BRIDGE-SELFTEST device: %s" %
              next((ln for ln in self.log_text().splitlines()
                    if ln.startswith("DEVICE ")), "?"))
        try:
            self.case_mail()
            self.case_vadd()
            self.case_stale()
            self.case_badmagic()
            self.case_failure()
            self.case_bound()
        finally:
            self.stop()
        print("BRIDGE-SELFTEST %s" % ("PASS" if self.fails == 0 else "FAIL"))
        if self.fails:
            print("--- daemon log tail ---")
            print("\n".join(self.log_text().splitlines()[-25:]))
        return 0 if self.fails == 0 else 1

    # ---------------------------------------------------------------- cases
    def case_mail(self):
        n = 39
        self.write_job(pack_job(1, 1, n, in_bytes=0, out_bytes=n,
                                grid=(1, 1, 1), block=(64, 1, 1),
                                param_count=2))
        d = self.wait_done(1)
        if d is None:
            self.check("(a) mail n=39", False, "no done record")
            return
        try:
            with open(self.out, "rb") as fh:
                got = fh.read()
        except OSError:
            got = b""
        ok = (d["status"] == ST_OK and d["bad"] == 0 and
              d["out_bytes_written"] == n and got == MAIL and
              d["kernel_ns"] > 0 and d["device"])
        print('BRIDGE-SELFTEST (a) mail n=%d gpu="%s" kernel_ns=%d '
              'status=%d bad=%d out_bytes=%d mail="%s"'
              % (n, d["device"], d["kernel_ns"], d["status"], d["bad"],
                 d["out_bytes_written"], got.decode(errors="replace")))
        self.check("(a) mail n=39", ok,
                   "out_exact=%s" % ("yes" if got == MAIL else "no"))

    def case_vadd(self):
        n = 1 << 20
        a = array.array("f", ((i % 1000) * 0.5 for i in range(n)))
        b = array.array("f", ((i % 997) * 0.25 for i in range(n)))
        self.write_in(a.tobytes() + b.tobytes())
        self.write_job(pack_job(2, 2, n, in_bytes=8 * n, out_bytes=4 * n,
                                grid=((n + 255) // 256, 1, 1),
                                block=(256, 1, 1), param_count=5))
        d = self.wait_done(2)
        if d is None:
            self.check("(b) vadd n=1<<20", False, "no done record")
            return
        with open(self.out, "rb") as fh:
            c = array.array("f")
            c.frombytes(fh.read(4 * n))
        bad = sum(1 for i in range(n) if c[i] != a[i] + b[i])
        checksum = sum(c)
        delta = abs(checksum - VADD_CHECKSUM)
        ok = (d["status"] == ST_OK and d["bad"] == 0 and bad == 0 and
              delta <= CHECKSUM_TOL and d["out_bytes_written"] == 4 * n and
              abs(d["checksum"] - VADD_CHECKSUM) <= CHECKSUM_TOL and
              d["kernel_ns"] > 0)
        print("BRIDGE-SELFTEST (b) vadd n=%d gpu=\"%s\" kernel_ns=%d "
              "status=%d bad=%d checksum=%.6f delta=%.3e out_bytes=%d"
              % (n, d["device"], d["kernel_ns"], d["status"], d["bad"],
                 checksum, delta, d["out_bytes_written"]))
        self.check("(b) vadd n=1<<20", ok, "file_bad=%d" % bad)

    def case_stale(self):
        with open(self.out, "rb") as fh:
            before = fh.read()
        self.write_job(pack_job(2, 1, 39, in_bytes=0, out_bytes=39,
                                grid=(1, 1, 1), block=(64, 1, 1),
                                param_count=2))
        time.sleep(0.6)
        with open(self.out, "rb") as fh:
            after = fh.read()
        lines = [ln for ln in self.log_text().splitlines()
                 if ln.startswith("CUBB2 seq=2 ")]
        d = None
        try:
            with open(self.done, "rb") as fh:
                d = parse_done(fh.read(DONE_SIZE))
        except (OSError, AssertionError):
            pass
        ok = (len(lines) == 1 and before == after and d is not None and
              d["seq"] == 2 and d["status"] == ST_OK)
        self.check("(c) stale seq=2 ignored", ok,
                   "exec_lines=%d out_unchanged=%s"
                   % (len(lines), "yes" if before == after else "no"))

    def case_badmagic(self):
        self.write_job(pack_job(100, 1, 39, in_bytes=0, out_bytes=39,
                                magic=0xDEADBEEF))
        time.sleep(0.4)
        hits = [ln for ln in self.log_text().splitlines()
                if "seq=100" in ln or "seq=100 " in ln]
        self.check("(d) invalid magic seq=100 ignored", len(hits) == 0,
                   "exec_lines=%d" % len(hits))

    def case_failure(self):
        self.write_job(pack_job(3, 9, 8, in_bytes=0, out_bytes=8))
        d = self.wait_done(3)
        ok = (d is not None and d["status"] != ST_OK and
              d["err_code"] == 9)
        self.check("(e) bad opcode done non-zero",
                   ok,
                   "status=%s err_code=%s" % (d["status"], d["err_code"])
                   if d else "no done record")

    def case_bound(self):
        n = (1 << 26) + 1
        self.write_job(pack_job(4, 2, n, in_bytes=8 * n, out_bytes=4 * n))
        d = self.wait_done(4)
        ok = (d is not None and d["status"] != ST_OK)
        self.check("(f) oversized n bound done non-zero", ok,
                   "status=%s err_code=%s" % (d["status"], d["err_code"])
                   if d else "no done record")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", default="/tmp/opencode/cuda-bridge-bridge-test")
    ap.add_argument("--bin", default=os.path.join(HERE, "cuda_bridge"))
    args = ap.parse_args()
    if not os.path.exists(args.bin):
        print("bridge_selftest: binary not found: %s (run make first)"
              % args.bin)
        return 2
    return Harness(args.bin, args.dir).run()


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