"""Find self-contained arbitrage transactions in a sample of Solana blocks.

A transaction counts as arbitrage when it succeeded, touched at least one known
DEX program, and left its fee payer with more SOL/USDC/USDT while every other
token balance of the payer ended unchanged (a closed cycle). Net result includes
fees and tips, because they are paid from the payer's native SOL balance.
Split-leg and multi-transaction strategies (including sandwiches) are excluded.

Usage: python3 scan_blocks.py <evidence_dir> <blocks> <span_slots>
"""
import json, os, re, sys, time, urllib.request, collections
URL = os.environ.get("SOLANA_RPC_URL", "https://api.mainnet-beta.solana.com")
E, NBLOCKS, SPAN = sys.argv[1], int(sys.argv[2]), int(sys.argv[3])
WSOL = "So11111111111111111111111111111111111111112"
STABLE = {"EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v", "Es9vMFrzaCERmJfrF4H2FYD4KCoNkY11McCe8BenwNYB"}
DEX = {"675kPX9MHTjS2zt1qfr1NYHuzeLXfQM9H24wFSUt1Mp8", "CAMMCzo5YL8w4VFF8KVHrK22GGUsp5VTaW7grrKgrWqK",
       "CPMMoo8L3F4NbTegBCKVNunggL7H1ZpdTHKxQB5qKP1C", "whirLbMiicVdio4qvUfM5KAg6Ct8VwpYzGff3uctyCc",
       "LBUZKhRxPF3XUpBCjp4YzTKgLccjZhTSDM9YuVaPwxo", "Eo7WjKq67rjJQSZxS6z3YkapzY3eMj6Xy8X5EQVn5UaB",
       "cpamdpZCGKUy5JxQXB4dcpGPiikHawvSWAd6mEn1sGG", "pAMMBay6oceH9fJKBRHGP5D4bD4sWpmSwMn52FMfXEA",
       "PhoeNiXZ8ByJGLkxNfZRnkUfjvmuYqLR89jjFHGqdXY", "JUP6LkbZbjS1jKKwapdHNy74zcZ3tLUZoi5QNyVTaV4"}
SYSTEMISH = {"ComputeBudget111111111111111111111111111111", "11111111111111111111111111111111",
             "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA", "TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb",
             "ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL", "SysvarRent111111111111111111111111111111111",
             "Sysvar1nstructions1111111111111111111111111", "SysvarC1ock11111111111111111111111111111111"}

def rpc(method, params):
    for i in range(6):
        try:
            req = urllib.request.Request(URL, json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": params}).encode(), {"content-type": "application/json"})
            r = json.load(urllib.request.urlopen(req, timeout=90))
            if "result" in r: return r["result"]
            if r.get("error", {}).get("code") in (-32007, -32009): return None  # skipped / missing slot
        except Exception: pass
        time.sleep(2 * (i + 1))
    return None

def amt(b): return int(b["uiTokenAmount"]["amount"]) / 10 ** b["uiTokenAmount"]["decimals"]

tip = rpc("getSlot", [{"commitment": "finalized"}]) - 100
slots = [tip - SPAN + round(i * SPAN / NBLOCKS) for i in range(NBLOCKS)]
json.dump({"frozenAtUtc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "slots": slots}, open(f"{E}/sample_slots.json", "w"))
out = open(f"{E}/arb_transactions.jsonl", "w"); stats = collections.Counter(); t0 = time.time()
for n, slot in enumerate(slots, 1):
    blk = rpc("getBlock", [slot, {"encoding": "json", "transactionDetails": "accounts", "maxSupportedTransactionVersion": 1, "rewards": False}])
    if not blk: stats["missing_blocks"] += 1; continue
    stats["blocks"] += 1
    for tx in blk["transactions"]:
        stats["txs"] += 1; m = tx["meta"]
        if m["err"] is not None: continue
        keys = [k["pubkey"] for k in tx["transaction"]["accountKeys"]]
        dex = DEX & set(keys)
        if not dex: continue
        payer = keys[0]
        pre = {b["accountIndex"]: b for b in m.get("preTokenBalances") or []}
        post = {b["accountIndex"]: b for b in m.get("postTokenBalances") or []}
        sol = (m["postBalances"][0] - m["preBalances"][0]) / 1e9; stable = 0.0; other = False
        for i in set(pre) | set(post):
            b = post.get(i) or pre[i]
            if b.get("owner") != payer: continue
            lam = (m["postBalances"][i] - m["preBalances"][i]) / 1e9
            d = (amt(post[i]) if i in post else 0) - (amt(pre[i]) if i in pre else 0)
            if b["mint"] == WSOL: sol += lam
            elif b["mint"] in STABLE: stable += d; sol += lam
            elif abs(d) > 0: other = True; break
            else: sol += lam
        if other or (sol <= 0 and stable <= 0): continue
        if sol < 0 and stable > 0 and stable < 1: pass
        bots = sorted(set(keys) - DEX - SYSTEMISH - {payer})
        out.write(json.dumps({"slot": slot, "blockTime": blk.get("blockTime"), "sig": tx["transaction"]["signatures"][0], "payer": payer,
                              "solDelta": sol, "stableDelta": stable, "fee": m["fee"], "cu": m.get("computeUnitsConsumed"),
                              "dex": sorted(dex), "programs": [k for k in bots if k in {kk["pubkey"] for kk in tx["transaction"]["accountKeys"] if not kk["writable"]}]}) + "\n")
        stats["arb_candidates"] += 1
    if n % 25 == 0: print(f"{n}/{len(slots)} blocks, {stats['arb_candidates']} arb candidates, {time.time()-t0:.0f}s", flush=True)
out.close(); json.dump(stats, open(f"{E}/scan_stats.json", "w")); print("DONE", dict(stats), flush=True)