#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Olympian LED
"""Tests for step_kit.py: the self-test cases, `measure` end to end against a fake SCPI scope on 127.0.0.1, the
restore, and the kit's files kept free of private words. Run: python3 -I test_step_kit.py"""
import json
import re
import socket
import sys
import tempfile
import threading
import unittest
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import numpy as np  # noqa: E402

import step_kit as kit  # noqa: E402


class FakeClock:
    """Never really sleeps: sleep() moves the time on."""

    def __init__(self, start=1792000000.0):
        self.t = 0.0
        self.w = start

    def now(self):
        return self.t

    def wall(self):
        return self.w + self.t

    def sleep(self, s):
        if s > 0:
            self.t += s


def short_form(word):
    """A mixed-case SCPI word as the scope answers it: CHANnel1 -> CHAN1, FIFTy -> FIFT, NORMal -> NORM."""
    if word == word.upper():
        return word
    m = re.match(r"([A-Z]+)[a-z]*(\d*)$", word)
    return (m.group(1) + m.group(2)) if m else word.upper()


def depth_points(word):
    m = re.fullmatch(r"(\d+)([kM]?)", word)
    if not m:
        return None
    return int(m.group(1)) * {"": 1, "k": 1000, "M": 1_000_000}[m.group(2)]


SERIAL = "UNITSERIAL0042"


