#!/usr/bin/env python3
"""3-node CPU test network for the difficulty attacks (ports 27700 and above, appdir /tmp/igneum-diff-attacks).

Node 1 (gRPC 27700, P2P 27701, wRPC JSON 27702) is the honest seed, node 2 (27710/27711) honest, node 3
(27720/27721) the attacker's node: it is started with IGNEUM_ATTACK_TS_OFFSET_MS so every template it hands
out is stamped at the given offset (the diff-attacks branch of vendor/igneum-node reads that variable in the
template builder only; the validation rules are untouched). Miners are `igneum-miner mine --engine igneum-pow`
CPU miners with their own payout labels, started and stopped by pid.

    python3 testnet.py --rule igneum --scenario ts --minutes 15 --out /tmp/igneum-diff-attacks/ts-igneum
    python3 testnet.py --rule kaspa  --scenario hop --minutes 15 --out /tmp/igneum-diff-attacks/hop-kaspa
    python3 testnet.py --analyse /tmp/igneum-diff-attacks/ts-igneum

Scenarios
  ts   honest miner A on node 1 all run; forger F mines on node 2 (honest stamps) for the first third, then on
       node 3 (stamps 130 s ahead) for the rest. Threads chosen for about a 50% share.
  hop  honest miner A on node 1 all run; hopper B (same threads) on node 2, on while the difficulty read from
       node 1 is below its trailing mean (1% band), off when above, re-evaluated every second; from minute 3.
Everything started here is stopped at the end (SIGTERM, then SIGKILL) by pid; nothing else is touched.
"""
import argparse
import asyncio
import csv
import json
import os
import signal
import subprocess
import sys
import time

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
BIN = os.path.join(ROOT, "vendor", "igneum-node-diff-attacks", "target", "release")
IGNEUMD = os.path.join(BIN, "igneumd")
MINER = os.path.join(BIN, "igneum-miner")
GENESIS_BITS = 0x1f010000   # 2^16 expected hashes per block: about 1 block/s for one 4-thread CPU miner on this Mac
NODES = [
    {"name": "node1", "grpc": 27700, "p2p": 27701, "json": 27702},
    {"name": "node2", "grpc": 27710, "p2p": 27711, "json": None},
    {"name": "node3", "grpc": 27720, "p2p": 27721, "json": None},
]
started = []   # Popen objects we own


def log(out, msg):
    line = "%s %s" % (time.strftime("%H:%M:%S"), msg)
    print(line, flush=True)
    with open(os.path.join(out, "schedule.log"), "a") as f:
        f.write(line + "\n")


def start_node(n, out, rule, env_extra=None):
    override = os.path.join(out, "override.json")
    with open(override, "w") as f:
        json.dump({"difficulty_rule": rule, "genesis_bits": GENESIS_BITS}, f)
    appdir = os.path.join(out, n["name"])
    args = [IGNEUMD, "--devnet", "--nodnsseed", "--disable-upnp", "--enable-unsynced-mining", "--nologfiles",
            "--override-params-file=" + override, "--appdir=" + appdir,
            "--rpclisten=127.0.0.1:%d" % n["grpc"], "--listen=127.0.0.1:%d" % n["p2p"]]
    if n["json"]:
        args.append("--rpclisten-json=127.0.0.1:%d" % n["json"])
    if n["name"] != "node1":
        args.append("--connect=127.0.0.1:%d" % NODES[0]["p2p"])
    env = dict(os.environ)
    env.update(env_extra or {})
    lf = open(os.path.join(out, n["name"] + ".log"), "w")
    p = subprocess.Popen(["nice", "-n", "19"] + args, stdout=lf, stderr=subprocess.STDOUT, env=env)
    started.append(p)
    log(out, "%s pid %d gRPC %d%s" % (n["name"], p.pid, n["grpc"], " env " + json.dumps(env_extra) if env_extra else ""))
    return p


def start_miner(out, label, grpc, threads, secs):
    lf = open(os.path.join(out, "miner-%s.log" % label), "a")
    args = ["nice", "-n", "19", MINER, "mine", "grpc://127.0.0.1:%d" % grpc, str(threads), str(secs), label,
            "--engine", "igneum-pow", "--payout-label", label, "--status-secs", "10"]
    p = subprocess.Popen(args, stdout=lf, stderr=subprocess.STDOUT)
    started.append(p)
    log(out, "miner %s pid %d on %d threads %d for %d s" % (label, p.pid, grpc, threads, secs))
    return p


