"""Exercise the compiled program over JSON-RPC in a disposable local Surfpool SVM."""
import argparse
import base64
import hashlib
import json
import struct
import subprocess
import time
import urllib.error
import urllib.request
from pathlib import Path

from solders.hash import Hash
from solders.instruction import AccountMeta, Instruction
from solders.keypair import Keypair
from solders.pubkey import Pubkey
from solders.system_program import CreateAccountParams, TransferParams, create_account, transfer
from solders.transaction import Transaction

ROOT = Path(__file__).resolve().parent
SYSTEM = Pubkey.default()
LOADER = Pubkey.from_string("BPFLoaderUpgradeab1e11111111111111111111111")
SYS_RENT = Pubkey.from_string("SysvarRent111111111111111111111111111111111")
SYS_CLOCK = Pubkey.from_string("SysvarC1ock11111111111111111111111111111111")
AUDIENCE = hashlib.sha256(b"https://solmate.io").digest()
BIN = ROOT / "build/solmate_registry.so"


class RpcError(RuntimeError):
    pass


class Rpc:
    def __init__(self, url, interval=0):
        self.url = url
        self.interval = interval
        self.last_call = 0

    def call(self, method, *params):
        payload = json.dumps(dict(jsonrpc="2.0", id=1, method=method, params=params)).encode()
        req = urllib.request.Request(self.url, payload, {"Content-Type": "application/json"})
        for attempt in range(5):
            time.sleep(max(0, self.last_call + self.interval - time.monotonic()))
            self.last_call = time.monotonic()
            try:
                with urllib.request.urlopen(req, timeout=25) as res:
                    body = json.load(res)
                break
            except urllib.error.HTTPError as error:
                if error.code != 429 or attempt == 4:
                    raise
                retry_after = error.headers.get("Retry-After", "")
                delay = int(retry_after) if retry_after.isdigit() else 2 ** (attempt + 1)
                time.sleep(max(1, delay))
        if "error" in body:
            raise RpcError(json.dumps(body["error"]))
        return body["result"]

    def confirmed(self, sig, timeout=35):
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            st = self.call("getSignatureStatuses", [sig])["value"][0]
            if st and st["err"]:
                raise RpcError(str(st["err"]))
            if st and st["confirmationStatus"] in ("confirmed", "finalized"):
                return sig
            time.sleep(0.2)
        raise TimeoutError(sig)

    def send(self, instructions, payer, others=()):
        h = self.call("getLatestBlockhash", {"commitment": "confirmed"})["value"]["blockhash"]
        tx = Transaction.new_signed_with_payer(instructions, payer.pubkey(), [payer, *others], Hash.from_string(h))
        encoded = base64.b64encode(bytes(tx)).decode()
        sig = self.call("sendTransaction", encoded, {"encoding": "base64", "preflightCommitment": "confirmed"})
        return self.confirmed(sig)

    def account(self, key):
        info = self.call("getAccountInfo", str(key), {"commitment": "confirmed", "encoding": "base64"})["value"]
        if info:
            info["data"] = base64.b64decode(info["data"][0])
        return info


def meta(key, signer=False, write=False):
    return AccountMeta(key, signer, write)


def account_ix(rpc, payer, target, length):
    return create_account(CreateAccountParams(
        from_pubkey=payer.pubkey(), to_pubkey=target.pubkey(), space=length, owner=LOADER,
        lamports=rpc.call("getMinimumBalanceForRentExemption", length),
    ))


def read_upload(rpc, buffer_address, authority, binary):
    info = rpc.account(buffer_address)
    header = struct.pack("<IB", 1, 1) + bytes(authority)
    if not (info and info["owner"] == str(LOADER) and not info["executable"]
            and info["data"][:37] == header and len(info["data"]) == 37 + len(binary)):
        raise ValueError("Upload buffer does not match the expected loader, authority and size.")
    uploaded = info["data"][37:]
    for offset in range(0, len(binary), 800):
        chunk = uploaded[offset:offset + 800]
        if any(chunk) and chunk != binary[offset:offset + 800]:
            raise ValueError("Upload buffer contains bytes from a different program.")
    return info


