#!/usr/bin/env python3
"""Live receive-only TCI acceptance test. Uses only Python's standard library.

The test never sends START/STOP, TRX with a value, TUNE with a value, TX audio,
CW text, or any other command capable of keying the transmitter.
"""

import argparse
import math
import os
import select
import socket
import struct
import sys
import time


HDR = struct.Struct("<16I")
STREAM_NAMES = {0: "IQ", 1: "RX_AUDIO", 2: "TX_AUDIO", 3: "TX_CHRONO", 4: "LINEOUT"}
SAMPLE_BYTES = {0: 2, 1: 3, 2: 4, 3: 4}


def split_commands(text):
    return [part.strip() + ";" for part in text.split(";") if part.strip()]


def parse_command(command):
    body = command.rstrip(";")
    if ":" not in body:
        return body.lower(), []
    name, args = body.split(":", 1)
    return name.lower(), args.split(",")


class WS:
    def __init__(self, host, port):
        self.sock = socket.create_connection((host, port), timeout=4)
        self.sock.settimeout(None)
        key = "dGhlIHNhbXBsZSBub25jZQ=="
        req = (f"GET / HTTP/1.1\r\nHost: {host}:{port}\r\n"
               "Upgrade: websocket\r\nConnection: Upgrade\r\n"
               f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n")
        self.sock.sendall(req.encode("ascii"))
        raw = b""
        while b"\r\n\r\n" not in raw:
            raw += self.sock.recv(4096)
        head, self.buf = raw.split(b"\r\n\r\n", 1)
        if b" 101 " not in head.split(b"\r\n", 1)[0]:
            raise RuntimeError(head.decode("latin1", "replace"))

    def close(self):
        # Finish the WebSocket protocol explicitly. A bare TCP EOF is also
        # legal, but the server's EOF path is tested separately and must not
        # make every acceptance run consume a leaked client slot.
        try:
            self.send_frame(8, b"")
            ready, _, _ = select.select([self.sock], [], [], 0.25)
            if ready:
                self.sock.recv(4096)
        except OSError:
            pass
        try:
            self.sock.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
        self.sock.close()

    def send_frame(self, opcode, payload=b""):
        mask = os.urandom(4)
        n = len(payload)
        if n <= 125:
            head = bytes((0x80 | opcode, 0x80 | n))
        elif n <= 65535:
            head = bytes((0x80 | opcode, 0xFE)) + struct.pack("!H", n)
        else:
            head = bytes((0x80 | opcode, 0xFF)) + struct.pack("!Q", n)
        masked = bytes(v ^ mask[i & 3] for i, v in enumerate(payload))
        self.sock.sendall(head + mask + masked)

    def send_text(self, command):
        self.send_frame(1, command.encode("utf-8"))

    def _need(self, n, deadline):
        while len(self.buf) < n:
            left = deadline - time.monotonic()
            if left <= 0:
                return False
            ready, _, _ = select.select([self.sock], [], [], left)
            if not ready:
                return False
            block = self.sock.recv(65536)
            if not block:
                raise ConnectionError("TCI connection closed")
            self.buf += block
        return True

    def recv_frame(self, timeout=1.0):
        deadline = time.monotonic() + timeout
        if not self._need(2, deadline):
            return None
        b0, b1 = self.buf[0], self.buf[1]
        pos, n = 2, b1 & 0x7f
        if n == 126:
            if not self._need(4, deadline):
                return None
            n, pos = struct.unpack("!H", self.buf[2:4])[0], 4
        elif n == 127:
            if not self._need(10, deadline):
                return None
            n, pos = struct.unpack("!Q", self.buf[2:10])[0], 10
        if not self._need(pos + n, deadline):
            return None
        payload = self.buf[pos:pos + n]
        self.buf = self.buf[pos + n:]
        opcode = b0 & 0x0f
        if opcode == 9:
            self.send_frame(10, payload)
            return self.recv_frame(max(0.0, deadline - time.monotonic()))
        return opcode, payload


class LiveTest:
    def __init__(self, ws, seconds):
        self.ws = ws
        self.seconds = seconds
        self.passed = 0
        self.failed = 0
        self.warned = 0
        self.text = []
        self.binary = []
        self.initial = []

    def ok(self, label, condition, detail=""):
        if condition:
            self.passed += 1
            print(f"  ok   {label}" + (f" — {detail}" if detail else ""))
        else:
            self.failed += 1
            print(f"  FAIL {label}" + (f" — {detail}" if detail else ""))

    def warn(self, label, detail=""):
        self.warned += 1
        print(f"  WARN {label}" + (f" — {detail}" if detail else ""))

    def pump(self, seconds):
        deadline = time.monotonic() + seconds
        while time.monotonic() < deadline:
            frame = self.ws.recv_frame(min(0.2, deadline - time.monotonic()))
            if frame is None:
                continue
            opcode, payload = frame
            if opcode == 1:
                self.text.extend(split_commands(payload.decode("utf-8", "replace")))
            elif opcode == 2:
                self.binary.append(payload)

    def wait_text(self, name, seconds=1.5):
        deadline = time.monotonic() + seconds
        while time.monotonic() < deadline:
            for i, command in enumerate(self.text):
                got, args = parse_command(command)
                if got == name:
                    self.text.pop(i)
                    return command, args
            self.pump(min(0.1, deadline - time.monotonic()))
        return None, []

    def init(self):
        deadline = time.monotonic() + 5
        while time.monotonic() < deadline:
            self.pump(0.2)
            ready = next((x for x in self.text if parse_command(x)[0] == "ready"), None)
            if ready:
                self.initial = list(self.text)
                self.text.clear()
                return
        raise RuntimeError("READY not received")

    def initial_by_name(self, name):
        return [parse_command(x)[1] for x in self.initial if parse_command(x)[0] == name]

    def initial_checks(self):
        names = [parse_command(x)[0] for x in self.initial]
        required = {
            "protocol", "device", "receive_only", "trx_count", "channel_count",
            "vfo_limits", "if_limits", "modulations_list", "dds", "vfo", "if",
            "modulation", "rx_filter_band", "agc_mode", "tx_enable", "trx",
            "tune", "drive", "volume", "mute", "iq_samplerate",
            "audio_samplerate", "tx_frequency", "app_focus", "ready",
        }
        missing = sorted(required.difference(names))
        self.ok("initial state contains required commands", not missing,
                "missing=" + ",".join(missing) if missing else f"commands={len(names)}")
        self.ok("READY is last initialization command",
                bool(names) and names[-1] == "ready", names[-1] if names else "empty")

    def query(self, command, response, echo_same=False):
        self.text.clear()
        self.ws.send_text(command)
        got, args = self.wait_text(response)
        self.ok(command.rstrip(";"), got is not None, got or "no response")
        if echo_same and got is not None:
            self.text.clear()
            self.ws.send_text(got)
            echoed, _ = self.wait_text(response)
            self.ok("SET same " + command.rstrip(";"), echoed is not None,
                    echoed or "no confirmation")
        return args

    def command_checks(self, receivers, channels):
        print("Commands: receive-only queries")
        for rx in receivers:
            for name in ("dds", "modulation", "rx_filter_band", "agc_mode", "agc_gain",
                         "lock", "sql_enable", "sql_level", "rx_mute", "rx_nr_enable",
                         "rx_nb_enable", "rx_anf_enable", "rx_nb_param", "rx_bin_enable",
                         "rx_anc_enable", "rx_apf_enable", "rx_dse_enable", "rx_nf_enable",
                         "rit_enable", "rit_offset", "xit_enable", "xit_offset"):
                self.query(f"{name}:{rx};", name, echo_same=True)
            for ch in channels.get(rx, []):
                for name in ("vfo", "if", "rx_volume", "rx_balance"):
                    self.query(f"{name}:{rx},{ch};", name, echo_same=True)
            # Stub today, but its wire contract still has to be valid.
            self.query(f"rx_channel_enable:{rx},1;", "rx_channel_enable", echo_same=True)
        for command, response in (("trx:0;", "trx"), ("tune:0;", "tune"),
                                  ("drive:0;", "drive"), ("tune_drive:0;", "tune_drive"),
                                  ("split_enable:0;", "split_enable"), ("volume;", "volume"),
                                  ("mute;", "mute"), ("mon_volume;", "mon_volume"),
                                  ("mon_enable;", "mon_enable"),
                                  ("cw_macros_speed;", "cw_macros_speed"),
                                  ("cw_macros_delay;", "cw_macros_delay"),
                                  ("digl_offset;", "digl_offset"),
                                  ("digu_offset;", "digu_offset")):
            self.query(command, response,
                       echo_same=response not in ("trx", "tune"))

        # Exercise safe one-way commands without changing persistent radio state.
        self.ws.send_text("rx_sensors_enable:true,100;")
        got, _ = self.wait_text("rx_sensors", 2)
        self.ok("RX_SENSORS_ENABLE", got is not None, got or "no sensor report")
        self.ws.send_text("rx_sensors_enable:false;")
        self.ws.send_text("tx_sensors_enable:true,100;")
        got, _ = self.wait_text("tx_sensors", 2)
        self.ok("TX_SENSORS_ENABLE while RX", got is not None,
                got or "no sensor report")
        self.ws.send_text("tx_sensors_enable:false;")
        self.query("cw_keyer_speed;", "cw_macros_speed")
        self.ws.send_text("spot:ZZ0TCITEST,usb,14074000,4294967295,TCI RX test;")
        self.ws.send_text("spot_delete:ZZ0TCITEST;")
        self.ok("SPOT/SPOT_DELETE", True, "temporary marker removed")

        # Negative TX safety guard: malformed values must not key anything.
        self.ws.send_text("trx:0,not-a-bool,tci;")
        _, trx = self.wait_text("trx")
        self.ok("TRX malformed value rejected", len(trx) >= 2 and trx[1].lower() == "false",
                ",".join(trx))
        self.ws.send_text("tune:0,not-a-bool;")
        _, tune = self.wait_text("tune")
        self.ok("TUNE malformed value rejected", len(tune) >= 2 and tune[1].lower() == "false",
                ",".join(tune))

    def visible_change(self, label, response, changed, restored, dwell):
        """Apply one RX-only change, leave it visible, then restore it."""
        self.text.clear()
        self.ws.send_text(changed)
        got, args = self.wait_text(response, 2)
        self.ok(label + " changed", got is not None, got or "no confirmation")
        time.sleep(dwell)
        self.text.clear()
        self.ws.send_text(restored)
        got, args = self.wait_text(response, 2)
        self.ok(label + " restored", got is not None, got or "no confirmation")
        time.sleep(0.3)

    def visible_checks(self, receivers, channels, dwell):
        print(f"Visible UI sweep: dwell={dwell:.1f}s; every value is restored")
        for rx in receivers:
            dds = self.query(f"dds:{rx};", "dds")
            if len(dds) >= 2:
                old = int(float(dds[1]))
                self.visible_change(f"DDS rx={rx}", "dds",
                                    f"dds:{rx},{old + 5000};", f"dds:{rx},{old};", dwell)

            mode = self.query(f"modulation:{rx};", "modulation")
            filt = self.query(f"rx_filter_band:{rx};", "rx_filter_band")
            if len(mode) >= 2:
                old_mode = mode[1].lower()
                alternate = {"nfm": "am", "digu": "usb", "digl": "lsb",
                             "fmraw": "nfm", "dmr": "nfm"}.get(old_mode, "am")
                if alternate == old_mode:
                    alternate = "usb"
                self.visible_change(f"MODULATION rx={rx}", "modulation",
                                    f"modulation:{rx},{alternate};",
                                    f"modulation:{rx},{old_mode};", dwell)

            if len(filt) >= 3:
                lo, hi = int(filt[1]), int(filt[2])
                if hi - lo > 400:
                    self.visible_change(f"FILTER rx={rx}", "rx_filter_band",
                                        f"rx_filter_band:{rx},{lo + 100},{hi - 100};",
                                        f"rx_filter_band:{rx},{lo},{hi};", dwell)

            agc = self.query(f"agc_mode:{rx};", "agc_mode")
            if len(agc) >= 2:
                old_agc = agc[1].lower()
                new_agc = "fast" if old_agc != "fast" else "normal"
                self.visible_change(f"AGC rx={rx}", "agc_mode",
                                    f"agc_mode:{rx},{new_agc};",
                                    f"agc_mode:{rx},{old_agc};", dwell)

            nr = self.query(f"rx_nr_enable:{rx};", "rx_nr_enable")
            if len(nr) >= 2:
                old_nr = nr[1].lower()
                new_nr = "false" if old_nr == "true" else "true"
                self.visible_change(f"NR rx={rx}", "rx_nr_enable",
                                    f"rx_nr_enable:{rx},{new_nr};",
                                    f"rx_nr_enable:{rx},{old_nr};", dwell)

            sql = self.query(f"sql_enable:{rx};", "sql_enable")
            if len(sql) >= 2:
                old_sql = sql[1].lower()
                new_sql = "false" if old_sql == "true" else "true"
                self.visible_change(f"SQL rx={rx}", "sql_enable",
                                    f"sql_enable:{rx},{new_sql};",
                                    f"sql_enable:{rx},{old_sql};", dwell)

            for ch in channels.get(rx, []):
                vfo = self.query(f"vfo:{rx},{ch};", "vfo")
                if len(vfo) >= 3:
                    old_vfo = int(float(vfo[2]))
                    self.visible_change(f"VFO rx={rx} ch={ch}", "vfo",
                                        f"vfo:{rx},{ch},{old_vfo + 2000};",
                                        f"vfo:{rx},{ch},{old_vfo};", dwell)
                vol = self.query(f"rx_volume:{rx},{ch};", "rx_volume")
                if len(vol) >= 3:
                    old_vol = int(float(vol[2]))
                    new_vol = old_vol - 12 if old_vol > -49 else old_vol + 12
                    new_vol = max(-60, min(0, new_vol))
                    self.visible_change(f"VOLUME rx={rx} ch={ch}", "rx_volume",
                                        f"rx_volume:{rx},{ch},{new_vol};",
                                        f"rx_volume:{rx},{ch},{old_vol};", dwell)

    def visible_vfo_b(self, channels, dwell):
        """Move hardware VFO-B; this is not a UI slice on the main pan."""
        rx, ch = 0, 1
        self.ok("TCI VFO-B channel is present", ch in channels.get(rx, []),
                str(channels))
        if ch not in channels.get(rx, []):
            return
        print(f"VFO-B only: +15 kHz for {dwell:.1f}s, then restore")
        vfo = self.query(f"vfo:{rx},{ch};", "vfo")
        if len(vfo) >= 3:
            old_vfo = int(float(vfo[2]))
            self.visible_change("VFO-B", "vfo",
                                f"vfo:{rx},{ch},{old_vfo + 15000};",
                                f"vfo:{rx},{ch},{old_vfo};", dwell)

    def sample_stats(self, payload, fmt, length):
        data = payload[HDR.size:]
        count = min(length, 4096)
        vals = []
        if fmt == 3:
            vals = struct.unpack_from(f"<{count}f", data)
        elif fmt == 0:
            vals = [v / 32768.0 for v in struct.unpack_from(f"<{count}h", data)]
        elif fmt == 1:
            vals = []
            for i in range(min(count, len(data) // 3)):
                v = data[i * 3] | (data[i * 3 + 1] << 8) | (data[i * 3 + 2] << 16)
                if v & 0x800000:
                    v -= 1 << 24
                vals.append(v / 8388608.0)
        elif fmt == 2:
            vals = [v / 2147483648.0 for v in struct.unpack_from(f"<{count}i", data)]
        finite = [v for v in vals if math.isfinite(v)]
        if not finite:
            return 0.0, 0.0, False
        rms = math.sqrt(sum(v * v for v in finite) / len(finite))
        return rms, max(abs(v) for v in finite), len(finite) == len(vals)

    def stream(self, rx, start, stop, wanted_type, requested_length=None,
               expected_rate=None):
        self.binary.clear()
        self.text.clear()
        self.ws.send_text(f"{start}:{rx};")
        self.pump(self.seconds)
        self.ws.send_text(f"{stop}:{rx};")
        self.pump(0.25)
        errors = [x for x in self.text if parse_command(x)[0] == "tci_error"]
        blocks = []
        for payload in self.binary:
            if len(payload) < HDR.size:
                continue
            h = HDR.unpack_from(payload)
            if h[6] == wanted_type and h[0] == rx:
                blocks.append((payload, h))
        label = f"{STREAM_NAMES[wanted_type]} rx={rx}"
        self.ok(label + " blocks", bool(blocks),
                errors[0] if errors else f"blocks={len(blocks)}")
        if not blocks:
            return
        bad = 0
        rms_values = []
        rates = set()
        for payload, h in blocks:
            receiver, rate, fmt, codec, crc, length, kind, chans = h[:8]
            rates.add(rate)
            # IQ length counts real values across I/Q; audio length is samples
            # per channel, so its payload additionally includes every channel.
            payload_samples = length if kind == 0 else length * chans
            expected = HDR.size + payload_samples * SAMPLE_BYTES.get(fmt, 0)
            if codec != 0 or crc != 0 or chans not in (1, 2) or expected != len(payload):
                bad += 1
            if requested_length is not None and length != requested_length:
                bad += 1
            rms, peak, finite = self.sample_stats(payload, fmt, length)
            if not finite:
                bad += 1
            rms_values.append(rms)
        if expected_rate is not None and rates != {expected_rate}:
            bad += 1
        self.ok(label + " headers", bad == 0,
                f"blocks={len(blocks)}, bad={bad}, rates={sorted(rates)}")
        rms = max(rms_values) if rms_values else 0.0
        if rms > 1e-7:
            self.ok(label + " signal", True, f"max RMS={rms:.6g}")
        else:
            self.warn(label + " signal is silent", "check mute/squelch and tuned signal")

    def stream_checks(self, receivers):
        print("Streams: IQ, demod audio and line-out")
        self.ws.send_text("iq_samplerate:48000;")
        self.wait_text("iq_samplerate")
        self.ws.send_text("audio_samplerate:12000;")
        self.wait_text("audio_samplerate")
        self.ws.send_text("audio_stream_channels:2;")
        self.ws.send_text("audio_stream_sample_type:float32;")
        self.ws.send_text("audio_stream_samples:512;")
        self.pump(0.2)
        for rx in receivers:
            self.stream(rx, "iq_start", "iq_stop", 0)
            self.stream(rx, "audio_start", "audio_stop", 1, 512)
            self.stream(rx, "line_out_start", "line_out_stop", 4, 512)

    def format_matrix(self, receivers):
        """Exercise every negotiated audio representation on one live RX."""
        rx = receivers[1] if len(receivers) > 1 else receivers[0]
        old_seconds = self.seconds
        self.seconds = 0.7
        print(f"Audio format/rate matrix on rx={rx}")
        try:
            self.ws.send_text("audio_samplerate:12000;")
            self.wait_text("audio_samplerate")
            self.ws.send_text("audio_stream_samples:128;")
            for channels in (1, 2):
                self.ws.send_text(f"audio_stream_channels:{channels};")
                for sample_type in ("int16", "int24", "int32", "float32"):
                    self.ws.send_text(f"audio_stream_sample_type:{sample_type};")
                    self.pump(0.05)
                    print(f"  config channels={channels}, type={sample_type}")
                    self.stream(rx, "audio_start", "audio_stop", 1, 128)
                    self.stream(rx, "line_out_start", "line_out_stop", 4, 128)
            self.ws.send_text("audio_stream_channels:2;")
            self.ws.send_text("audio_stream_sample_type:float32;")
            for rate in (8000, 12000, 24000, 48000):
                self.ws.send_text(f"audio_samplerate:{rate};")
                self.wait_text("audio_samplerate")
                self.ws.send_text("audio_stream_samples:128;")
                print(f"  config rate={rate}")
                self.stream(rx, "audio_start", "audio_stop", 1, 128)
        finally:
            self.seconds = old_seconds

    def iq_rate_matrix(self, receivers):
        """Compare negotiated IQ rate with every live block header."""
        rx = receivers[0]
        old_seconds = self.seconds
        self.seconds = 0.8
        print(f"IQ rate matrix on rx={rx}")
        try:
            for requested in (48000, 96000, 192000, 384000):
                self.text.clear()
                self.ws.send_text(f"iq_samplerate:{requested};")
                got, args = self.wait_text("iq_samplerate")
                effective = int(args[0]) if got and args else 0
                self.ok(f"IQ rate request {requested}", effective > 0,
                        got or "no response")
                if effective > 0:
                    self.stream(rx, "iq_start", "iq_stop", 0,
                                expected_rate=effective)
        finally:
            self.ws.send_text("iq_samplerate:48000;")
            self.wait_text("iq_samplerate")
            self.seconds = old_seconds


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=40001)
    ap.add_argument("--seconds", type=float, default=2.0,
                    help="capture time for each stream and receiver")
    ap.add_argument("--commands-only", action="store_true",
                    help="skip IQ/audio capture")
    ap.add_argument("--streams-only", action="store_true",
                    help="skip command matrix and capture only IQ/audio")
    ap.add_argument("--format-matrix", action="store_true",
                    help="test every audio sample type/channel count and rate")
    ap.add_argument("--iq-rate-matrix", action="store_true",
                    help="test negotiated IQ rates against binary headers")
    ap.add_argument("--visible", action="store_true",
                    help="visibly change RX parameters and restore every value")
    ap.add_argument("--vfo-b-visible", action="store_true",
                    help="move hardware VFO-B by 15 kHz and restore it")
    ap.add_argument("--dwell", type=float, default=1.5,
                    help="seconds to leave each visible test value active")
    ap.add_argument("--expect-channels", type=int, default=0,
                    help="require this exact total TCI channel count")
    args = ap.parse_args()

    print(f"Connecting to ws://{args.host}:{args.port} (RX-only safety mode)")
    ws = WS(args.host, args.port)
    test = LiveTest(ws, args.seconds)
    try:
        test.init()
        test.initial_checks()
        protocol = test.initial_by_name("protocol")
        test.ok("READY/init", bool(protocol), str(protocol))
        trx = test.initial_by_name("trx")
        if trx and len(trx[-1]) >= 2 and trx[-1][1].lower() == "true":
            raise RuntimeError("radio is already transmitting; aborting RX-only test")

        channels = {}
        for row in test.initial_by_name("vfo"):
            if len(row) >= 3:
                channels.setdefault(int(row[0]), set()).add(int(row[1]))
        channels = {rx: sorted(value) for rx, value in channels.items()}
        receivers = sorted({int(row[0]) for row in test.initial_by_name("dds") if row})
        print("Discovered receivers/channels:", receivers, channels)
        test.ok("live receivers found", bool(receivers), str(receivers))
        slice_count = sum(map(len, channels.values()))
        if args.expect_channels > 0:
            test.ok(f"expected {args.expect_channels} TCI channels visible",
                    slice_count == args.expect_channels, str(channels))
        else:
            test.ok("at least one slice/channel visible", slice_count > 0,
                    str(channels))

        if args.vfo_b_visible:
            test.visible_vfo_b(channels, args.dwell)
        elif args.iq_rate_matrix:
            test.iq_rate_matrix(receivers)
        elif args.format_matrix:
            test.format_matrix(receivers)
        elif args.visible:
            test.visible_checks(receivers, channels, args.dwell)
        elif not args.streams_only:
            test.command_checks(receivers, channels)
        if (not args.commands_only and not args.visible and
                not args.vfo_b_visible and
                not args.format_matrix and not args.iq_rate_matrix):
            test.stream_checks(receivers)
        print("Skipped by RF-safety policy: START, STOP, TRX/TUNE setters, TX audio, "
              "CW text/keying, SPOT_CLEAR and SET_IN_FOCUS")
    finally:
        # Idempotent receive-side cleanup only. Never touch TRX/TUNE.
        for rx in range(8):
            for name in ("iq_stop", "audio_stop", "line_out_stop"):
                try:
                    ws.send_text(f"{name}:{rx};")
                except OSError:
                    pass
        try:
            ws.send_text("rx_sensors_enable:false;")
        except OSError:
            pass
        ws.close()

    print(f"\nResult: {test.passed + test.failed} checks, "
          f"failed={test.failed}, warnings={test.warned}")
    return 1 if test.failed else 0


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