#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""可回滚的 I210 网卡参数实验：停环 → 改高级属性 → 恢复 OP → 标定跳拍率。

默认只测、不永久破坏：原值会写入 restore json；`--restore` 按快照还原。
改网卡会触发短暂链路重协商，期间从站可能掉 OP，脚本会等 link+OP 恢复。
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
import urllib.request

ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))
LAB = os.path.join(ROOT, "lab", "skip_meter")
NIC = "以太网 2"
BASE = os.environ.get("ECAT_BASE", "http://127.0.0.1:8090")


def post(path, body=None, timeout=8.0):
    data = json.dumps(body or {}).encode()
    req = urllib.request.Request(
        BASE + path, data=data, headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode())


def ps(cmd, timeout=60):
    r = subprocess.run(
        ["powershell", "-NoProfile", "-Command", cmd],
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
        timeout=timeout,
    )
    return (r.stdout or "").strip(), (r.stderr or "").strip()


def nic_snapshot():
    out, err = ps(
        f"Get-NetAdapterAdvancedProperty -Name '{NIC}' | "
        "Select-Object RegistryKeyword,DisplayValue,RegistryValue | "
        "ConvertTo-Json -Compress"
    )
    if not out:
        return []
    data = json.loads(out)
    if isinstance(data, dict):
        return [data]
    return data


def nic_set(keyword: str, value: str):
    # value as string; Intel accepts "0"/"1"/numbers
    cmd = (
        f"Set-NetAdapterAdvancedProperty -Name '{NIC}' "
        f"-RegistryKeyword '{keyword}' -RegistryValue '{value}' -NoRestart"
    )
    out, err = ps(cmd)
    return out, err


def nic_restart():
    out, err = ps(
        f"Restart-NetAdapter -Name '{NIC}' -Confirm:$false; Start-Sleep -Seconds 2; "
        f"Get-NetAdapter -Name '{NIC}' | Select-Object Status,LinkSpeed | ConvertTo-Json -Compress"
    )
    return out, err


def wait_link(timeout=20.0):
    t0 = time.time()
    while time.time() - t0 < timeout:
        out, _ = ps(
            f"Get-NetAdapter -Name '{NIC}' | Select-Object Status,LinkSpeed | ConvertTo-Json -Compress"
        )
        try:
            j = json.loads(out)
            if isinstance(j, list):
                j = j[0]
            if str(j.get("Status", "")).lower() == "up":
                return True, j
        except Exception:
            pass
        time.sleep(0.5)
    return False, None


def wait_ecat_healthy(timeout=25.0):
    t0 = time.time()
    last = None
    while time.time() - t0 < timeout:
        try:
            h = post("/api/native/health")
            last = h
            if h.get("loop") == 1 and int(h.get("wkc") or 0) >= 3 and h.get("op") == 1:
                # also require cycles advancing
                c0 = h.get("cycles") or 0
                time.sleep(0.4)
                h2 = post("/api/native/health")
                c1 = h2.get("cycles") or 0
                if c1 > c0:
                    return True, h2
        except Exception as e:
            last = {"err": str(e)}
        time.sleep(0.5)
    return False, last


def stop_bus():
    logs = []
    for path in ("/api/motor/stop", "/api/op/stop"):
        try:
            r = post(path)
            logs.append((path, r.get("ok"), (r.get("log") or r)[:2]))
        except Exception as e:
            logs.append((path, False, str(e)))
    return logs


def start_bus():
    logs = []
    try:
        r = post("/api/op/start")
        logs.append(("/api/op/start", r.get("ok"), (r.get("log") or [""])[:2]))
    except Exception as e:
        logs.append(("/api/op/start", False, str(e)))
    try:
        r = post("/api/motor/start")
        logs.append(("/api/motor/start", r.get("ok"), (r.get("log") or [""])[:2]))
    except Exception as e:
        logs.append(("/api/motor/start", False, str(e)))
    return logs


def meter(tag: str, secs: float):
    # reuse skip_rate_meter as subprocess for isolation
    cmd = [
        sys.executable,
        os.path.join(ROOT, "scripts", "skip_rate_meter.py"),
        "--tag",
        tag,
        "--secs",
        str(secs),
    ]
    r = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8", errors="replace")
    # last json file for tag
    files = sorted(
        [f for f in os.listdir(LAB) if f.startswith(tag + "_") and f.endswith(".json")]
    )
    data = None
    if files:
        with open(os.path.join(LAB, files[-1]), encoding="utf-8") as f:
            data = json.load(f)
    return data, r.stdout, r.stderr


