#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# client_test.py - external producer/consumer check for the CUBB v1 transport.
#
# Simulates what vgsp.ko will do: write a CUBB job record into the job file,
# wait for the daemon to update the done file with the matching seq, then
# verify the fields. Uses a fake guest-RAM file (never /dev/shm/vgpu-guest-ram).
#
#   python3 client_test.py

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

HERE = os.path.dirname(os.path.abspath(__file__))
BIN = os.path.join(HERE, "cuda_bridge")
DIR = os.environ.get("CUBB_TEST_DIR", "/tmp/opencode/cuda-bridge-test")
RAM = os.path.join(DIR, "transport-ram.bin")
JOB = os.path.join(DIR, "transport-job.bin")
DONE = os.path.join(DIR, "transport-done.bin")
LOG = os.path.join(DIR, "transport-daemon.log")

N = 1 << 20
MASK = 0xFFFFFFFF

JOB_FMT = "<4sIQ10I9Q"
DONE_FMT = "<4sIQIIQQII64sQII"
assert struct.calcsize(JOB_FMT) == 128
assert struct.calcsize(DONE_FMT) == 128


def job_record(seq, opcode, n, mail_out_gpa=0, a_gpa=0, b_gpa=0, c_gpa=0):
    buffer_bytes = n * 4
    mail_bytes = n if opcode == 1 else 0
    return struct.pack(
        JOB_FMT,
        b"CUBB",
        1,
        seq,
        opcode,
        (n + 255) // 256 if opcode != 1 else 1,
        1,
        1,
        256,
        1,
        1,
        0,
        n,
        0,
        mail_out_gpa,
        mail_bytes,
        a_gpa,
        buffer_bytes,
        b_gpa,
        buffer_bytes,
        c_gpa,
        buffer_bytes,
        0,
    )


def parse_done(data):
    (
        magic,
        version,
        seq,
        status,
        bad_count,
        checksum_bits,
        kernel_ns,
        name_len,
        opcode,
        name,
        host_ns,
        err_code,
        _reserved,
    ) = struct.unpack(DONE_FMT, data)
    checksum = struct.unpack("<d", struct.pack("<Q", checksum_bits))[0]
    return {
        "magic": magic,
        "version": version,
        "seq": seq,
        "status": status,
        "bad_count": bad_count,
        "checksum": checksum,
        "kernel_ns": kernel_ns,
        "device": name.split(b"\0", 1)[0].decode(),
        "opcode": opcode,
        "host_ns": host_ns,
        "err_code": err_code,
    }


def fill_ram():
    fd = os.open(RAM, os.O_RDWR | os.O_CREAT)
    os.ftruncate(fd, 4 << 30)
    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)])
    os.pwrite(fd, a.tobytes(), 0x10000000)
    os.pwrite(fd, b.tobytes(), 0x20000000)
    os.close(fd)


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


def wait_done(seq, timeout=60.0):
    t0 = time.time()
    while time.time() - t0 < timeout:
        try:
            with open(DONE, "rb") as fh:
                data = fh.read(128)
        except FileNotFoundError:
            data = b""
        if len(data) == 128 and data[:4] == b"CUBD":
            if struct.unpack_from("<Q", data, 8)[0] == seq:
                return data
        time.sleep(0.02)
    return None


def main():
    os.makedirs(DIR, exist_ok=True)
    fill_ram()
    for path in (JOB, DONE):
        try:
            os.unlink(path)
        except FileNotFoundError:
            pass

    with open(LOG, "wb") as log:
        daemon = subprocess.Popen(
            [BIN, "--job", JOB, "--done", DONE, "--ram", RAM, "--poll-ms", "5"],
            stdout=log,
            stderr=subprocess.STDOUT,
        )
    fails = 0
    try:
        time.sleep(0.5)

        write_job(job_record(1000, 2, N, a_gpa=0x10000000, b_gpa=0x20000000,
                             c_gpa=0x30000000))
        data = wait_done(1000)
        if data is None:
            print("TRANSPORT seq=1000 vadd: no done record")
            fails += 1
        else:
            d = parse_done(data)
            with open(RAM, "rb") as fh:
                fh.seek(0x30000000)
                c = array.array("f")
                c.frombytes(fh.read(N * 4))
            bad = sum(
                1
                for i in range(N)
                if c[i] != (i % 1000) * 0.5 + (i % 997) * 0.25
            )
            checksum = sum(c)
            ok = (
                d["status"] == 0
                and d["bad_count"] == 0
                and bad == 0
                and abs(checksum - 392344090.5) < 1e-4
                and d["kernel_ns"] > 0
            )
            if not ok:
                fails += 1
            print(
                "TRANSPORT seq=%d op=%d status=%d bad=%d kernel_ns=%d "
                "checksum=%.6f device=\"%s\" file_bad=%d %s"
                % (
                    d["seq"],
                    d["opcode"],
                    d["status"],
                    d["bad_count"],
                    d["kernel_ns"],
                    d["checksum"],
                    d["device"],
                    bad,
                    "PASS" if ok else "FAIL",
                )
            )

        write_job(job_record(1001, 1, len(b"the little gpu man says you've got mail"),
                             mail_out_gpa=0x40000000))
        data = wait_done(1001)
        want = b"the little gpu man says you've got mail"
        if data is None:
            print("TRANSPORT seq=1001 mail: no done record")
            fails += 1
        else:
            d = parse_done(data)
            with open(RAM, "rb") as fh:
                fh.seek(0x40000000)
                got = fh.read(len(want))
            ok = d["status"] == 0 and d["bad_count"] == 0 and got == want
            if not ok:
                fails += 1
            print(
                "TRANSPORT seq=%d op=%d status=%d bad=%d mail=\"%s\" %s"
                % (
                    d["seq"],
                    d["opcode"],
                    d["status"],
                    d["bad_count"],
                    got.decode(),
                    "PASS" if ok else "FAIL",
                )
            )
    finally:
        daemon.terminate()
        try:
            daemon.wait(timeout=10)
        except subprocess.TimeoutExpired:
            daemon.kill()
            daemon.wait()

    print("TRANSPORT %s" % ("PASS" if fails == 0 else "FAIL"))
    return 0 if fails == 0 else 1


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