========================

#!/usr/bin/env python3
"""
Custom Bitcoin Script: SHA256 Hash Preimage Puzzle
Anyone who knows X such that SHA256(X) = TARGET_HASH can spend this output.
"""

import hashlib
import struct

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

def hash160(data: bytes) -> bytes:
    return hashlib.new('ripemd160', sha256(data)).digest()

# ========================
# STEP 1: Define the secret
# ========================
SECRET = b"The secret to unlock this Bitcoin puzzle"
TARGET_HASH = sha256(SECRET)
print(f"Secret: {SECRET}")
print(f"Target SHA256: {TARGET_HASH.hex()}")

# ========================
# STEP 2: Build the locking script
# ========================
def build_locking_script(target_hash: bytes) -> bytes:
    """
    Build scriptPubKey:
    OP_SHA256 <target_hash> OP_EQUAL

    Only someone who knows the preimage of target_hash can spend.
    """
    assert len(target_hash) == 32, "SHA256 hash must be 32 bytes"

    script = bytearray()

    # OP_SHA256 = 0xa8
    script.append(0xa8)

    # Push 32 bytes (0x20 = decimal 32)
    script.append(0x20)
    script.extend(target_hash)

    # OP_EQUAL = 0x87
    script.append(0x87)

    return bytes(script)

locking_script = build_locking_script(TARGET_HASH)
print(f"\nLocking script ({len(locking_script)} bytes):")
print(f"  Hex: {locking_script.hex()}")
print(f"  ASM: OP_SHA256 <{TARGET_HASH.hex()[:16]}...> OP_EQUAL")

# ========================
# STEP 3: Build the unlocking script
# ========================
def build_unlocking_script(secret: bytes) -> bytes:
    """
    Build scriptSig:
    <secret_preimage>

    Just push the secret onto the stack.
    """
    script = bytearray()

    # Push the secret bytes
    secret_len = len(secret)
    if secret_len <= 75:
        script.append(secret_len)  # Direct push
    elif secret_len <= 255:
        script.append(0x4c)  # OP_PUSHDATA1
        script.append(secret_len)
    else:
        script.append(0x4d)  # OP_PUSHDATA2
        script.extend(struct.pack('<H', secret_len))

    script.extend(secret)

    return bytes(script)

unlocking_script = build_unlocking_script(SECRET)
print(f"\nUnlocking script ({len(unlocking_script)} bytes):")
print(f"  Hex: {unlocking_script.hex()}")

# ========================
# STEP 4: Simulate script execution
# ========================
def simulate_execution(unlocking: bytes, locking: bytes) -> bool:
    """
    Simulate Bitcoin Script execution.
    Returns True if script succeeds, False otherwise.
    """
    stack = []

    print("\n=== SCRIPT EXECUTION TRACE ===")

    # Parse and execute unlocking script
    print("\nPhase 1: Unlocking script")
    pos = 0
    while pos < len(unlocking):
        opcode = unlocking[pos]
        pos += 1

        if 1 <= opcode <= 75:  # Direct push
            data = unlocking[pos:pos+opcode]
            pos += opcode
            stack.append(data)
            print(f"  PUSH {len(data)} bytes: {data[:20].hex()}...")
        elif opcode == 0x4c:  # OP_PUSHDATA1
            length = unlocking[pos]
            pos += 1
            data = unlocking[pos:pos+length]
            pos += length
            stack.append(data)
            print(f"  OP_PUSHDATA1 {length} bytes")

    print(f"  Stack after Phase 1: [{len(stack[-1])} bytes on top]")

    # Parse and execute locking script
    print("\nPhase 2: Locking script")
    pos = 0
    while pos < len(locking):
        opcode = locking[pos]
        pos += 1

        if opcode == 0xa8:  # OP_SHA256
            if not stack:
                print("  ERROR: Stack underflow at OP_SHA256")
                return False
            top = stack.pop()
            result = sha256(top)
            stack.append(result)
            print(f"  OP_SHA256: {top[:10].hex()}... → {result.hex()[:16]}...")

        elif opcode == 0x20:  # Push 32 bytes
            data = locking[pos:pos+32]
            pos += 32
            stack.append(data)
            print(f"  PUSH 32 bytes: {data.hex()[:16]}...")

        elif opcode == 0x87:  # OP_EQUAL
            if len(stack) < 2:
                print("  ERROR: Stack underflow at OP_EQUAL")
                return False
            b = stack.pop()
            a = stack.pop()
            result = bytes([1]) if a == b else bytes([0])
            stack.append(result)
            print(f"  OP_EQUAL: {a.hex()[:16]}... == {b.hex()[:16]}... → {result.hex()}")

    # Final check
    if not stack:
        print("\nFINAL: Empty stack → FAIL")
        return False

    top = stack[-1]
    success = top != b'' and top != bytes([0])
    print(f"\nFINAL: Stack top = {top.hex()} → {'SUCCESS ✓' if success else 'FAIL ✗'}")
    return success

# Run the simulation
result = simulate_execution(unlocking_script, locking_script)

# ========================
# STEP 5: Test failure case
# ========================
print("\n=== TESTING WRONG SECRET ===")
wrong_unlocking = build_unlocking_script(b"wrong secret")
failure = simulate_execution(wrong_unlocking, locking_script)
print(f"Wrong secret result: {'SUCCESS' if failure else 'FAIL (expected)'}")

# ========================
# STEP 6: Wrap in P2SH for relay compatibility
# ========================
def wrap_in_p2sh(redeem_script: bytes) -> tuple:
    """
    Wrap a custom script in P2SH for standard relay.
    Returns (p2sh_scriptPubKey, p2sh_scriptSig_prefix)
    """
    script_hash = hash160(redeem_script)

    # P2SH locking script: OP_HASH160 <hash> OP_EQUAL
    p2sh_lock = bytes([0xa9, 0x14]) + script_hash + bytes([0x87])

    print(f"\nP2SH wrapping:")
    print(f"  Redeem script hash: {script_hash.hex()}")
    print(f"  P2SH scriptPubKey: {p2sh_lock.hex()}")
    print(f"  P2SH address prefix: 3... (mainnet) or 2... (testnet)")

    return p2sh_lock

p2sh_output = wrap_in_p2sh(locking_script)