#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""B6 bench leg: IEC 60870-5-104 master client (GI + COV round-trip).

Task card: docs/任务规划/W1-采集保护收尾.md (B6 DoD); protocol baseline:
docs/子系统设计/IEC104上行设计说明.md §3/§4. Constants mirror
App/protocol/iec104/iec104_asdu.h / iec104_link.h.

Flow:
  1. TCP connect device :2404, STARTDT act/con (enter RUN).
  2. General interrogation C_IC_NA_1 (QOI=20): expect actcon, background
     scan (YX double-point group SQ=1 -> YX single-point group SQ=1 ->
     YC group SQ=1 -> frequency single point), actterm.
  3. Deterministic COV trigger: bench analog inputs are floating (phasor
     rms is 0..0.3 raw-code noise only; the standard source belongs to
     B10), so we use the two-step C_SE_NC_1 write to temporarily lower
     the OC1 current setting (comm addr 0x8858, default 10.0 A) to
     0.05 A -> protectTask (20 ms cadence) asserts OC1 -> YX16/YX15
     spontaneous (COT=3) uplink. This is exactly the "104 YX linkage"
     leg left open by B5. Then restore 10.0 A -> YX16 returns.
  4. YK11 (0x600A) two-step select/execute -> positive actcon (reset
     request consumed asynchronously by protectTask, B5 contract).
  5. C_RD read back 0x8858 ~= 10.0 (restore confirmation).

Exit codes: 0 = all checks PASS; 1 = any check FAIL; 2 = connect or
protocol error. Output is pure ASCII (redirect stdout to archive the
round-trip log as bench evidence).

Usage: python iec104_master.py [--ip 192.168.1.100] [--port 2404]
                               [--addr 0x8858] [--low 0.05]
                               [--restore 10.0]