class FakeScope:
    """A Rigol-dialect scope on a local socket: settings it answers in short forms, a stock-like 1 Mpts single that
    gives 250,000 points, and a negative-going pulser on every channel whose falling edge is a 1.1 GHz Gaussian and
    whose return is a 150 MHz pole. Stock habits kept: a hidden channel ignores its scale and offset, the depth is
    taken only while running and only as a word, and a waveform source on a hidden channel is refused."""

    def __init__(self, model="MHO934", state="RUN", glitch=False, max_rate=None, fp=1_092_233.0, fall_f3=1.1e9):
        self.model = model
        self.mho = model.startswith("MHO")
        self.nch = int(model[5]) if model[5].isdigit() else 4
        self.max_rate = max_rate or (4e9 if self.mho else 1.25e9)
        self.run = state
        self.glitch = glitch
        self.fp, self.fall_f3 = fp, fall_f3
        self.errors = []
        self.lines = []
        self.v = {}
        for i in range(1, self.nch + 1):
            c = ":CHANnel%d:" % i
            self.v.update({c + "DISPlay": "1", c + "COUPling": "AC" if i == 1 else "DC", c + "BWLimit": "20M" if i == 1 else "OFF",
                           c + "PROBe": "1.0000E+01" if i == 1 else "1.0000E+00", c + "SCALe": "1.0000E+00",
                           c + "OFFSet": "5.0000E-01" if i == 1 else "0.0000E+00"})
            if self.mho:
                self.v[c + "IMPedance"] = "OMEG"
        self.v.update({":ACQuire:TYPE": "AVER", ":ACQuire:MDEPth": "1.0000E+04", ":TIMebase:MAIN:SCALe": "5.0000E-07",
                       ":TIMebase:MAIN:OFFSet": "1.0000E-06", ":TRIGger:MODE": "PULS", ":TRIGger:EDGE:SOURce": "CHAN2",
                       ":TRIGger:SWEep": "AUTO", ":WAVeform:SOURce": "CHAN3", ":WAVeform:MODE": "NORM",
                       ":WAVeform:FORMat": "BYTE", ":WAVeform:STARt": "1", ":WAVeform:STOP": "1000"})
        self.levels = {"CHAN%d" % i: "%.4E" % (0.1 * i) for i in range(1, self.nch + 1)}
        self.initial = (dict(self.v), dict(self.levels), self.run)
        self.changed = set()
        self.record = None
        self.n_single = 0
        self.rng = np.random.default_rng(99)
        self.srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.srv.bind(("127.0.0.1", 0))
        self.srv.listen(4)
        self.port = self.srv.getsockname()[1]
        self.alive = True
        threading.Thread(target=self._accept, daemon=True).start()

    # -- the socket
    def close(self):
        self.alive = False
        try:
            self.srv.close()
        except OSError:
            pass

    def _accept(self):
        while self.alive:
            try:
                c, _ = self.srv.accept()
            except OSError:
                return
            threading.Thread(target=self._serve, args=(c,), daemon=True).start()

    def _serve(self, c):
        buf = b""
        with c:
            while True:
                try:
                    d = c.recv(65536)
                except OSError:
                    return
                if not d:
                    return
                buf += d
                while b"\n" in buf:
                    line, _, buf = buf.partition(b"\n")
                    r = self.handle(line.decode().strip())
                    if r is None:
                        continue
                    c.sendall(r if isinstance(r, bytes) else (r + "\n").encode())

    # -- the scope
    def shown(self, i):
        return self.v.get(":CHANnel%d:DISPlay" % i) == "1"

    def srate(self):
        on = sum(self.shown(i) for i in range(1, self.nch + 1))
        top = self.max_rate / (1 if on <= 1 else 2 if on == 2 else 4)
        depth = float(self.v[":ACQuire:MDEPth"])
        raw = depth / (12 * float(self.v[":TIMebase:MAIN:SCALe"]))
        ladder = [top / d for d in (1, 2, 4, 5, 8, 10, 20, 40, 50, 100, 200, 400, 500, 1000, 10_000, 100_000)]
        return max([r for r in ladder if r <= raw * 1.0001] or [ladder[-1]])

    def note(self, key, now):
        init_v, init_l, _ = self.initial
        if key.startswith("LEVEL "):
            if now != init_l.get(key[6:]):
                self.changed.add(key)
        elif now != init_v.get(key):
            self.changed.add(key)

    def single(self):
        self.n_single += 1
        dt = 1.0 / self.srate()
        depth = int(float(self.v[":ACQuire:MDEPth"]))
        pts = depth // 4 if depth >= 1_000_000 else depth  # the stock's RAW read of a 1 Mpts single: 250,000 points
        src = 1  # the pulser is on CH1
        scale = float(self.v[":CHANnel%d:SCALe" % src])
        offset = float(self.v[":CHANnel%d:OFFSet" % src])
        off = int(self.rng.integers(0, 4000))
        x = kit.synth_pulser(pts + 4000, dt, 1.0 / self.fp, "pole", 150e6, lo=-0.5, hi=0.0, noise=1.5e-3,
                             fall=("gauss", self.fall_f3), seed=1000 + self.n_single)[off:off + pts]
        if self.glitch and scale < 0.5:
            x = x.copy()
            x[pts // 2] += 2.0
        yinc = 10 * scale / 65536
        yref = 32768.0
        yorig = offset / yinc
        codes = np.round(x / yinc + yorig + yref)
        lim = 4.1 * scale / yinc
        codes = np.clip(codes, max(0, yref - lim), min(65535, yref + lim)).astype("<u2")
        self.record = dict(codes=codes, pre="1,2,%d,1,%.10E,%.6E,0,%.10E,%.6E,%d" % (pts, dt, -pts * dt / 2, yinc, yorig, yref))
        self.run = "STOP"

    def handle(self, line):
        self.lines.append(line)
        if line.endswith("?"):
            return self.query(line[:-1])
        head, _, val = line.partition(" ")
        val = val.strip()
        if head == ":RUN":
            self.run = "RUN"
        elif head == ":STOP":
            self.run = "STOP"
        elif head == ":SINGle":
            self.single()
        elif head == ":TFORce":
            pass
        elif head == ":TRIGger:EDGE:LEVel":
            src = self.v[":TRIGger:EDGE:SOURce"]
            self.levels[src] = "%.4E" % float(val)
            self.note("LEVEL " + src, self.levels[src])
        elif head in self.v:
            self.set(head, val)
        else:
            self.errors.append('-113,"Undefined header"')
        return None

    def set(self, head, val):
        m = re.match(r":CHANnel(\d):(SCALe|OFFSet)$", head)
        if m and not self.shown(int(m.group(1))):
            return  # the stock ignores a hidden channel's scale and offset
        if head.endswith(":DISPlay"):
            now = {"1": "1", "ON": "1", "0": "0", "OFF": "0"}.get(val.upper())
            if now is None:
                self.errors.append('-224,"Illegal parameter value"')
                return
        elif head == ":ACQuire:MDEPth":
            pts = depth_points(val)
            if pts is None:
                self.errors.append('-222,"Data out of range"')  # 1.0000E+04 is refused: a word is needed
                return
            if self.run == "STOP":
                return  # taken only while running
            now = "%.4E" % pts
        elif head == ":WAVeform:SOURce":
            now = short_form(val)
            n = now[-1]
            if n.isdigit() and not self.shown(int(n)):
                self.errors.append('-200,"Execution error"')
                return
        elif head in (":WAVeform:STARt", ":WAVeform:STOP"):
            now = str(int(float(val)))
        else:
            try:
                now = "%.4E" % float(val)
            except ValueError:
                now = short_form(val)
        self.v[head] = now
        self.note(head, now)

    def query(self, head):
        if head == "*IDN":
            return "RIGOL TECHNOLOGIES,%s,%s,00.01.00" % (self.model, SERIAL)
        if head == "*OPC":
            return "1"
        if head == ":SYSTem:ERRor":
            return self.errors.pop(0) if self.errors else '0,"No error"'
        if head == ":TRIGger:STATus":
            return self.run if self.run == "STOP" else ("WAIT" if self.v[":TRIGger:SWEep"] == "NORM" else "AUTO")
        if head == ":ACQuire:SRATe":
            return "%.4E" % self.srate()
        if head == ":TRIGger:EDGE:LEVel":
            return self.levels.get(self.v[":TRIGger:EDGE:SOURce"], "0.0000E+00")
        if head == ":WAVeform:PREamble":
            return self.record["pre"] if self.record else "1,2,0,1,1E-9,0,0,1,0,0"
        if head == ":WAVeform:DATA":
            a, b = int(self.v[":WAVeform:STARt"]), int(self.v[":WAVeform:STOP"])
            data = self.record["codes"][a - 1:b].tobytes()
            return b"#9%09d" % len(data) + data + b"\n"
        if head in self.v:
            return self.v[head]
        self.errors.append('-113,"Undefined header"')
        return None  # an unknown query is never answered

    def same_as_initial(self):
        init_v, init_l, init_run = self.initial
        diff = {k: (init_v[k], self.v[k]) for k in init_v if init_v[k] != self.v[k]}
        diff.update({"LEVEL " + k: (init_l[k], self.levels[k]) for k in init_l if init_l[k] != self.levels[k]})
        return diff


class KitCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.dir = Path(self.tmp.name)

    def tearDown(self):
        self.tmp.cleanup()

    def measure(self, scope, *extra):
        out = self.dir / "run.json"
        argv = ["measure", "--host", "127.0.0.1:%d" % scope.port, "--ch", "1", "--out", str(out), "--timeout", "3"] + list(extra)
        rc = kit.main(argv, clock=FakeClock())
        doc = json.loads(out.read_text(encoding="utf-8")) if out.exists() else None
        return rc, doc, out


# ----------------------------------------------------------------------------------------------------------------
class SelfTest(unittest.TestCase):
    def test_every_case(self):
        lines = []
        fails, results = kit.selftest(echo=lines.append)
        self.assertEqual(fails, [], "\n".join(lines))
        names = [r["name"] for r in results]
        for want in ("pole 400 MHz", "gauss 900 MHz", "pole 300 MHz", "gauss 1300 MHz at 4 GSa/s", "locked period refused",
                     "fast falling edge picked", "slow phase walk warned"):
            self.assertTrue(any(n.startswith(want) for n in names), want)
        for r in results:
            if r.get("err_pct") is not None and not r["name"].startswith("slow"):
                self.assertLess(abs(r["err_pct"]), 1.0, r["name"])

    def test_locked_pulser_refused(self):
        dt = 0.8e-9
        rec = kit.synth_pulser(400_000, dt, 1000 * dt, "pole", 400e6, jitter=0.0, drift_ppm=0.0, noise=1e-3)
        with self.assertRaisesRegex(ValueError, "never sampled"):
            kit.step_analyse([rec], dt)

    def test_summarise_skips_a_one_bin_dip(self):
        f = [1e8, 2e8, 3e8, 4e8, 5e8, 6e8]
        db = [0.0, -0.2, -4.0, -0.5, -2.0, -5.0]  # 300 MHz is a one-bin dip, not the corner
        s = kit.summarise(f, db, d0=0.0)
        self.assertGreater(s["f3db"], 5e8)

    def test_depth_words(self):
        self.assertEqual(kit.depth_word("1.0000E+06"), "1M")
        self.assertEqual(kit.depth_word("1.0000E+04"), "10k")
        self.assertEqual(kit.parse_depth("1M"), 1_000_000)
        self.assertEqual(kit.parse_depth("250k"), 250_000)

    def test_refuses_reset_and_screenshot(self):
        s = kit.Scpi("127.0.0.1", 1)
        for cmd in ("*RST", ":DISPlay:DATA? PNG", ":SYSTem:SETup?", ":CALibrate:STARt"):
            with self.assertRaises(kit.KitError):
                s._send(cmd)


class Measure(KitCase):
    def check_restored(self, scope, doc):
        self.assertEqual(scope.same_as_initial(), {})
        self.assertTrue(doc["restore"]["complete"], doc["restore"])
        written = doc["restore"]["written"]
        for k in sorted(scope.changed):
            if k.startswith("LEVEL "):
                self.assertTrue(any(w.startswith(":TRIGger:EDGE:LEVel") and w.endswith("(%s)" % k[6:]) for w in written), (k, written))
            else:
                self.assertTrue(any(w.split(" ")[0] == k for w in written), (k, written))

    def test_end_to_end_stock_mho(self):
        scope = FakeScope()
        try:
            rc, doc, out = self.measure(scope, "--save-records", str(self.dir / "recs.npz"), "--csv", str(self.dir / "run"))
        finally:
            scope.close()
        self.assertEqual(rc, 0, doc.get("refused") or doc.get("error"))
        r = doc["results"]
        self.assertEqual(doc["edge"], "fall")
        self.assertEqual(doc["points_per_record"], 250_000)
        self.assertEqual(doc["preamble"]["points"], 250_000)
        self.assertAlmostEqual(doc["rate"], 4e9, delta=4e6)
        self.assertEqual(r["records"], 8)
        self.assertLess(abs(r["f3db"] / 1.1e9 - 1), 0.02, r["f3db"])
        self.assertGreaterEqual(r["phase_turns"], kit.MIN_TURNS)
        self.assertEqual(doc["unit"], {"model": "MHO934", "firmware": "00.01.00", "series": "MHO900"})
        self.assertNotIn(SERIAL, out.read_text(encoding="utf-8"))
        # every SCPI line, with its time
        sent = [e["send"] for e in doc["trace"]]
        self.assertIn(":TRIGger:SWEep NORMal", sent)
        self.assertIn(":CHANnel1:IMPedance FIFTy", sent)
        self.assertTrue(all("t" in e and "s" in e for e in doc["trace"]))
        self.assertEqual(sum(1 for x in sent if x == ":WAVeform:DATA?"), 9)  # the pulser check, then the 8 records
        rb = doc["settings"]["read_back"]
        self.assertEqual(rb[":CHANnel1:IMPedance"], "FIFT")
        self.assertEqual(rb[":TRIGger:SWEep"], "NORM")
        self.assertEqual([rb[":CHANnel%d:DISPlay" % i] for i in (1, 2, 3, 4)], ["1", "0", "0", "0"])
        self.check_restored(scope, doc)
        # analyse on the saved records gives the same numbers
        rc2 = kit.main(["analyse", str(self.dir / "recs.npz"), "--out", str(self.dir / "again.json")])
        self.assertEqual(rc2, 0)
        again = json.loads((self.dir / "again.json").read_text(encoding="utf-8"))
        self.assertEqual(again["edge"], "fall")
        self.assertEqual(again["results"], r)
        # the CSVs
        step = (self.dir / "run-step.csv").read_text(encoding="utf-8").splitlines()
        resp = (self.dir / "run-response.csv").read_text(encoding="utf-8").splitlines()
        self.assertEqual(step[0], "time_ps,level")
        self.assertEqual(resp[0], "freq_hz,db,snr_db")
        self.assertEqual(len(resp) - 1, len(r["freqs"]))
        self.assertEqual(len(step) - 1, len(r["step"]))
        # the report reads the JSON
        self.assertEqual(kit.main(["report", str(out)]), 0)

    def test_found_stopped_is_left_stopped(self):
        scope = FakeScope(state="STOP")
        try:
            rc, doc, _ = self.measure(scope, "--records", "3", "--edge", "fall")
        finally:
            scope.close()
        self.assertEqual(rc, 0, doc.get("refused") or doc.get("error"))
        self.assertEqual(scope.run, "STOP")
        self.check_restored(scope, doc)

    def test_clipped_record_refused_and_restored(self):
        scope = FakeScope(glitch=True)
        try:
            rc, doc, _ = self.measure(scope, "--records", "2")
        finally:
            scope.close()
        self.assertEqual(rc, 1)
        self.assertIn("screen's edge", doc["refused"])
        self.assertNotIn("results", doc)
        self.check_restored(scope, doc)

    def test_wrong_rate_refused_and_restored(self):
        scope = FakeScope(max_rate=2e9)
        try:
            rc, doc, _ = self.measure(scope, "--records", "2")
        finally:
            scope.close()
        self.assertEqual(rc, 1)
        self.assertIn("no time base takes", doc["refused"])
        self.check_restored(scope, doc)

    def test_slow_phase_walk_warned(self):
        scope = FakeScope(fp=1_003_017.0)
        try:
            rc, doc, _ = self.measure(scope, "--records", "2")
        finally:
            scope.close()
        self.assertEqual(rc, 0, doc.get("refused") or doc.get("error"))
        self.assertTrue(any("sweep the sample interval" in w for w in doc.get("warnings", [])))
        self.assertLess(doc["results"]["phase_turns"], kit.MIN_TURNS)

    def test_dho_at_50_ohm_refused_before_any_setting(self):
        scope = FakeScope(model="DHO924S")
        try:
            rc, _, _ = self.measure(scope, "--imp", "50")
        finally:
            scope.close()
        self.assertEqual(rc, 2)
        self.assertEqual([x for x in scope.lines if not x.endswith("?")], [])

    def test_other_models_refused(self):
        scope = FakeScope(model="XYZ1234")
        try:
            rc, _, _ = self.measure(scope)
        finally:
            scope.close()
        self.assertEqual(rc, 2)
        self.assertEqual([x for x in scope.lines if not x.endswith("?")], [])


class PublicText(unittest.TestCase):
    # the words are kept reversed, so this file passes its own check
    WORDS = [w[::-1] for w in ("bda", "adirf", "soavon", "dqca-avon", "dqca", "telpirt", "rif_hcneb", "feoc", "retsiger",
                               "1x0", "1-ohm", "2-ohm", "3-ohm", "1-epocs", "akzsik", "hcet kn", "/:x", "ratsavon",
                               "lacol.secived")]

    def files(self):
        fs = sorted(p for p in HERE.iterdir() if p.is_file() and p.suffix in (".py", ".md", "") and not p.name.startswith("."))
        self.assertTrue(any(p.name == "step_kit.py" for p in fs))
        return fs

    def test_no_private_words(self):
        for p in self.files():
            text = p.read_text(encoding="utf-8").lower()
            for w in self.WORDS:
                self.assertNotIn(w, text, "%s carries %r" % (p.name, w))

    def test_no_addresses_but_loopback(self):
        rx = re.compile(r"(?<![\d.])(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})(?!\.?\d)")
        for p in self.files():
            for ip in rx.findall(p.read_text(encoding="utf-8")):
                self.assertEqual(ip, "127.0.0.1", "%s carries %s" % (p.name, ip))

    def test_spdx_header(self):
        for p in self.files():
            if p.name in ("LICENSE", "SHA256SUMS"):  # the published copy carries its sums beside the files
                continue
            head = p.read_text(encoding="utf-8")[:300]
            self.assertIn("SPDX-License-Identifier: MIT", head, p.name)
            self.assertIn("Copyright (c) 2026 Olympian LED", head, p.name)


if __name__ == "__main__":
    unittest.main(verbosity=2)