import os
import json
import argparse
import sys
from eth_hash.auto import keccak

def parse_arguments():
    parser = argparse.ArgumentParser(description="Dynamic Merkle Tree for Ghost Preclaims")
    parser.add_argument(
        "--network-id",
        type=int,
        required=True,
        help="EVM numeric Network ID to mix into preimage (e.g., 11155111)"
    )
    return parser.parse_args()

def keccak256(data: bytes) -> bytes:
    return keccak(data)

EMPTY_HASH = keccak256(b'')

def next_power_of_two(n: int) -> int:
    if n <= 1:
        return 1
    return 1 << (n - 1).bit_length()

def generate_preimage(token_id, evm_address, shares, network_id) -> bytes:
    token_bytes = token_id.to_bytes(32, byteorder='big')
    shares_bytes = shares.to_bytes(32, byteorder='big')

    clean_addr = evm_address.replace('0x', '')
    raw_addr_bytes = bytes.fromhex(clean_addr)
    addr_bytes = b'\x00' * 12 + raw_addr_bytes

    network_bytes = network_id.to_bytes(32, byteorder='big')

    return token_bytes + shares_bytes + addr_bytes + network_bytes

def generate_tree(raw_values, network_id):
    max_index = max(item['index'] for item in raw_values)
    num_of_leaves = max_index + 1
    padded_leaves = next_power_of_two(num_of_leaves)
    total_capacity = (2 * padded_leaves) - 1

    merkle_tree = [EMPTY_HASH] * total_capacity

    for item in raw_values:
        idx = item['index']
        preimage = generate_preimage(item['token_id'], item['address'], item['shares'], network_id)
        merkle_tree[idx] = keccak256(preimage)

    layer_start = 0
    current_layer_len = padded_leaves
    write_ptr = padded_leaves

    while current_layer_len > 1:
        for i in range(0, current_layer_len, 2):
            layer_index = layer_start + i
            left_bytes = merkle_tree[layer_index]
            right_bytes = merkle_tree[layer_index | 1]

            combined = left_bytes + right_bytes
            merkle_tree[write_ptr] = keccak256(combined)
            write_ptr += 1

        layer_start += current_layer_len
        current_layer_len >>= 1

    return merkle_tree, padded_leaves

def get_merkle_proof(tree, index, padded_leaves):
    proof = []
    current_idx = index
    layer_start = 0
    current_layer_len = padded_leaves

    while current_layer_len > 1:
        sibling_idx = current_idx ^ 1
        proof.append(tree[sibling_idx].hex())

        local_pair_idx = (current_idx - layer_start) // 2
        next_layer_start = layer_start + current_layer_len
        current_idx = next_layer_start + local_pair_idx

        layer_start = next_layer_start
        current_layer_len >>= 1

    return proof

def main():
    args = parse_arguments()
    network_id = args.network_id

    snapshot_file = os.path.join("snapshots", f"{network_id}.json")

    if not os.path.exists(snapshot_file):
        print(f"[-] File not found: {snapshot_file}")
        sys.exit(1)

    print(f"[+] Reading data snapshot: {snapshot_file}")
    with open(snapshot_file, "r") as f:
        snapshot_data = json.load(f)

    if not snapshot_data:
        print("[-] Snapshot file is empty.")
        sys.exit(1)

    total_shares = 0
    for data in snapshot_data:
        total_shares += data["shares"]

    tree, padded_leaves = generate_tree(snapshot_data, args.network_id)
    root_hash = tree[-1].hex()

    proof_len = 0
    output_packages = {}
    for item in snapshot_data:
        proof = get_merkle_proof(tree, item['index'], padded_leaves)
        proof_len = len(proof)

        output_packages[item['address']] = {
            "index": item["index"],
            "token_id": item['token_id'],
            "shares": item['shares'],
            "merkle_proof": [f"0x{p}" for p in proof]
        }

    message = f"=== FINAL MERKLE ROOT [{network_id}]: {root_hash} ==="
    separator = "=" * len(message)

    print(f"\n{separator}")
    print(message)
    print(f"{separator}\n")

    print(f"[+] Total padded leaves: {padded_leaves}")
    print(f"[+] Merkle tree size: {len(tree)}")
    print(f"[+] Proof size: {proof_len}")
    print(f"[+] Total shares: {total_shares}")

    output_file = os.path.join("preclaims", f"{network_id}.json")
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    with open(output_file, "w") as f:
        json.dump(output_packages, f, indent=2)

    print(f"\n[+] Proof are stored into: {output_file}")

if __name__ == "__main__":
    main()