"""
import argparse
import collections
import socket
import struct
import sys
import time

# ---- constants (mirror iec104_asdu.h / iec104_link.h) -----------------------
U_STARTDT_ACT, U_STARTDT_CON = 0x07, 0x0B
U_TESTFR_ACT, U_TESTFR_CON = 0x43, 0x83

T_M_SP_NA_1, T_M_DP_NA_1, T_M_ME_NC_1 = 1, 3, 13
T_C_SE_NC_1, T_C_DC_NA_1 = 50, 52
T_C_IC_NA_1, T_C_RD_NA_1 = 100, 102

COT_BG, COT_SPONT, COT_INIT, COT_REQ = 2, 3, 4, 5
COT_ACT, COT_ACTCON, COT_ACTTERM = 6, 7, 8
COT_PN = 0x40                 # COT high byte bit6: negative confirm
RQT_GROUP1 = 20
CA = 1                        # common address fixed to 1 in V0.1
SCAN_COTS = (COT_BG, COT_INIT)

YX_NAMES = {15: "YX15 fault-total", 16: "YX16 OC-I", 17: "YX17 OC-II",
            18: "YX18 OC-III", 23: "YX23 overload", 31: "YX31 prot-block"}

T0 = time.time()


def log(msg):
    print("[%9.3fs] %s" % (time.time() - T0, msg), flush=True)


def hexs(b):
    return " ".join("%02X" % x for x in b)


# ---- APCI / ASDU builders ----------------------------------------------------
def u_frame(code):
    return bytes([0x68, 0x04, code, 0x00, 0x00, 0x00])


def s_frame(vr):
    return bytes([0x68, 0x04, 0x01, 0x00, (vr << 1) & 0xFE, (vr >> 7) & 0x7F])


def i_frame(vs, vr, body):
    ctrl = bytes([(vs << 1) & 0xFE, (vs >> 7) & 0x7F,
                  (vr << 1) & 0xFE, (vr >> 7) & 0x7F])
    return bytes([0x68, len(body) + 4]) + ctrl + body


def put_ioa(ioa):
    return bytes([ioa & 0xFF, (ioa >> 8) & 0xFF, (ioa >> 16) & 0xFF])


def get_ioa(b, off=0):
    return b[off] | (b[off + 1] << 8) | (b[off + 2] << 16)


def asdu(type_id, vsq, cot, cot_hi, ioa=None, payload=b""):
    h = bytes([type_id, vsq, cot & 0x3F, cot_hi, CA & 0xFF, (CA >> 8) & 0xFF])
    if ioa is not None:
        h += put_ioa(ioa)
    return h + payload


class Master(object):
    """Minimal 104 master. k/w/t are covered by "S after every recv batch"
    (far faster than t1=15s / t2=10s), so no timers are needed here."""

    def __init__(self, ip, port):
        self.sock = socket.create_connection((ip, port), timeout=10)
        self.sock.settimeout(1.0)
        self.rx = b""
        self.vs = 0
        self.vr = 0
        self.vr_dirty = False
        self.events = collections.deque()   # (kind, detail) tuples
        self.yx = {}                        # ioa -> (val, iv)
        self.yc = {}                        # ioa -> (val, qds)

    # ---- send ----
    def _send(self, data, tag):
        log("TX  %s  %s" % (tag, hexs(data)))
        self.sock.sendall(data)

    def send_i(self, body, tag):
        self._send(i_frame(self.vs, self.vr, body), tag)
        self.vs = (self.vs + 1) & 0x7FFF

    def _send_s(self):
        if self.vr_dirty:
            self._send(s_frame(self.vr), "S  N(R)=%d" % self.vr)
            self.vr_dirty = False

    # ---- receive pump ----
    def pump_until(self, deadline):
        """Pump RX until deadline (non-blocking semantics for wait_event)."""
        while time.time() < deadline:
            try:
                chunk = self.sock.recv(512)
            except socket.timeout:
                self._send_s()
                continue
            if not chunk:
                raise RuntimeError("peer closed (EOF)")
            self.rx += chunk
            while len(self.rx) >= 2 and self.rx[0] == 0x68:
                ln = self.rx[1]
                if ln < 4 or ln > 253:
                    raise RuntimeError("bad APDU length %d" % ln)
                if len(self.rx) < 2 + ln:
                    break
                apdu, self.rx = self.rx[:2 + ln], self.rx[2 + ln:]
                self._handle(apdu)
            if self.rx and self.rx[0] != 0x68:
                raise RuntimeError("stream desync at 0x%02X" % self.rx[0])
            self._send_s()

    def wait_event(self, pred, timeout):
        deadline = time.time() + timeout
        while True:
            for i, e in enumerate(self.events):
                if pred(e):
                    return self.events[i]
            if time.time() >= deadline:
                return None
            self.pump_until(min(deadline, time.time() + 1.0))

    # ---- frame handling ----
    def _handle(self, apdu):
        c = apdu[2:6]
        b0, b1, b2, b3 = c[0], c[1], c[2], c[3]
        if (b0 & 0x01) == 0:                       # I format
            ns = (b0 >> 1) | ((b1 & 0x7F) << 7)
            nr = (b2 >> 1) | ((b3 & 0x7F) << 7)
            if ns != self.vr:
                raise RuntimeError("N(S)=%d out of order (expect %d)"
                                   % (ns, self.vr))
            self.vr = (self.vr + 1) & 0x7FFF
            self.vr_dirty = True
            self._asdu(apdu[6:], nr)
        elif (b0 & 0x03) == 0x01:                  # S format
            nr = (b2 >> 1) | ((b3 & 0x7F) << 7)
            log("RX  S  N(R)=%d" % nr)
        else:                                      # U format
            if b0 == U_STARTDT_CON:
                log("RX  U  STARTDT con")
                self.events.append(("startdt_con", None))
            elif b0 == U_TESTFR_ACT:
                self._send(u_frame(U_TESTFR_CON), "U  TESTFR con")
            elif b0 == U_TESTFR_CON:
                log("RX  U  TESTFR con")
            else:
                log("RX  U  code=0x%02X" % b0)

    def _asdu(self, a, nr):
        if len(a) < 6:
            raise RuntimeError("short ASDU")
        type_id, vsq, cot, cot_hi = a[0], a[1], a[2], a[3]
        pn = (cot_hi & COT_PN) != 0
        sq = (vsq & 0x80) != 0
        n = vsq & 0x7F
        obj = a[6:]

        if type_id in (T_M_SP_NA_1, T_M_DP_NA_1):
            pts = []
            if sq:
                base = get_ioa(obj)
                for i in range(n):
                    pts.append((base + i, obj[3 + i]))
            else:
                off = 0
                for _ in range(n):
                    pts.append((get_ioa(obj, off), obj[off + 3]))
                    off += 4
            for ioa, raw in pts:
                val = (raw & 0x03) if type_id == T_M_DP_NA_1 else (raw & 0x01)
                iv = (raw & 0x80) != 0
                self.yx[ioa] = (val, iv)
                self.events.append(("yx", (ioa, val, iv, cot)))
                if ioa in YX_NAMES or cot == COT_SPONT:
                    log("RX  I  %s ioa=%d val=%d iv=%d cot=%d  [N(R)=%d]"
                        % ("M_DP" if type_id == T_M_DP_NA_1 else "M_SP",
                           ioa, val, 1 if iv else 0, cot, nr))
        elif type_id == T_M_ME_NC_1:
            pts = []
            if sq:
                base = get_ioa(obj)
                for i in range(n):
                    v = struct.unpack_from("<f", obj, 3 + i * 5)[0]
                    pts.append((base + i, v, obj[3 + i * 5 + 4]))
            else:
                off = 0
                for _ in range(n):
                    v = struct.unpack_from("<f", obj, off + 3)[0]
                    pts.append((get_ioa(obj, off), v, obj[off + 3 + 4]))
                    off += 8
            for ioa, v, qds in pts:
                self.yc[ioa] = (v, qds)
                self.events.append(("yc", (ioa, v, qds, cot)))
                if cot != COT_BG and cot != COT_INIT:
                    log("RX  I  M_ME_NC ioa=0x%04X val=%.3f qds=0x%02X cot=%d"
                        % (ioa, v, qds, cot))
        else:
            self.events.append(("asdu", (type_id, cot, pn, bytes(obj))))
            log("RX  I  type=%d cot=%d%s body=%s"
                % (type_id, cot, " PN" if pn else "", hexs(obj)))

    def close(self):
        try:
            self.sock.settimeout(2.0)
            self.sock.close()
        except Exception:
            pass


# ---- test flow ----------------------------------------------------------------
class Checks(object):
    def __init__(self):
        self.items = []

    def add(self, name, ok, detail=""):
        self.items.append((name, bool(ok), detail))
        log("CHECK %-28s %s %s" % (name, "PASS" if ok else "FAIL", detail))


def actcon_of(type_id, want_pn=None):
    return lambda e: (e[0] == "asdu" and e[1][0] == type_id
                      and e[1][1] == COT_ACTCON
                      and (want_pn is None or e[1][2] == want_pn))


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--ip", default="192.168.1.100")
    ap.add_argument("--port", type=int, default=2404)
    ap.add_argument("--addr", default="0x8858",
                    help="setpoint comm addr to drive COV (OC1 current)")
    ap.add_argument("--low", type=float, default=0.05,
                    help="temporary low OC1 setting (A) to force pickup")
    ap.add_argument("--restore", type=float, default=10.0,
                    help="OC1 default setting to restore (A)")
    args = ap.parse_args()
    addr = int(args.addr, 0)

    ck = Checks()
    log("iec104 master client -> %s:%d  (COV trigger: write 0x%X %.2f then "
        "restore %.2f)" % (args.ip, args.port, addr, args.low, args.restore))

    try:
        m = Master(args.ip, args.port)
    except Exception as e:
        log("CONNECT FAIL: %r" % (e,))
        return 2

    rc = 2
    try:
        # 1. STARTDT
        m._send(u_frame(U_STARTDT_ACT), "U  STARTDT act")
        ck.add("startdt_con", m.wait_event(
            lambda e: e[0] == "startdt_con", 5.0) is not None)

        # 2. first general interrogation
        m.send_i(asdu(T_C_IC_NA_1, 1, COT_ACT, 0, ioa=0,
                      payload=bytes([RQT_GROUP1])),
                 "I  C_IC act (QOI=20)")
        ck.add("gi_actcon", m.wait_event(actcon_of(T_C_IC_NA_1, False),
                                         5.0) is not None)

        def scan_yx(ioa):
            return lambda e: (e[0] == "yx" and e[1][0] == ioa
                              and e[1][3] in SCAN_COTS)

        ck.add("scan_yx_dp(ioa4)", m.wait_event(scan_yx(4), 10.0) is not None)
        ck.add("scan_yx_sp(ioa16)", m.wait_event(scan_yx(16), 10.0) is not None)
        ck.add("scan_yc(ioa0x4018)", m.wait_event(
            lambda e: (e[0] == "yc" and e[1][0] == 0x4018
                       and e[1][3] in SCAN_COTS), 10.0) is not None)
        ck.add("scan_freq(ioa0x401D)", m.wait_event(
            lambda e: (e[0] == "yc" and e[1][0] == 0x401D
                       and e[1][3] in SCAN_COTS), 10.0) is not None)
        ck.add("gi_actterm", m.wait_event(
            lambda e: (e[0] == "asdu" and e[1][0] == T_C_IC_NA_1
                       and e[1][1] == COT_ACTTERM), 10.0) is not None)

        v16 = m.yx.get(16)
        ck.add("baseline_yx16_off", v16 is not None and v16 == (0, False),
               "yx16=%r" % (v16,))

        # 3. COV up: lower OC1 setting (two-step C_SE_NC_1)
        def se(select):
            pl = struct.pack("<f", args.low) + bytes([0x80 if select else 0x00])
            return asdu(T_C_SE_NC_1, 1, COT_ACT, 0, ioa=addr, payload=pl)

        m.send_i(se(True), "I  C_SE select 0x%X = %.3f" % (addr, args.low))
        ck.add("se_select_con", m.wait_event(actcon_of(T_C_SE_NC_1, False),
                                             5.0) is not None)
        m.send_i(se(False), "I  C_SE execute 0x%X = %.3f" % (addr, args.low))
        ck.add("se_execute_con", m.wait_event(actcon_of(T_C_SE_NC_1, False),
                                              5.0) is not None)
        m.events.clear()      # drop any echo/ack noise before COV watch

        def spont(ioa, val):
            return lambda e: (e[0] == "yx" and e[1][0] == ioa
                              and e[1][1] == val and e[1][3] == COT_SPONT)

        ck.add("spont_yx16_ON", m.wait_event(spont(16, 1), 8.0) is not None)
        ck.add("spont_yx15_ON", m.wait_event(spont(15, 1), 3.0) is not None)

        # 4. COV down: restore OC1 default setting
        def se_restore(select):
            pl = struct.pack("<f", args.restore) + bytes([0x80 if select else 0x00])
            return asdu(T_C_SE_NC_1, 1, COT_ACT, 0, ioa=addr, payload=pl)

        m.send_i(se_restore(True), "I  C_SE select 0x%X = %.3f (restore)"
                 % (addr, args.restore))
        ck.add("se_restore_select_con",
               m.wait_event(actcon_of(T_C_SE_NC_1, False), 5.0) is not None)
        m.send_i(se_restore(False), "I  C_SE execute 0x%X = %.3f (restore)"
                 % (addr, args.restore))
        ck.add("se_restore_execute_con",
               m.wait_event(actcon_of(T_C_SE_NC_1, False), 5.0) is not None)
        m.events.clear()

        got_spont_off = m.wait_event(spont(16, 0), 5.0) is not None
        if got_spont_off:
            ck.add("yx16_OFF", True, "via spontaneous COT=3")
        else:
            # floating-input chatter may have returned YX16 before the
            # restore write; fall back to a second GI background scan.
            log("no spont OFF within 5s -> second GI background scan")
            m.send_i(asdu(T_C_IC_NA_1, 1, COT_ACT, 0, ioa=0,
                          payload=bytes([RQT_GROUP1])),
                     "I  C_IC act #2 (QOI=20)")
            ok_term = m.wait_event(
                lambda e: (e[0] == "asdu" and e[1][0] == T_C_IC_NA_1
                           and e[1][1] == COT_ACTTERM), 15.0) is not None
            ck.add("yx16_OFF", ok_term and m.yx.get(16) == (0, False),
                   "via 2nd GI scan, yx16=%r" % (m.yx.get(16),))

        # 5. YK11 two-step (positive confirm; reset consumed by protectTask)
        m.send_i(asdu(T_C_DC_NA_1, 1, COT_ACT, 0, ioa=0x600A,
                      payload=bytes([0x82])), "I  C_DC YK11 select (QDC=82H)")
        ck.add("yk11_select_con", m.wait_event(actcon_of(T_C_DC_NA_1, False),
                                               5.0) is not None)
        m.send_i(asdu(T_C_DC_NA_1, 1, COT_ACT, 0, ioa=0x600A,
                      payload=bytes([0x02])), "I  C_DC YK11 execute (QDC=02H)")
        ck.add("yk11_execute_con", m.wait_event(actcon_of(T_C_DC_NA_1, False),
                                                5.0) is not None)

        # 6. C_RD read back the restored setpoint
        m.send_i(asdu(T_C_RD_NA_1, 1, COT_REQ, 0, ioa=addr),
                 "I  C_RD 0x%X" % addr)
        e = m.wait_event(lambda e: (e[0] == "yc" and e[1][0] == addr
                                    and e[1][3] == COT_REQ), 5.0)
        ck.add("rd_setpoint_10.0",
               e is not None and abs(e[1][1] - args.restore) < 0.1,
               "read=%.3f" % (e[1][1],) if e else "no response")

        m.close()
        rc = 0 if all(ok for _, ok, _ in ck.items) else 1
    except Exception as e:
        log("PROTOCOL ERROR: %r" % (e,))
        rc = 2

    n_ok = sum(1 for _, ok, _ in ck.items)
    log("SUMMARY: %d/%d checks PASS -> %s"
        % (n_ok, len(ck.items), "PASS" if rc == 0 else ("FAIL" if rc == 1 else "ERROR")))
    return rc


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