#!/usr/bin/env python3
# Isolated regtest chat round-trip: starts a throwaway regtest dragonxd (no network, own datadir),
# funds a shielded note, then uses build/linux/bin/chat_regtest_smoke to compose a chunked (multi-
# output) message + a Control memo, broadcasts them via z_sendmany, mines, harvests via
# z_listreceivedbyaddress, and feeds the memos back to the tool to verify they reassemble/decrypt.
# Nothing touches ~/.hush — a dedicated datadir, -regtest, and no-network flags keep it sandboxed.
import json, os, shutil, signal, subprocess, sys, time

BIN = os.environ.get("DRAGONX_REGTEST_DAEMON_DIR", "/home/d/dragonx/release/dragonx-1.3.0-linux-amd64")
DAEMON, CLI = f"{BIN}/dragonxd", f"{BIN}/dragonx-cli"
TOOL = os.environ.get("DRAGONX_CHAT_SMOKE", os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build/linux/bin/chat_regtest_smoke"))
RUN = os.environ.get("DRAGONX_REGTEST_RUN", "/tmp/dragonx-chat-regtest")
DD = f"{RUN}/dd"
PORT = 29201
USER, PW = "u", "p"
proc = None

def log(m): print(f"[regtest] {m}", flush=True)

def cli(method, *args, check=True):
    cmd = [CLI, "-regtest", f"-datadir={DD}", f"-conf={DD}/DRAGONX.conf", method, *[str(a) for a in args]]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        if check: raise RuntimeError(f"cli {method} failed: {r.stderr.strip()[:300]}")
        return None
    out = r.stdout.strip()
    try: return json.loads(out)
    except Exception: return out

def wait_rpc(timeout=60):
    t0 = time.time()
    while time.time() - t0 < timeout:
        if cli("getinfo", check=False) is not None: return True
        if proc and proc.poll() is not None:
            raise RuntimeError(f"daemon exited early (rc={proc.returncode}); see {DD}/regtest/debug.log")
        time.sleep(1)
    raise RuntimeError("daemon RPC never came up")

def wait_op(opid, timeout=180):
    t0 = time.time()
    while time.time() - t0 < timeout:
        res = cli("z_getoperationresult", json.dumps([opid]), check=False)
        if isinstance(res, list) and res:
            st = res[0].get("status")
            if st == "success": return res[0]
            if st == "failed": raise RuntimeError(f"op {opid} failed: {res[0].get('error')}")
        time.sleep(1)
    raise RuntimeError(f"op {opid} timed out")

def mine(n):
    cli("generate", n)
    time.sleep(1)

def tool(mode, *args, stdin=None):
    r = subprocess.run([TOOL, mode, *args], input=stdin, capture_output=True, text=True)
    sys.stderr.write(r.stderr)
    return r.stdout, r.returncode

def main():
    global proc
    if os.path.isdir(RUN): shutil.rmtree(RUN)
    os.makedirs(DD)
    with open(f"{DD}/DRAGONX.conf", "w") as f:
        f.write(f"rpcuser={USER}\nrpcpassword={PW}\nrpcport={PORT}\nusemnemonic=1\n")
    # Keep the real HOME so the daemon finds ~/.zcash-params (read-only); isolation is via the explicit
    # -datadir + -regtest + no-network flags, so no mainnet data is touched.
    log(f"starting regtest dragonxd (datadir {DD}, no network)")
    proc = subprocess.Popen(
        [DAEMON, "-regtest", f"-datadir={DD}", f"-conf={DD}/DRAGONX.conf", f"-rpcport={PORT}",
         f"-port={PORT+100}", "-listen=0", "-dnsseed=0", "-discover=0", "-maxconnections=0",
         "-usemnemonic=1", f"-asmap={BIN}/asmap.dat"],
        env=os.environ, stdout=open(f"{RUN}/daemon.out", "w"), stderr=subprocess.STDOUT)
    wait_rpc()
    log("daemon up; mining to maturity")
    mine(10)

    # Fund a shielded note: shield coinbase into a fresh z-addr (if auto-shield already did it, we just
    # find a funded note below).
    zfrom = cli("z_getnewaddress")
    op = cli("z_shieldcoinbase", "*", zfrom, check=False)
    if isinstance(op, dict) and op.get("opid"):
        try: wait_op(op["opid"])
        except Exception as e: log(f"shieldcoinbase: {e}")
    mine(3)
    # Find a funded z-note to send FROM.
    uns = cli("z_listunspent", 1) or []
    funded = [u for u in uns if float(u.get("amount", 0)) > 0.001]
    if not funded:
        log(f"no funded z-note yet; z_gettotalbalance={cli('z_gettotalbalance')}")
        raise RuntimeError("could not fund a shielded note")
    zfrom = funded[0]["address"]
    log(f"funded z-note: {zfrom} = {funded[0]['amount']}")

    z = cli("z_getnewaddress")                       # chat recipient (self for this round-trip)
    mnem = cli("z_exportmnemonic")
    if isinstance(mnem, dict): mnem = mnem.get("mnemonic") or mnem.get("seed_phrase") or ""
    assert mnem, "no mnemonic"
    pub, rc = tool("pubkey", mnem); pub = pub.strip()
    log(f"chat pubkey {pub[:16]}…, recipient {z[:16]}…")

    big = ("The quick brown fox jumps over the lazy dog. " * 34).strip()   # ~1560 chars -> multi-part
    results = {}

    def roundtrip(label, compose_args, expect_kind):
        out, rc = tool(*compose_args, )
        if rc != 0: raise RuntimeError(f"{label}: compose rc={rc}")
        recip, memos = None, []
        for line in out.splitlines():
            if line.startswith("RECIP "): recip = line[6:]
            elif line.startswith("HEADER "): memos.append(line[7:])
            elif line.startswith("PAYLOAD "): memos.append(line[8:])
        recipients = [{"address": recip, "amount": "0.00000000", "memo": "utf8:" + m} for m in memos]
        log(f"{label}: {len(memos)} memo outputs -> z_sendmany from {zfrom[:12]}…")
        op = cli("z_sendmany", zfrom, json.dumps(recipients), 1, "0.0001")
        opres = wait_op(op["opid"] if isinstance(op, dict) else op)
        sent_txid = (opres.get("result") or {}).get("txid")
        mine(2)
        recv = cli("z_listreceivedbyaddress", recip, 0) or []
        # Harvest is per-transaction in production; filter to THIS tx's outputs so each extract gets
        # exactly one message/control's memos (the recipient z-addr is reused across roundtrips).
        memostrs = [e.get("memoStr", "") for e in recv if e.get("memoStr") and e.get("txid") == sent_txid]
        log(f"{label}: tx {str(sent_txid)[:12]}… -> {len(memostrs)} of {len(recv)} received memos are this tx")
        ext, rc = tool("extract", mnem, stdin="\n".join(memostrs) + "\n")
        results[label] = (ext.strip(), rc)
        return ext.strip(), rc

    # 1) multi-part message round-trip
    msg_out, msg_rc = roundtrip("MESSAGE", ("compose", mnem, pub, z, z, "cid-1", big), "msg")
    rec = ""
    for line in msg_out.splitlines():
        if line.startswith("PLAINTEXT "): rec = line[10:]
    ok_msg = (rec == big)
    log(f"MESSAGE round-trip: {'OK' if ok_msg else 'FAIL'} (recovered {len(rec)}/{len(big)} bytes)")

    # 2) Control(del) memo round-trip
    ctl_out, ctl_rc = roundtrip("CONTROL", ("control", mnem, z, "cid-7", "del", "peerlabel"), "ctl")
    ok_ctl = "CONTROL op=del" in ctl_out and "cid=cid-7" in ctl_out
    log(f"CONTROL round-trip: {'OK' if ok_ctl else 'FAIL'} ({ctl_out[:80]})")

    print("\n==== RESULT ====")
    print(f"multi-part MESSAGE through live daemon: {'PASS' if ok_msg else 'FAIL'}")
    print(f"Control(del) memo through live daemon:  {'PASS' if ok_ctl else 'FAIL'}")
    return 0 if (ok_msg and ok_ctl) else 1

if __name__ == "__main__":
    rc = 1
    try:
        rc = main()
    except Exception as e:
        log(f"ERROR: {e}")
    finally:
        if proc and proc.poll() is None:
            log("stopping daemon")
            cli("stop", check=False); time.sleep(3)
            if proc.poll() is None: proc.terminate(); time.sleep(2)
            if proc.poll() is None: proc.kill()
    sys.exit(rc)