#!/usr/bin/env python3
"""Live instrument test for SDG GUI bug #10.

Bug #10
-------
Sweep fallback from WAV to SPAC was timing-fragile because the code queried
SYST:ERR? immediately after the first SWWV write. This script tests the new
synchronised approach on a real instrument via raw socket SCPI.

What it checks
--------------
1. Instrument identity and basic socket comms.
2. Which direct sweep form the instrument accepts:
   - WAV,<LIN|LOG>
   - SPAC,<LIN|LOG>
3. The production-style fallback logic:
   *CLS -> SWWV write -> *OPC? -> SYST:ERR? -> optional SPAC fallback
4. Stale-error immunity:
   an old queued error should not poison the next sweep apply because the
   helper clears the error state before the real attempt.

Notes
-----
- This mirrors the current GUI sweep token set (WAV / STAR) on purpose,
  because the target is bug #10 regression testing of the GUI path.
- The programming guide documents START and SWMD for sweep. This script does
  not try to resolve that wider compatibility question.
"""

from __future__ import annotations

import argparse
import socket
import sys
import time
from dataclasses import dataclass
from typing import List, Optional

DEFAULT_PORT = 5025
SOCKET_TIMEOUT = 4.0


class SDGLan:
    def __init__(self) -> None:
        self.sock: Optional[socket.socket] = None

    def connect(self, host: str, port: int = DEFAULT_PORT) -> None:
        self.close()
        self.sock = socket.create_connection((host, port), timeout=SOCKET_TIMEOUT)
        self.sock.settimeout(SOCKET_TIMEOUT)

    def close(self) -> None:
        try:
            if self.sock is not None:
                self.sock.close()
        finally:
            self.sock = None

    def write(self, cmd: str) -> None:
        if self.sock is None:
            raise RuntimeError("Not connected")
        if not cmd.endswith("\n"):
            cmd += "\n"
        self.sock.sendall(cmd.encode("ascii"))

    def drain(self) -> bytes:
        if self.sock is None:
            raise RuntimeError("Not connected")
        data = bytearray()
        while True:
            try:
                chunk = self.sock.recv(65536)
                if not chunk:
                    break
                data.extend(chunk)
            except socket.timeout:
                break
        return bytes(data)

    def _recv_line(self) -> str:
        if self.sock is None:
            raise RuntimeError("Not connected")
        buf = bytearray()
        while True:
            chunk = self.sock.recv(4096)
            if not chunk:
                break
            buf.extend(chunk)
            if b"\n" in buf:
                break
        return buf.decode("ascii", errors="ignore")

    def query_retry(self, cmd: str, retries: int = 2) -> str:
        last_err: Optional[Exception] = None
        for _ in range(retries + 1):
            try:
                self.drain()
                self.write(cmd)
                return self._recv_line().strip()
            except socket.timeout as exc:
                last_err = exc
                try:
                    self.write("*CLS")
                except Exception:
                    pass
        if last_err is not None:
            raise last_err
        raise RuntimeError("query_retry failed without exception")


@dataclass
class ApplyResult:
    ok: bool
    cmd: str
    opc: str
    err: str
    used_fallback: bool = False
    fallback_cmd: str = ""
    fallback_opc: str = ""
    fallback_err: str = ""
    note: str = ""


def print_section(title: str) -> None:
    print("\n" + "=" * 78)
    print(title)
    print("=" * 78)


def sync_apply(io: SDGLan, cmd: str) -> ApplyResult:
    try:
        try:
            io.write("*CLS")
        except Exception:
            pass
        io.write(cmd)
        opc = io.query_retry("*OPC?")
        err = io.query_retry("SYST:ERR?")
        return ApplyResult(ok=err.startswith("0"), cmd=cmd, opc=opc, err=err)
    except Exception as exc:
        return ApplyResult(ok=False, cmd=cmd, opc="", err=str(exc), note="exception during sync_apply")


def apply_with_bug10_logic(io: SDGLan, channel: str, parts: List[str]) -> ApplyResult:
    cmd = f"{channel}:SWWV " + ",".join(parts)
    first = sync_apply(io, cmd)
    if first.ok:
        first.note = "primary form accepted"
        return first

    parts2 = [p.replace("WAV,", "SPAC,") for p in parts]
    if parts2 == parts:
        first.note = "primary form failed and no fallback token found"
        return first

    cmd2 = f"{channel}:SWWV " + ",".join(parts2)
    second = sync_apply(io, cmd2)
    return ApplyResult(
        ok=second.ok,
        cmd=first.cmd,
        opc=first.opc,
        err=first.err,
        used_fallback=True,
        fallback_cmd=cmd2,
        fallback_opc=second.opc,
        fallback_err=second.err,
        note="fallback accepted" if second.ok else "both forms failed",
    )


def build_gui_style_parts(start_hz: int, stop_hz: int, sweep_type: str = "LIN", time_s: float = 1.0,
                          direction: str = "UP", trigger_source: str = "INT") -> List[str]:
    return [
        "STATE,ON",
        f"WAV,{sweep_type}",
        f"STAR,{start_hz}",
        f"STOP,{stop_hz}",
        f"TIME,{time_s}",
        f"DIR,{direction}",
        f"TRSR,{trigger_source}",
    ]


