#!/usr/bin/env python3
"""
derive.py - reproduce the Quantum Bounty vault addresses from their public seeds.

Pure Python 3.8+, standard library only. Implements the pieces of secp256k1,
ed25519, BIP-340 tagged hashes, bech32m (BIP-350) and base58 that are needed
to show that both vault addresses are "nothing-up-my-sleeve" points: public
keys derived from a hash, for which no private key is known to anyone.

Usage:
    python3 derive.py            # human readable report
    python3 derive.py --json     # machine readable
    python3 derive.py --check    # exit 1 if the published values do not match

Bitcoin (BIP-341 style):
    x  = SHA256(seed || counter_be32), counter incremented until lift_x(x) is
         a valid curve point P (BIP-340 lift_x, even y)
    t  = tagged_hash("TapTweak", P.x)          (no script tree)
    Q  = P + t*G
    address = bech32m("bc", witness v1, Q.x)

Solana (ed25519):
    b  = SHA256(seed || counter_be32), counter incremented until b decodes as a
         valid ed25519 point (RFC 8032 decoding)
    P' = 8 * P                                  (clear the cofactor)
    check: L * P' == identity                   (P' is in the prime order group)
    address = base58(encode(P'))
"""

import argparse
import hashlib
import json
import sys

# --------------------------------------------------------------------------
# Published values (what the website shows). derive.py recomputes them.
# --------------------------------------------------------------------------

BTC_SEED = "QUANTUMBOUNTY/BTC/v1/000000000019d6689c085ae165831e934ff763ae46a2a6c172b3f1b60a8ce26f"
SOL_SEED = "QUANTUMBOUNTY/SOL/v1/5eykt4UsFv8P8NJdTREpY1vzqKqZKvdpKuc147dw2N9d"

EXPECTED = {
    "btc": {
        "counter": 0,
        "internal_key_x": "89cab8a147d54575a8fe4b9248435fe71b9811ea513e8982895e621d5f362b15",
        "tap_tweak": "d234ceabc8758b652f3586af800b48fd36b136434aa9e80e27a813b54a3a4db4",
        "output_key_x": "bc5b325601cc529103187ee1511110038190c1c4f946454334ff6f337efbd4f1",
        "address": "bc1ph3dny4spe3ffzqcc0ms4zygsqwqepswyl9ry2se5lahnxlhm6ncs2kzylk",
    },
    "sol": {
        "counter": 2,
        "public_key": "ae689445fb688e39de4a599ad48a568f04a7313b8237edb43b35d7baf8cc1f12",
        "address": "CjpRvmzbWBF497hSGL3cBdxq6q1jSnZBU6LedpRhMU4y",
    },
}


def sha256(data: bytes) -> bytes:
    return hashlib.sha256(data).digest()


def seed_hash(seed: str, counter: int) -> bytes:
    """SHA256(seed || 4-byte big-endian counter)."""
    return sha256(seed.encode("utf-8") + counter.to_bytes(4, "big"))


# --------------------------------------------------------------------------
# secp256k1 (affine arithmetic, None is the point at infinity)
# --------------------------------------------------------------------------

SECP_P = 2 ** 256 - 2 ** 32 - 977
SECP_N = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
SECP_G = (
    0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798,
    0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8,
)


def secp_add(p1, p2):
    if p1 is None:
        return p2
    if p2 is None:
        return p1
    x1, y1 = p1
    x2, y2 = p2
    if x1 == x2:
        if (y1 + y2) % SECP_P == 0:
            return None
        lam = (3 * x1 * x1) * pow(2 * y1, SECP_P - 2, SECP_P) % SECP_P
    else:
        lam = (y2 - y1) * pow(x2 - x1, SECP_P - 2, SECP_P) % SECP_P
    x3 = (lam * lam - x1 - x2) % SECP_P
    y3 = (lam * (x1 - x3) - y1) % SECP_P
    return (x3, y3)


def secp_mul(k: int, point):
    result = None
    addend = point
    while k:
        if k & 1:
            result = secp_add(result, addend)
        addend = secp_add(addend, addend)
        k >>= 1
    return result