def stop(p, out, what):
    if p.poll() is None:
        p.send_signal(signal.SIGTERM)
        try:
            p.wait(10)
        except subprocess.TimeoutExpired:
            p.kill()
            p.wait(5)
    log(out, "stopped %s pid %d code %s" % (what, p.pid, p.returncode))


async def rpc(port, method, params=None):
    import websockets
    async with websockets.connect("ws://127.0.0.1:%d" % port, max_size=None) as ws:
        await ws.send(json.dumps({"id": 1, "method": method, "params": params or {}}))
        m = json.loads(await ws.recv())
        if m.get("error"):
            raise RuntimeError(m["error"])
        return m["params"]


def dag_info(port):
    return asyncio.run(rpc(port, "getBlockDagInfo"))


def export(port, path):
    info = dag_info(port)
    low = info.get("pruningPointHash") or info["pruningPointHash"]
    seen = set()
    rows = []
    while True:
        r = asyncio.run(rpc(port, "getBlocks", {"lowHash": low, "includeBlocks": True, "includeTransactions": False}))
        blocks = r.get("blocks") or []
        new = 0
        for b in blocks:
            h = b["header"]
            hs = h.get("hash") or (b.get("verboseData") or {}).get("hash")
            if hs in seen:
                continue
            seen.add(hs)
            new += 1
            v = b.get("verboseData") or {}
            rows.append({"hash": hs, "daa_score": h["daaScore"], "blue_score": h["blueScore"], "timestamp_ms": h["timestamp"],
                         "bits": "0x%08x" % int(h["bits"]), "difficulty": v.get("difficulty"), "is_chain_block": 1 if v.get("isChainBlock") else 0,
                         "parents": len((h.get("parentsByLevel") or [[]])[0]), "selected_parent": v.get("selectedParentHash"),
                         "vote_key_hash": h.get("voteKeyHash"), "blue_work": h.get("blueWork")})
        if not blocks or new == 0:
            break
        low = (blocks[-1]["header"].get("hash") or blocks[-1]["verboseData"]["hash"])
        if len(blocks) < 2:
            break
    rows.sort(key=lambda r: (int(r["blue_score"]), int(r["daa_score"]), int(r["timestamp_ms"])))
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)
    return rows, info


def wait_grpc(out, n, secs=60):
    lp = os.path.join(out, n["name"] + ".log")
    for _ in range(secs * 2):
        try:
            if "GRPC Server starting" in open(lp).read():
                return True
        except FileNotFoundError:
            pass
        time.sleep(0.5)
    return False


def run(args):
    out = args.out
    os.makedirs(out, exist_ok=True)
    total = args.minutes * 60
    log(out, "start rule=%s scenario=%s minutes=%d genesis_bits=%#x" % (args.rule, args.scenario, args.minutes, GENESIS_BITS))
    try:
        env3 = {"IGNEUM_ATTACK_TS_OFFSET_MS": str(args.ts_offset_ms)} if args.scenario == "ts" else None
        start_node(NODES[0], out, args.rule)
        wait_grpc(out, NODES[0])
        start_node(NODES[1], out, args.rule)
        start_node(NODES[2], out, args.rule, env3)
        for n in NODES[1:]:
            wait_grpc(out, n)
        time.sleep(3)
        t0 = time.time()
        log(out, "miners from now; block 1 expected shortly")
        a = start_miner(out, "A", NODES[0]["grpc"], args.threads, total)
        if args.scenario == "ts":
            third = total // 3
            f1 = start_miner(out, "F", NODES[1]["grpc"], args.forger_threads, third)
            time.sleep(third)
            stop(f1, out, "miner F (honest leg)")
            f2 = start_miner(out, "F", NODES[2]["grpc"], args.forger_threads, total - third - 2)
            log(out, "forger now on node3 with offset %d ms" % args.ts_offset_ms)
            while time.time() - t0 < total:
                time.sleep(5)
            stop(f2, out, "miner F (forging leg)")
        elif args.scenario == "hop":
            warm = 180
            time.sleep(warm)
            hop = None
            hist = []
            on_since = None
            on_total = 0.0
            switches = 0
            while time.time() - t0 < total:
                try:
                    d = float(dag_info(NODES[0]["json"])["difficulty"])
                except Exception as e:
                    log(out, "dag info failed: %s" % e)
                    time.sleep(1)
                    continue
                now = time.time()
                hist.append((now, d))
                hist = [x for x in hist if x[0] >= now - 600]
                avg = sum(x[1] for x in hist) / len(hist)
                want = (d < avg * 0.99) if hop is None or hop.poll() is not None else (d <= avg * 1.01)
                running = hop is not None and hop.poll() is None
                if want and not running:
                    hop = start_miner(out, "B", NODES[1]["grpc"], args.threads, int(total - (now - t0)) + 5)
                    on_since = now
                    switches += 1
                    log(out, "hop ON  D=%.0f avg=%.0f" % (d, avg))
                elif not want and running:
                    stop(hop, out, "miner B")
                    on_total += now - on_since
                    on_since = None
                    switches += 1
                    log(out, "hop OFF D=%.0f avg=%.0f" % (d, avg))
                with open(os.path.join(out, "hop.csv"), "a") as f:
                    f.write("%.3f,%.1f,%.1f,%d\n" % (now, d, avg, 1 if running else 0))
                time.sleep(1)
            if hop is not None and hop.poll() is None:
                stop(hop, out, "miner B")
                on_total += time.time() - on_since
            log(out, "hopper on %.0f s of %d, %d switches" % (on_total, total - warm, switches))
        else:
            while time.time() - t0 < total:
                time.sleep(5)
        stop(a, out, "miner A")
        time.sleep(2)
        rows, info = export(NODES[0]["json"], os.path.join(out, "record.csv"))
        log(out, "dag %s" % json.dumps({k: info.get(k) for k in ("blockCount", "headerCount", "difficulty", "tipHashes", "virtualDaaScore")}))
        log(out, "wrote %d rows" % len(rows))
    finally:
        for p in reversed(started):
            stop(p, out, "pid")
        log(out, "end")


