#!/usr/bin/env python3
"""
Merkle Root Calculator for WitsCoin

Computes the Merkle root for a set of transactions, matching the server's algorithm.

Usage:
    python3 merkle_root.py              # Run worked example
    python3 merkle_root.py hash1 hash2  # Compute root from given hashes
"""

import hashlib
import sys
import uuid


def sha256(data: str) -> str:
    """SHA-256 hash of a string, returned as hex."""
    return hashlib.sha256(data.encode('utf-8')).hexdigest()


def hash_transaction(tx_id: str, from_addr: str, to_addr: str, amount: int, nonce: int) -> str:
    """Hash a transaction the way the server does: SHA-256(id + from + to + amount + nonce)"""
    data = f"{tx_id}{from_addr}{to_addr}{amount}{nonce}"
    return sha256(data)


def hash_pair(left: str, right: str) -> str:
    """Hash two hashes together. Sort them first so order doesn't matter."""
    if left > right:
        left, right = right, left
    return sha256(left + right)


def compute_merkle_root(hashes: list) -> str:
    """
    Build a Merkle tree from a list of hashes and return the root.

    - If empty, return 64 zeros
    - If one hash, that's the root
    - Otherwise pair up, hash pairs, repeat (duplicate last if odd)
    """
    if not hashes:
        return "0" * 64

    if len(hashes) == 1:
        return hashes[0]

    current = hashes[:]
    level = 0

    while len(current) > 1:
        level += 1
        # Duplicate last if odd
        if len(current) % 2 == 1:
            current.append(current[-1])

        next_level = []
        for i in range(0, len(current), 2):
            parent = hash_pair(current[i], current[i + 1])
            next_level.append(parent)

        current = next_level

    return current[0]


def compute_reward_tx_id(wallet_address: str, timestamp: int) -> str:
    """Compute the deterministic reward transaction ID using UUID v5 with OID namespace."""
    namespace_oid = uuid.UUID('6ba7b812-9dad-11d1-80b4-00c04fd430c8')
    return str(uuid.uuid5(namespace_oid, f"{wallet_address}{timestamp}"))


def compute_mining_merkle_root(wallet_address: str, timestamp: int, pending_txs: list) -> str:
    """
    Compute the Merkle root for a block being mined.

    Includes the reward transaction (coinbase -> miner) plus any pending transactions.
    """
    # Reward transaction
    reward_tx_id = compute_reward_tx_id(wallet_address, timestamp)
    reward_hash = hash_transaction(reward_tx_id, "coinbase", wallet_address, 50, 0)

    all_hashes = [reward_hash]

    # Add pending transactions
    for tx in pending_txs:
        tx_hash = hash_transaction(tx['id'], tx['from'], tx['to'], tx['amount'], 0)
        all_hashes.append(tx_hash)

    return compute_merkle_root(all_hashes)


def demo():
    """Worked example with 3 transactions showing every step."""
    print("MERKLE ROOT - WORKED EXAMPLE")
    print("=" * 60)
    print()

    # These are the same transactions shown on the docs page
    txs = [
        ("550e8400-e29b-41d4-a716-446655440000", "wc_a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6", "joshua_wits_coin_2026", 100, 1),
        ("550e8400-e29b-41d4-a716-446655440001", "wc_b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7", "wc_a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6", 50, 2),
        ("550e8400-e29b-41d4-a716-446655440002", "wc_c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8", "joshua_wits_coin_2026", 75, 3),
    ]

    # Step 1: hash each transaction
    print("Step 1: Hash each transaction  SHA-256(id + from + to + amount + nonce)")
    print()
    hashes = []
    for i, (tx_id, frm, to, amt, nonce) in enumerate(txs, 1):
        h = hash_transaction(tx_id, frm, to, amt, nonce)
        hashes.append(h)
        print(f"  TX {i}: {h}")
    print()

    # Step 2: pair up (3 hashes -> duplicate last)
    print("Step 2: Pair and hash (3 hashes, so duplicate last)")
    print()
    hashes_padded = hashes + [hashes[-1]]  # duplicate last

    hash_a = hash_pair(hashes_padded[0], hashes_padded[1])
    hash_b = hash_pair(hashes_padded[2], hashes_padded[3])

    smaller1 = min(hashes_padded[0], hashes_padded[1])
    larger1 = max(hashes_padded[0], hashes_padded[1])
    print(f"  Pair 1: sort then concat")
    print(f"    smaller: {smaller1}")
    print(f"    larger:  {larger1}")
    print(f"    Hash A:  {hash_a}")
    print()
    print(f"  Pair 2: tx3 duplicated (same hash twice)")
    print(f"    Hash B:  {hash_b}")
    print()

    # Step 3: final root
    print("Step 3: Hash the two remaining hashes together")
    print()
    smaller2 = min(hash_a, hash_b)
    larger2 = max(hash_a, hash_b)
    root = hash_pair(hash_a, hash_b)
    print(f"    smaller: {smaller2}")
    print(f"    larger:  {larger2}")
    print(f"    Root:    {root}")
    print()
    print("=" * 60)
    print(f"MERKLE ROOT: {root}")
    print("=" * 60)
    print()

    # Also show the mining case
    print()
    print("MINING EXAMPLE - Block with only reward transaction")
    print("=" * 60)
    print()
    wallet = "wc_ee964162ad97b399f7cbbc540148ad8e"
    timestamp = 1777581991
    print(f"  Wallet:    {wallet}")
    print(f"  Timestamp: {timestamp}")
    print()

    reward_tx_id = compute_reward_tx_id(wallet, timestamp)
    print(f"  Reward TX ID: {reward_tx_id}")
    print(f"    = uuid5(OID, \"{wallet}{timestamp}\")")
    print()

    reward_hash = hash_transaction(reward_tx_id, "coinbase", wallet, 50, 0)
    print(f"  TX Hash: {reward_hash}")
    print(f"    = SHA-256(\"{reward_tx_id}coinbase{wallet}500\")")
    print()
    print(f"  Only 1 transaction, so Merkle root = TX hash")
    print(f"  Merkle Root: {reward_hash}")


if __name__ == "__main__":
    if len(sys.argv) > 1:
        # Compute merkle root from hashes passed as arguments
        hashes = sys.argv[1:]
        root = compute_merkle_root(hashes)
        print(f"Merkle Root: {root}")
    else:
        demo()