def apply_batch(changes: list[tuple[str, str]]):
    results = []
    for k, v in changes:
        out, err = nic_set(k, v)
        results.append({"key": k, "value": v, "out": out[:300], "err": err[:300]})
    return results


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--restore", action="store_true", help="restore from snap json")
    ap.add_argument("--snap", default=None, help="snapshot json for restore")
    ap.add_argument("--exp", default=None, help="named experiment")
    ap.add_argument("--secs", type=float, default=16.0)
    ap.add_argument("--list-exp", action="store_true")
    args = ap.parse_args()

    os.makedirs(LAB, exist_ok=True)

    # experiment table: list of (keyword, value)
    EXPS = {
        "itr_off": [
            ("*InterruptModeration", "0"),
            ("ITR", "0"),
        ],
        "itr_low": [
            ("*InterruptModeration", "1"),
            ("ITR", "400"),  # valid: 65535,3600,2000,950,400,200,0
        ],
        "itr_200": [
            ("*InterruptModeration", "1"),
            ("ITR", "200"),
        ],
        "itr_950": [
            ("*InterruptModeration", "1"),
            ("ITR", "950"),
        ],
        "eee_off": [
            ("EEELinkAdvertisement", "0"),
        ],
        "flow_off": [
            ("*FlowControl", "0"),
        ],
        "txbuf_max": [
            ("*TransmitBuffers", "4096"),
        ],
        "rxbuf_max": [
            ("*ReceiveBuffers", "1024"),
        ],
        "combo_lowlat": [
            ("*InterruptModeration", "0"),
            ("ITR", "0"),
            ("EEELinkAdvertisement", "0"),
            ("*FlowControl", "0"),
            ("*TransmitBuffers", "2048"),
            ("*ReceiveBuffers", "512"),
        ],
    }

    if args.list_exp:
        print(json.dumps({k: v for k, v in EXPS.items()}, indent=2))
        return 0

    # always keep a fresh snapshot before mutating
    snap = nic_snapshot()
    snap_path = os.path.join(LAB, f"nic_before_{int(time.time())}.json")
    with open(snap_path, "w", encoding="utf-8") as f:
        json.dump(snap, f, ensure_ascii=False, indent=2)
    print(f"NIC snapshot → {snap_path}")

    if args.restore:
        path = args.snap or snap_path
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
        # stop, apply all, restart
        print("restore: stop bus")
        print(stop_bus())
        n = 0
        for row in data:
            kw = row.get("RegistryKeyword")
            val = row.get("RegistryValue")
            if not kw or val is None:
                continue
            if isinstance(val, list):
                val = val[0] if val else ""
            val = str(val)
            apply_batch([(kw, val)])
            n += 1
        print("restore applied", n, "keys; restarting NIC")
        print(nic_restart())
        print(wait_link())
        print("start bus")
        print(start_bus())
        ok, h = wait_ecat_healthy()
        print("healthy", ok, h)
        return 0

    exp = args.exp or "itr_off"
    if exp not in EXPS:
        print("unknown exp", exp, "use --list-exp")
        return 2
    changes = EXPS[exp]

    print(f"\n=== EXP {exp} {changes} ===")
    print("stop bus...")
    print(stop_bus())
    time.sleep(0.3)
    print("apply...")
    print(apply_batch(changes))
    print("restart nic...")
    print(nic_restart())
    ok, info = wait_link()
    print("link", ok, info)
    print("start bus...")
    print(start_bus())
    ok, h = wait_ecat_healthy(timeout=30)
    print("healthy", ok, h)
    if not ok:
        print("ABORT: bus not healthy after change")
        return 1

    # settle longer: NIC restart + re-enumeration leaves ~2–5s of elevated skip
    time.sleep(8.0)
    data, out, err = meter(f"exp_{exp}", args.secs)
    print(out)
    if err:
        print("meter stderr", err[:500])
    if data:
        summary = {
            "exp": exp,
            "changes": changes,
            "late_per_s": data.get("late_per_s"),
            "skip_per_s": data.get("skip_per_s"),
            "late_frac_pct": data.get("late_frac_pct"),
            "send_stall_frac": data.get("send_stall_frac"),
            "rx_stall_frac": data.get("rx_stall_frac"),
            "achieved_hz": data.get("achieved_hz"),
            "n_late": data.get("n_late"),
        }
        sp = os.path.join(LAB, f"summary_{exp}_{int(time.time())}.json")
        with open(sp, "w", encoding="utf-8") as f:
            json.dump(summary, f, ensure_ascii=False, indent=2)
        print("SUMMARY", json.dumps(summary, ensure_ascii=False))
        print("WROTE", sp)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())