def analyse(out):
    """Per-30-s table in wall time. Honest miner A's header stamps are real time (its node's clock), so every
    chain block is placed at the stamp of the nearest earlier A block in blue order; the forger's own stamps are
    reported as an offset against that clock."""
    rows = list(csv.DictReader(open(os.path.join(out, "record.csv"))))
    rows = [r for r in rows if int(r["blue_score"]) > 0]
    chain = sorted([r for r in rows if r["is_chain_block"] == "1"], key=lambda r: int(r["blue_score"]))
    keys = {}
    summary = {}
    for fn in os.listdir(out):
        if fn.startswith("miner-") and fn.endswith(".log"):
            label = fn[6:-4]
            for line in open(os.path.join(out, fn)):
                if "vote_key_hash=" in line and label not in keys:
                    keys[label] = line.split("vote_key_hash=")[1].split()[0]
                if "MINER SUMMARY" in line:
                    summary.setdefault(label, []).append(line.strip())
    who = {v: k for k, v in keys.items()}
    t = None
    t0 = None
    bins = {}
    offs = []
    for r in chain:
        label = who.get(r["vote_key_hash"], "?")
        ts = int(r["timestamp_ms"]) / 1000.0
        if label == "A":
            t = ts
        if t is None:
            continue
        if t0 is None:
            t0 = t
        b = int((t - t0) // 30)
        d = bins.setdefault(b, {"n": 0, "F": 0, "dsum": 0.0, "off": []})
        d["n"] += 1
        d["F"] += 1 if label == "F" else 0
        d["dsum"] += float(r["difficulty"] or 0)
        if label == "F":
            d["off"].append(ts - t)
    print("blocks %d, chain blocks %d, miners %s" % (len(rows), len(chain), sorted(keys)))
    print("| from s | chain blocks | blocks/s | F share | mean difficulty | F stamp offset s |")
    print("|---|---|---|---|---|---|")
    for b in sorted(bins):
        d = bins[b]
        off = ("%.0f" % (sum(d["off"]) / len(d["off"]))) if d["off"] else ""
        print("| %d | %d | %.2f | %.2f | %s | %s |" % (b * 30, d["n"], d["n"] / 30.0, d["F"] / d["n"], "{:,.0f}".format(d["dsum"] / d["n"]), off))
    for k, v in sorted(summary.items()):
        for line in v:
            print(k, line)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rule", default="igneum-dual", help="igneum-dual or kaspa-sampled")
    ap.add_argument("--scenario", default="ts", help="ts or hop")
    ap.add_argument("--minutes", type=int, default=15)
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--forger-threads", type=int, default=4)
    ap.add_argument("--ts-offset-ms", type=int, default=130000)
    ap.add_argument("--out", default="/tmp/igneum-diff-attacks/run")
    ap.add_argument("--analyse", default="")
    args = ap.parse_args()
    if args.analyse:
        analyse(args.analyse)
        return
    if args.rule == "igneum":
        args.rule = "igneum-dual"
    if args.rule == "kaspa":
        args.rule = "kaspa-sampled"
    run(args)


if __name__ == "__main__":
    main()