#!/usr/bin/env python3
"""AgentGenesis PoW solver.

Finds `nonce` such that
    uint256(keccak256(abi.encode(minter, agentTokenId, nonce))) < 2**(256-POW_BITS)
which is what AgentGenesis.mint(agentTokenId, nonce, referrer) checks on-chain.

Usage:
  python3 pow_solver.py --minter 0xYourWallet --agent moomo            # name without .ai
  python3 pow_solver.py --minter 0x.. --agent 0x7853...6e --workers 4  # or the tokenId
  python3 pow_solver.py ... --bits 22 --json
Optional on-chain check (eth_call checkPoW on the deployed contract):
  python3 pow_solver.py ... --contract 0xGenesis --rpc https://bsc-rpc.publicnode.com

Needs: pip install pycryptodome   (or eth-hash[pycryptodome] / pysha3)
Speed: ~200k hashes/s per core in pure python; 22 bits = ~4.2M hashes on average.
The nonce is bound to YOUR wallet + agent id, so nobody can front-run/steal it.
"""
import argparse, json, multiprocessing as mp, os, sys, time, urllib.request


def _keccak_fn():
    try:
        from Crypto.Hash import keccak
        return lambda b: keccak.new(data=b, digest_bits=256).digest()
    except ImportError:
        pass
    try:
        from eth_hash.auto import keccak
        return keccak
    except ImportError:
        pass
    try:
        import sha3
        return lambda b: sha3.keccak_256(b).digest()
    except ImportError:
        sys.exit("install pycryptodome:  pip install pycryptodome")


def agent_token_id(agent: str) -> int:
    if agent.startswith("0x") and len(agent) > 20:
        return int(agent, 16)
    if agent.isdigit():
        return int(agent)
    label = agent[:-3] if agent.endswith(".ai") else agent
    return int.from_bytes(_keccak_fn()(label.encode()), "big")


def _work(prefix, target, start, step, q):
    k = _keccak_fn()
    n = start
    tried = 0
    while True:
        for _ in range(50000):
            if int.from_bytes(k(prefix + n.to_bytes(32, "big")), "big") < target:
                q.put((n, tried))
                return
            n += step
            tried += 1


def solve(minter: str, tid: int, bits: int, workers: int):
    prefix = bytes(12) + bytes.fromhex(minter[2:].lower().rjust(40, "0")) + tid.to_bytes(32, "big")
    target = 1 << (256 - bits)
    base = int.from_bytes(os.urandom(8), "big") << 64  # random start per run
    t = time.time()
    if workers <= 1:
        k = _keccak_fn(); n = base
        while int.from_bytes(k(prefix + n.to_bytes(32, "big")), "big") >= target:
            n += 1
        return n, time.time() - t, n - base + 1
    q = mp.Queue()
    ps = [mp.Process(target=_work, args=(prefix, target, base + i, workers, q), daemon=True) for i in range(workers)]
    for p in ps: p.start()
    n, tried = q.get()
    for p in ps: p.terminate()
    return n, time.time() - t, tried * workers


def onchain_check(rpc, contract, minter, tid, nonce):
    # checkPoW(address,uint256,uint256) selector
    sel = _keccak_fn()(b"checkPoW(address,uint256,uint256)")[:4].hex()
    data = "0x" + sel + minter[2:].lower().rjust(64, "0") + "%064x" % tid + "%064x" % nonce
    req = urllib.request.Request(rpc, json.dumps({"jsonrpc": "2.0", "id": 1, "method": "eth_call",
                                                  "params": [{"to": contract, "data": data}, "latest"]}).encode(),
                                 {"Content-Type": "application/json", "User-Agent": "pow"})
    r = json.load(urllib.request.urlopen(req, timeout=20))
    return int(r.get("result", "0x0"), 16) == 1


def main():
    a = argparse.ArgumentParser()
    a.add_argument("--minter", required=True, help="wallet that will send the mint tx")
    a.add_argument("--agent", required=True, help="Agent ID name (moomo / moomo.ai) or tokenId")
    a.add_argument("--bits", type=int, default=22, help="must equal contract POW_BITS")
    a.add_argument("--workers", type=int, default=os.cpu_count() or 1)
    a.add_argument("--contract"); a.add_argument("--rpc", default="https://bsc-rpc.publicnode.com")
    a.add_argument("--json", action="store_true")
    o = a.parse_args()
    tid = agent_token_id(o.agent)
    nonce, dt, hashes = solve(o.minter, tid, o.bits, o.workers)
    out = {"minter": o.minter, "agentTokenId": "0x%064x" % tid, "nonce": nonce, "seconds": round(dt, 1),
           "hashes": hashes}
    if o.contract:
        out["onchainValid"] = onchain_check(o.rpc, o.contract, o.minter, tid, nonce)
    print(json.dumps(out) if o.json else "\n".join("%s: %s" % kv for kv in out.items()))


if __name__ == "__main__":
    main()