def deploy(rpc, payer, program_key, binary, buffer_address=None):
    """Deploy via the standard upgradeable loader. Caller supplies all test keys."""
    who = payer.pubkey()
    program = program_key.pubkey()
    data_key = Pubkey.find_program_address([bytes(program)], LOADER)[0]
    uploaded = b""
    if buffer_address is None:
        buffer = Keypair()
        buffer_address = buffer.pubkey()
        print("Deployment buffer:", buffer_address, flush=True)
        init = Instruction(LOADER, struct.pack("<I", 0), [meta(buffer_address, write=True), meta(who)])
        rpc.send([account_ix(rpc, payer, buffer, 37 + len(binary)), init], payer, [buffer])
    else:
        uploaded = read_upload(rpc, buffer_address, who, binary)["data"][37:]
    for offset in range(0, len(binary), 800):
        chunk = binary[offset:offset + 800]
        if uploaded[offset:offset + len(chunk)] == chunk:
            continue
        write = Instruction(LOADER, struct.pack("<IIQ", 1, offset, len(chunk)) + chunk,
                            [meta(buffer_address, write=True), meta(who, signer=True)])
        rpc.send([write], payer)
    launch = Instruction(LOADER, struct.pack("<IQ", 2, len(binary)), [
        meta(who, True, True), meta(data_key, write=True), meta(program, write=True),
        meta(buffer_address, write=True), meta(SYS_RENT), meta(SYS_CLOCK), meta(SYSTEM), meta(who, True),
    ])
    sig = rpc.send([account_ix(rpc, payer, program_key, 36), launch], payer, [program_key])
    print("Deployment transaction:", sig, flush=True)
    program_info, data_info = rpc.account(program), rpc.account(data_key)
    assert program_info["executable"] and program_info["owner"] == str(LOADER)
    assert program_info["data"] == struct.pack("<I", 2) + bytes(data_key)
    assert data_info["owner"] == str(LOADER) and data_info["data"][45:] == binary
    assert data_info["data"][12:45] == b"\x01" + bytes(who)
    time.sleep(0.6)
    return sig


def state_key(program, vault):
    return Pubkey.find_program_address([b"vault", bytes(vault)], program)[0]


def grant_key(program, vault, delegate, asset=SYSTEM):
    return Pubkey.find_program_address([b"grant", bytes(vault), bytes(delegate), AUDIENCE,
                                        bytes([asset != SYSTEM]), bytes(asset)], program)[0]


def grant_ix(program, vault, delegate, expiry, asset=SYSTEM):
    fields = bytes([0, asset != SYSTEM]) + bytes(delegate) + AUDIENCE + bytes(asset) + struct.pack("<q", expiry)
    return Instruction(program, fields, [meta(vault, True, True), meta(state_key(program, vault), write=True),
                                         meta(grant_key(program, vault, delegate, asset), write=True), meta(SYSTEM)])


def must_fail(rpc, instructions, payer):
    try:
        rpc.send(instructions, payer)
    except RpcError as err:
        if "InstructionError" not in str(err):
            raise
        return
    raise AssertionError("a disallowed operation succeeded")