def lift_x(x: int):
    """BIP-340 lift_x: the curve point with this x coordinate and even y, or None."""
    if x >= SECP_P:
        return None
    c = (pow(x, 3, SECP_P) + 7) % SECP_P
    y = pow(c, (SECP_P + 1) // 4, SECP_P)
    if pow(y, 2, SECP_P) != c:
        return None
    if y % 2 != 0:
        y = SECP_P - y
    return (x, y)


def tagged_hash(tag: str, msg: bytes) -> bytes:
    """BIP-340 tagged hash: SHA256(SHA256(tag) || SHA256(tag) || msg)."""
    tag_hash = sha256(tag.encode("utf-8"))
    return sha256(tag_hash + tag_hash + msg)


# --------------------------------------------------------------------------
# bech32 / bech32m (BIP-173 / BIP-350)
# --------------------------------------------------------------------------

BECH32_CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
BECH32M_CONST = 0x2BC830A3


def bech32_polymod(values):
    gen = [0x3B6A57B2, 0x26508E6D, 0x1EA119FA, 0x3D4233DD, 0x2A1462B3]
    chk = 1
    for v in values:
        top = chk >> 25
        chk = ((chk & 0x1FFFFFF) << 5) ^ v
        for i in range(5):
            if (top >> i) & 1:
                chk ^= gen[i]
    return chk


def bech32_hrp_expand(hrp):
    return [ord(c) >> 5 for c in hrp] + [0] + [ord(c) & 31 for c in hrp]


def convertbits(data, frombits, tobits, pad=True):
    acc = 0
    bits = 0
    ret = []
    maxv = (1 << tobits) - 1
    for value in data:
        acc = (acc << frombits) | value
        bits += frombits
        while bits >= tobits:
            bits -= tobits
            ret.append((acc >> bits) & maxv)
    if pad and bits:
        ret.append((acc << (tobits - bits)) & maxv)
    return ret


def bech32m_encode(hrp: str, witver: int, witprog: bytes) -> str:
    data = [witver] + convertbits(witprog, 8, 5)
    values = bech32_hrp_expand(hrp) + data
    polymod = bech32_polymod(values + [0, 0, 0, 0, 0, 0]) ^ BECH32M_CONST
    checksum = [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]
    return hrp + "1" + "".join(BECH32_CHARSET[d] for d in data + checksum)


# --------------------------------------------------------------------------
# ed25519 (affine twisted Edwards arithmetic, RFC 8032 encoding)
# --------------------------------------------------------------------------

ED_Q = 2 ** 255 - 19
ED_L = 2 ** 252 + 27742317777372353535851937790883648493
ED_D = (-121665 * pow(121666, ED_Q - 2, ED_Q)) % ED_Q
ED_I = pow(2, (ED_Q - 1) // 4, ED_Q)  # sqrt(-1)
ED_IDENTITY = (0, 1)


def ed_add(p1, p2):
    x1, y1 = p1
    x2, y2 = p2
    dxy = ED_D * x1 * x2 * y1 * y2 % ED_Q
    x3 = (x1 * y2 + x2 * y1) * pow(1 + dxy, ED_Q - 2, ED_Q) % ED_Q
    y3 = (y1 * y2 + x1 * x2) * pow(1 - dxy, ED_Q - 2, ED_Q) % ED_Q
    return (x3, y3)


def ed_mul(k: int, point):
    result = ED_IDENTITY
    addend = point
    while k:
        if k & 1:
            result = ed_add(result, addend)
        addend = ed_add(addend, addend)
        k >>= 1
    return result


def ed_decode(b: bytes):
    """RFC 8032 section 5.1.3 point decoding. Returns (x, y) or None."""
    if len(b) != 32:
        return None
    y = int.from_bytes(b, "little")
    sign = y >> 255
    y &= (1 << 255) - 1
    if y >= ED_Q:
        return None
    y2 = y * y % ED_Q
    u = (y2 - 1) % ED_Q
    v = (ED_D * y2 + 1) % ED_Q
    # candidate root x = (u/v)^((q+3)/8) computed as u*v^3*(u*v^7)^((q-5)/8)
    x = u * pow(v, 3, ED_Q) * pow(u * pow(v, 7, ED_Q), (ED_Q - 5) // 8, ED_Q) % ED_Q
    if (v * x * x - u) % ED_Q == 0:
        pass
    elif (v * x * x + u) % ED_Q == 0:
        x = x * ED_I % ED_Q
    else:
        return None
    if x == 0 and sign == 1:
        return None
    if (x & 1) != sign:
        x = ED_Q - x
    return (x, y)


def ed_encode(point) -> bytes:
    x, y = point
    return (y | ((x & 1) << 255)).to_bytes(32, "little")


# --------------------------------------------------------------------------
# base58 (Bitcoin / Solana alphabet)
# --------------------------------------------------------------------------

B58_ALPHABET = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"


def base58_encode(b: bytes) -> str:
    n = int.from_bytes(b, "big")
    out = ""
    while n > 0:
        n, rem = divmod(n, 58)
        out = B58_ALPHABET[rem] + out
    pad = 0
    for byte in b:
        if byte == 0:
            pad += 1
        else:
            break
    return "1" * pad + out


# --------------------------------------------------------------------------
# Derivations
# --------------------------------------------------------------------------

def derive_btc(seed: str = BTC_SEED, start_counter: int = 0, max_tries: int = 1000) -> dict:
    counter = start_counter
    while counter < start_counter + max_tries:
        h = seed_hash(seed, counter)
        point = lift_x(int.from_bytes(h, "big"))
        if point is not None:
            break
        counter += 1
    else:
        raise RuntimeError("no valid secp256k1 x coordinate found")

    px = point[0].to_bytes(32, "big")
    t_bytes = tagged_hash("TapTweak", px)
    t = int.from_bytes(t_bytes, "big")
    if t >= SECP_N:
        raise RuntimeError("tweak out of range (astronomically unlikely)")
    q = secp_add(point, secp_mul(t, SECP_G))
    qx = q[0].to_bytes(32, "big")
    return {
        "chain": "bitcoin",
        "seed": seed,
        "counter": counter,
        "internal_key_x": px.hex(),
        "tap_tweak": t_bytes.hex(),
        "output_key_x": qx.hex(),
        "address": bech32m_encode("bc", 1, qx),
    }


def derive_sol(seed: str = SOL_SEED, start_counter: int = 0, max_tries: int = 1000) -> dict:
    counter = start_counter
    while counter < start_counter + max_tries:
        h = seed_hash(seed, counter)
        point = ed_decode(h)
        if point is not None:
            cleared = ed_mul(8, point)
            if cleared != ED_IDENTITY:
                break
        counter += 1
    else:
        raise RuntimeError("no valid ed25519 point found")

    if ed_mul(ED_L, cleared) != ED_IDENTITY:
        raise RuntimeError("cleared point is not in the prime order subgroup")
    pub = ed_encode(cleared)
    return {
        "chain": "solana",
        "seed": seed,
        "counter": counter,
        "raw_hash": h.hex(),
        "public_key": pub.hex(),
        "address": base58_encode(pub),
    }


def check(result: dict) -> list:
    """Return a list of mismatches between a derivation and the published values."""
    expected = EXPECTED[{"bitcoin": "btc", "solana": "sol"}[result["chain"]]]
    return [k for k, v in expected.items() if result.get(k) != v]


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="Reproduce the Quantum Bounty vault addresses.")
    parser.add_argument("--json", action="store_true", help="print machine readable output")
    parser.add_argument("--check", action="store_true", help="exit 1 if results differ from the published values")
    parser.add_argument("--btc-seed", default=BTC_SEED, help="override the Bitcoin seed string")
    parser.add_argument("--sol-seed", default=SOL_SEED, help="override the Solana seed string")
    args = parser.parse_args(argv)

    btc = derive_btc(args.btc_seed)
    sol = derive_sol(args.sol_seed)
    using_defaults = args.btc_seed == BTC_SEED and args.sol_seed == SOL_SEED
    mismatches = {"btc": check(btc), "sol": check(sol)} if using_defaults else None

    if args.json:
        out = {"bitcoin": btc, "solana": sol}
        if mismatches is not None:
            out["matches_published"] = not (mismatches["btc"] or mismatches["sol"])
            out["mismatches"] = mismatches
        print(json.dumps(out, indent=2))
    else:
        print("QUANTUM BOUNTY - vault derivation")
        print()
        print("Bitcoin (Taproot, BIP-341 NUMS point)")
        print("  seed           :", btc["seed"])
        print("  counter        :", btc["counter"])
        print("  internal key x :", btc["internal_key_x"])
        print("  TapTweak t     :", btc["tap_tweak"])
        print("  output key x   :", btc["output_key_x"])
        print("  address        :", btc["address"])
        print()
        print("Solana (ed25519, cofactor cleared)")
        print("  seed           :", sol["seed"])
        print("  counter        :", sol["counter"])
        print("  public key     :", sol["public_key"])
        print("  address        :", sol["address"])
        print()
        if mismatches is not None:
            ok = not (mismatches["btc"] or mismatches["sol"])
            print("published values :", "MATCH" if ok else "MISMATCH " + json.dumps(mismatches))

    if args.check and mismatches is not None and (mismatches["btc"] or mismatches["sol"]):
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())