def set_safe_basic_wave(io: SDGLan, channel: str) -> None:
    io.write(f"{channel}:BSWV WVTP,SINE,FRQ,1000,AMP,2,OFST,0")
    io.query_retry("*OPC?")
    _ = io.query_retry("SYST:ERR?")


def read_back(io: SDGLan, channel: str) -> str:
    try:
        return io.query_retry(f"{channel}:SWWV?")
    except Exception as exc:
        return f"<readback failed: {exc}>"


def inject_stale_error(io: SDGLan, channel: str) -> None:
    # Force a known parser error, wait until the instrument has processed it,
    # but intentionally do not read SYST:ERR? afterwards. The next helper call
    # should clear it with *CLS before doing the real sweep write.
    io.write(f"{channel}:SWWV BADTOKEN,1")
    try:
        io.query_retry("*OPC?")
    except Exception:
        pass


def cleanup(io: SDGLan, channel: str) -> None:
    try:
        sync_apply(io, f"{channel}:SWWV STATE,OFF")
    except Exception:
        pass


def main() -> int:
    parser = argparse.ArgumentParser(description="Live regression test for SDG GUI bug #10 sweep fallback")
    parser.add_argument("--ip", required=True, help="Instrument IP address")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help=f"Instrument port (default: {DEFAULT_PORT})")
    parser.add_argument("--channel", default="C1", choices=["C1", "C2"], help="Target channel")
    parser.add_argument("--start", type=int, default=1000, help="Sweep start frequency in Hz")
    parser.add_argument("--stop", type=int, default=5000, help="Sweep stop frequency in Hz")
    parser.add_argument("--time", type=float, default=1.0, help="Sweep time in seconds")
    parser.add_argument("--type", default="LIN", choices=["LIN", "LOG"], help="Sweep type token used by GUI")
    args = parser.parse_args()

    io = SDGLan()
    try:
        print_section("CONNECT")
        io.connect(args.ip, args.port)
        print(f"Connected to {args.ip}:{args.port}")
        print(f"*IDN? -> {io.query_retry('*IDN?')}")

        print_section("PREPARE")
        set_safe_basic_wave(io, args.channel)
        print(f"{args.channel}: basic waveform set to safe SINE baseline")

        parts = build_gui_style_parts(
            start_hz=args.start,
            stop_hz=args.stop,
            sweep_type=args.type,
            time_s=args.time,
            direction="UP",
            trigger_source="INT",
        )
        parts_spac = [p.replace("WAV,", "SPAC,") for p in parts]

        print_section("INFORMATIONAL DIRECT FORM CHECKS")
        direct_wav = sync_apply(io, f"{args.channel}:SWWV " + ",".join(parts))
        print(f"WAV form ok      : {direct_wav.ok}")
        print(f"WAV form cmd     : {direct_wav.cmd}")
        print(f"WAV form *OPC?   : {direct_wav.opc}")
        print(f"WAV form SYST:ERR?: {direct_wav.err}")

        direct_spac = sync_apply(io, f"{args.channel}:SWWV " + ",".join(parts_spac))
        print(f"SPAC form ok     : {direct_spac.ok}")
        print(f"SPAC form cmd    : {direct_spac.cmd}")
        print(f"SPAC form *OPC?  : {direct_spac.opc}")
        print(f"SPAC form SYST:ERR?: {direct_spac.err}")

        print_section("TEST 1 - PRODUCTION-LIKE BUG #10 HELPER")
        helper = apply_with_bug10_logic(io, args.channel, parts)
        print(f"helper ok        : {helper.ok}")
        print(f"helper note      : {helper.note}")
        print(f"primary cmd      : {helper.cmd}")
        print(f"primary *OPC?    : {helper.opc}")
        print(f"primary SYST:ERR?: {helper.err}")
        if helper.used_fallback:
            print(f"fallback used    : True")
            print(f"fallback cmd     : {helper.fallback_cmd}")
            print(f"fallback *OPC?   : {helper.fallback_opc}")
            print(f"fallback SYST:ERR?: {helper.fallback_err}")
        else:
            print("fallback used    : False")
        print(f"readback         : {read_back(io, args.channel)}")

        print_section("TEST 2 - STALE ERROR IMMUNITY")
        inject_stale_error(io, args.channel)
        stale = apply_with_bug10_logic(io, args.channel, parts)
        print(f"stale-test ok    : {stale.ok}")
        print(f"stale-test note  : {stale.note}")
        print(f"primary cmd      : {stale.cmd}")
        print(f"primary *OPC?    : {stale.opc}")
        print(f"primary SYST:ERR?: {stale.err}")
        if stale.used_fallback:
            print("fallback used    : True")
            print(f"fallback cmd     : {stale.fallback_cmd}")
            print(f"fallback *OPC?   : {stale.fallback_opc}")
            print(f"fallback SYST:ERR?: {stale.fallback_err}")
        else:
            print("fallback used    : False")
        print(f"readback         : {read_back(io, args.channel)}")

        print_section("RESULT")
        critical_ok = helper.ok and stale.ok
        print(f"critical tests passed: {critical_ok}")
        print("Informational note: direct WAV/SPAC one-form checks may differ by model/firmware.")
        return 0 if critical_ok else 1
    finally:
        try:
            cleanup(io, args.channel)
        finally:
            io.close()


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