def exercise(rpc, payer, program):
    """Returns a signed grant and revoke receipt. Keys are isolated test fixtures."""
    who = payer.pubkey()
    other = Keypair()
    delegate, asset = other.pubkey(), Pubkey.new_unique()
    g, s = grant_key(program, who, delegate), state_key(program, who)
    now = rpc.call("getBlockTime", rpc.call("getSlot", {"commitment": "confirmed"}))
    expiry = now + 86400
    initial = grant_ix(program, who, delegate, expiry)
    # Pre-funding a future record must not prevent its authorized creation.
    rent0 = rpc.call("getMinimumBalanceForRentExemption", 0)
    rpc.send([transfer(TransferParams(from_pubkey=who, to_pubkey=g, lamports=rent0)),
              transfer(TransferParams(from_pubkey=who, to_pubkey=delegate, lamports=rent0 * 2))], payer)
    created = rpc.send([initial], payer)
    st, gr = rpc.account(s), rpc.account(g)
    assert st["owner"] == gr["owner"] == str(program)
    assert st["data"] == b"\x01" + bytes(who) + bytes(8)
    assert gr["data"] == bytes([2]) + bytes(who) + bytes(delegate) + AUDIENCE + bytes(33) + struct.pack("<Qq", 0, expiry)
    assert gr["lamports"] == rpc.call("getMinimumBalanceForRentExemption", 146)
    # The sole mint scope has its own deterministic grant.
    rpc.send([grant_ix(program, who, delegate, expiry, asset)], payer)
    narrow = grant_key(program, who, delegate, asset)
    assert rpc.account(narrow)["data"][97:130] == b"\x01" + bytes(asset)
    must_fail(rpc, [grant_ix(program, who, delegate, now - 1)], payer)
    must_fail(rpc, [grant_ix(program, who, delegate, now + 400 * 86400)], payer)
    invalid_scope = bytearray(initial.data)
    invalid_scope[1] = 2
    must_fail(rpc, [Instruction(program, bytes(invalid_scope), initial.accounts)], payer)
    must_fail(rpc, [Instruction(program, initial.data, [*initial.accounts[:3], meta(SYS_CLOCK)])], payer)
    must_fail(rpc, [Instruction(program, initial.data, [*initial.accounts[:2], meta(narrow, write=True), meta(SYSTEM)])], payer)
    must_fail(rpc, [Instruction(program, b"\x01", [meta(delegate, True, True), meta(g, write=True)])], other)
    must_fail(rpc, [Instruction(program, b"\x02", [meta(delegate, True), meta(s, write=True)])], other)
    unsigned_owner = Instruction(program, b"\x02", [meta(who), meta(s, write=True)])
    must_fail(rpc, [unsigned_owner], other)
    rpc.send([Instruction(program, b"\x02", [meta(who, True), meta(s, write=True)])], payer)
    assert rpc.account(s)["data"][-8:] == struct.pack("<Q", 1)
    assert rpc.account(g)["data"][130:138] == bytes(8)
    rpc.send([grant_ix(program, who, delegate, expiry + 60)], payer)
    assert rpc.account(g)["data"][130:] == struct.pack("<Qq", 1, expiry + 60)
    assert rpc.account(narrow)["data"][130:138] == bytes(8)
    revoke = Instruction(program, b"\x01", [meta(who, True, True), meta(g, write=True)])
    balance = rpc.account(who)["lamports"]
    # Sponsor the fee from a distinct test wallet to measure the exact refund.
    closed = rpc.send([revoke], other, [payer])
    assert rpc.account(who)["lamports"] == balance + gr["lamports"]
    assert rpc.account(g) is None
    # Close an old-epoch record and make sure later refunding cannot restore its grant.
    rpc.send([Instruction(program, b"\x01", [meta(who, True, True), meta(narrow, write=True)]),
              transfer(TransferParams(from_pubkey=who, to_pubkey=narrow, lamports=rent0))], payer)
    assert rpc.account(narrow)["owner"] == str(SYSTEM) and rpc.account(narrow)["data"] == b""
    assert rpc.account(s)["data"][-8:] == struct.pack("<Q", 1)
    return created, closed


def main():
    ap = argparse.ArgumentParser(__doc__)
    ap.add_argument("--surfpool", required=True, type=Path)
    ap.add_argument("--port", type=int, default=18987)
    ap.add_argument("--serve", action="store_true",
                    help="Keep the tested offline validator running until Ctrl+C.")
    ap.add_argument("--fund", type=Pubkey.from_string,
                    help="Optional public wallet address to fund on this offline validator.")
    args = ap.parse_args()
    rpc = Rpc(f"http://127.0.0.1:{args.port}")
    binary = BIN.read_bytes()
    cache = ROOT / ".build"
    cache.mkdir(exist_ok=True)
    args_ = [args.surfpool.resolve(), "start", "--offline", "--no-deploy", "--ci",
             "--host", "127.0.0.1", "--port", str(args.port), "--ws-port", str(args.port + 1)]
    with open(cache / "surfpool-tests.log", "w", encoding="utf-8") as log:
        child = subprocess.Popen(args_, cwd=cache, stdout=log, stderr=subprocess.STDOUT,
                                 creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        try:
            for _ in range(100):
                if child.poll() is not None:
                    raise RuntimeError("Surfpool exited; inspect .build/surfpool-tests.log")
                try:
                    if rpc.call("getHealth") == "ok":
                        break
                except (OSError, RpcError):
                    time.sleep(0.1)
            payer, program = Keypair(), Keypair()
            rpc.confirmed(rpc.call("requestAirdrop", str(payer.pubkey()), 2_000_000_000))
            receipt = deploy(rpc, payer, program, binary)
            granted, revoked = exercise(rpc, payer, program.pubkey())
            print("PASS: verified deployment bytes; create/renew; wallet/asset scopes; bad scope PDA; "
                  "expiry; unsigned/foreign vault; scoped and bulk revocation; rent refund and no revival.")
            print("Local receipts:", receipt, granted, revoked)
            if args.fund:
                rpc.confirmed(rpc.call("requestAirdrop", str(args.fund), 2_000_000_000))
            if args.serve:
                print("Local RPC:", rpc.url, "\nLocal program:", program.pubkey(),
                      "\nLocal genesis:", rpc.call("getGenesisHash"),
                      "\nPress Ctrl+C to stop the temporary chain.", flush=True)
                try:
                    while child.poll() is None:
                        time.sleep(1)
                except KeyboardInterrupt:
                    pass
        finally:
            child.terminate()
            child.wait(timeout=10)


if __name__ == "__main__":
    main()