┌─────────────────────────────────────────────────────────────────────────────┐
│                    ROLE OF HASHING IN BLOCKCHAIN                           │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │                    BLOCK IDENTIFICATION                              │   │
│  │  Each block is uniquely identified by its hash.                     │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │                    CHAIN INTEGRITY                                   │   │
│  │  The previous block's hash is stored in the next block.             │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │                    DATA INTEGRITY                                    │   │
│  │  Any change to transaction data changes the block hash.             │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │                    PROOF OF WORK                                     │   │
│  │  Miners find a nonce that produces a hash below target.             │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
│  ┌──────────────────────────────────────────────────────────────────────┐   │
│  │                    MERKLE TREES                                      │   │
│  │  Hashes of transactions are combined into a single root hash.       │   │
│  └──────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

                    ┌───────────────┐
                    │  Merkle Root  │
                    └───────────────┘
                           │
              ┌────────────┴────────────┐
              │                         │
         ┌────┴────┐               ┌────┴────┐
         │ Hash(H1,H2)│             │ Hash(H3,H4)│
         └──────────┘               └──────────┘
              │                         │
         ┌────┴────┐               ┌────┴────┐
         │         │               │         │
      ┌──┴──┐  ┌──┴──┐         ┌──┴──┐  ┌──┴──┐
      │ H1  │  │ H2  │         │ H3  │  │ H4  │
      └─────┘  └─────┘         └─────┘  └─────┘
      Tx1      Tx2              Tx3      Tx4

# ===================================================================
# MODULE 1, LESSON 2: CRYPTOGRAPHIC FOUNDATIONS OF BLOCKCHAIN
# ===================================================================

import hashlib
import binascii
import os
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec, utils
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
from typing import List, Tuple
import json
import time
import matplotlib.pyplot as plt
import numpy as np
import warnings
warnings.filterwarnings('ignore')

print("="*70)
print("CRYPTOGRAPHIC FOUNDATIONS OF BLOCKCHAIN")
print("="*70)

# ----------------------------------------------------------------
# PART A: HASH FUNCTIONS DEMONSTRATION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART A: Hash Functions Demonstration")
print("-"*60)

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

def hash_ripemd160(data: str) -> str:
    """Compute RIPEMD-160 hash of a string."""
    return hashlib.new('ripemd160', data.encode('utf-8')).hexdigest()

# Test data
test_messages = [
    "Hello, Blockchain!",
    "Hello, Blockchain!?",
    "Hello, Blockchain."  # small change
]

print("Hash Function Demonstration:")
print("-"*40)
for msg in test_messages:
    sha256_hash = hash_sha256(msg)
    ripemd_hash = hash_ripemd160(msg)
    print(f"Message: {msg}")
    print(f"  SHA-256: {sha256_hash[:16]}...")
    print(f"  RIPEMD-160: {ripemd_hash[:16]}...")
    print()

# Avalanche effect demonstration
msg1 = "Bitcoin"
msg2 = "Bitcoin"  # same
msg3 = "bitcoin"  # case change

print("Avalanche Effect (SHA-256):")
print("-"*40)
print(f"'{msg1}' -> {hash_sha256(msg1)[:20]}...")
print(f"'{msg2}' -> {hash_sha256(msg2)[:20]}... (same)")
print(f"'{msg3}' -> {hash_sha256(msg3)[:20]}... (different)")

# ----------------------------------------------------------------
# PART B: PUBLIC-KEY CRYPTOGRAPHY WITH ECC
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART B: Public-Key Cryptography with ECDSA")
print("-"*60)

from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature, encode_dss_signature

def generate_ec_key_pair():
    """Generate a private/public key pair using ECDSA (secp256k1)."""
    private_key = ec.generate_private_key(ec.SECP256K1())
    public_key = private_key.public_key()
    return private_key, public_key

def sign_message(private_key, message: str) -> bytes:
    """Sign a message with the private key."""
    signature = private_key.sign(
        message.encode('utf-8'),
        ec.ECDSA(hashes.SHA256())
    )
    return signature

def verify_signature(public_key, message: str, signature: bytes) -> bool:
    """Verify the signature with the public key."""
    try:
        public_key.verify(
            signature,
            message.encode('utf-8'),
            ec.ECDSA(hashes.SHA256())
        )
        return True
    except Exception:
        return False

# Generate key pair
private_key, public_key = generate_ec_key_pair()

# Serialize public key for display
public_key_bytes = public_key.public_bytes(
    encoding=serialization.Encoding.PEM,
    format=serialization.PublicFormat.SubjectPublicKeyInfo
)
print("Key pair generated (secp256k1).")
print(f"Public key (PEM first 30 chars): {public_key_bytes[:30].decode('utf-8')}...")

# Sign a transaction
message = "Transfer 0.5 BTC from Alice to Bob"
signature = sign_message(private_key, message)
print(f"\nMessage: {message}")
print(f"Signature (hex): {binascii.hexlify(signature).decode('utf-8')[:32]}...")

# Verify signature
valid = verify_signature(public_key, message, signature)
print(f"Signature valid? {valid}")

# Try tampering
tampered_message = "Transfer 0.5 BTC from Alice to Charlie"
valid_tampered = verify_signature(public_key, tampered_message, signature)
print(f"Tampered message signature valid? {valid_tampered}")

# ----------------------------------------------------------------
# PART C: MERKLE TREE IMPLEMENTATION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART C: Merkle Tree Implementation")
print("-"*60)

def compute_merkle_root(transactions: List[str]) -> str:
    """Compute Merkle root from a list of transaction strings."""
    # Hash each transaction
    leaves = [hash_sha256(tx) for tx in transactions]
    
    if len(leaves) == 0:
        return hash_sha256("empty")
    
    # If odd number of leaves, duplicate the last one
    if len(leaves) % 2 == 1:
        leaves.append(leaves[-1])
    
    # Build tree bottom-up
    while len(leaves) > 1:
        parent_level = []
        for i in range(0, len(leaves), 2):
            combined = leaves[i] + leaves[i+1]
            parent_hash = hash_sha256(combined)
            parent_level.append(parent_hash)
        leaves = parent_level
        
        # If odd, duplicate last
        if len(leaves) % 2 == 1 and len(leaves) > 1:
            leaves.append(leaves[-1])
    
    return leaves[0] if leaves else hash_sha256("empty")

def get_merkle_proof(transactions: List[str], tx_index: int) -> List[Tuple[str, str]]:
    """Generate a Merkle proof for a transaction at given index.
    Returns list of (sibling_hash, side) where side is 'left' or 'right'.
    """
    leaves = [hash_sha256(tx) for tx in transactions]
    if len(leaves) == 0 or tx_index >= len(leaves):
        return []
    
    proof = []
    level = leaves
    
    while len(level) > 1:
        # Find sibling index
        if tx_index % 2 == 0:
            # sibling is to the right if exists, else duplicate
            if tx_index + 1 < len(level):
                sibling = level[tx_index + 1]
                side = 'right'
            else:
                sibling = level[tx_index]  # duplicate itself (odd case)
                side = 'right'  # we treat as right
        else:
            sibling = level[tx_index - 1]
            side = 'left'
        
        proof.append((sibling, side))
        
        # Compute parent hash
        if tx_index % 2 == 0:
            if tx_index + 1 < len(level):
                combined = level[tx_index] + level[tx_index + 1]
            else:
                combined = level[tx_index] + level[tx_index]
        else:
            combined = level[tx_index - 1] + level[tx_index]
        
        tx_index = tx_index // 2
        
        # Build next level
        next_level = []
        for i in range(0, len(level), 2):
            if i+1 < len(level):
                combined = level[i] + level[i+1]
            else:
                combined = level[i] + level[i]
            next_level.append(hash_sha256(combined))
        level = next_level
    
    return proof

def verify_merkle_proof(tx_hash: str, proof: List[Tuple[str, str]], merkle_root: str) -> bool:
    """Verify a Merkle proof."""
    current_hash = tx_hash
    for sibling_hash, side in proof:
        if side == 'left':
            combined = sibling_hash + current_hash
        else:  # right
            combined = current_hash + sibling_hash
        current_hash = hash_sha256(combined)
    return current_hash == merkle_root

# Sample transactions
transactions = [
    "Alice pays Bob 1 BTC",
    "Bob pays Charlie 2 BTC",
    "Charlie pays David 0.5 BTC",
    "David pays Alice 1.5 BTC",
    "Eve pays Frank 3 BTC"  # odd number
]

print("Transactions:")
for i, tx in enumerate(transactions):
    print(f"  {i}: {tx}")

root = compute_merkle_root(transactions)
print(f"\nMerkle Root: {root[:20]}...")

# Generate proof for transaction at index 2
proof = get_merkle_proof(transactions, 2)
print(f"\nProof for transaction index 2 (size {len(proof)}):")
for i, (sibling, side) in enumerate(proof):
    print(f"  Level {i+1}: sibling hash = {sibling[:16]}..., side = {side}")

# Verify proof
tx_hash = hash_sha256(transactions[2])
is_valid = verify_merkle_proof(tx_hash, proof, root)
print(f"\nProof verification result: {is_valid}")

# Visualise Merkle tree structure
def print_merkle_tree(transactions: List[str], max_levels: int = 3):
    """Pretty print the Merkle tree."""
    leaves = [hash_sha256(tx)[:8] for tx in transactions]
    if len(leaves) % 2 == 1:
        leaves.append(leaves[-1])
    
    level = leaves
    levels = [level]
    while len(level) > 1:
        parent_level = []
        for i in range(0, len(level), 2):
            combined = level[i] + level[i+1]
            parent_hash = hash_sha256(combined)[:8]
            parent_level.append(parent_hash)
        levels.append(parent_level)
        level = parent_level
        if len(level) % 2 == 1 and len(level) > 1:
            level.append(level[-1])
    
    # Reverse for display (root first)
    levels.reverse()
    print("Merkle Tree Structure (truncated hashes):")
    for i, lvl in enumerate(levels):
        indent = "  " * i
        if i == 0:
            print(f"{indent}Root: {lvl[0]}")
        else:
            print(f"{indent}Level {i}: {' '.join(lvl)}")

print_merkle_tree(transactions)

# ----------------------------------------------------------------
# PART D: CRYPTOGRAPHIC SECURITY METRICS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: Cryptographic Security Metrics")
print("-"*60)

security_metrics = pd.DataFrame({
    'Security Aspect': [
        'Hash Collision Resistance',
        'Pre-image Resistance',
        'Private Key Space (secp256k1)',
        'Quantum Resistance',
        'Key Size (ECC)',
        'Key Size (RSA)'
    ],
    'Bitcoin (SHA-256/ECDSA)': [
        '~2^128 (collision)',
        '~2^256',
        '2^256',
        'Vulnerable (Shor)',
        '256 bits',
        'N/A'
    ],
    'Ethereum (Keccak-256/ECDSA)': [
        '~2^128',
        '~2^256',
        '2^256',
        'Vulnerable (Shor)',
        '256 bits',
        'N/A'
    ],
    'Post-Quantum Candidate': [
        '~2^256',
        '~2^256',
        '2^256 (or larger)',
        'Resistant',
        'Variable',
        '2048+ bits'
    ]
})

print("Cryptographic Security Metrics:")
print(security_metrics.to_string(index=False))

# ----------------------------------------------------------------
# PART E: ADDRESS GENERATION (Bitcoin-style)
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: Bitcoin Address Generation (Simplified)")
print("-"*60)

def generate_bitcoin_address(public_key_pem: bytes) -> str:
    """Generate a Bitcoin-style address from a public key (simplified)."""
    # Step 1: SHA-256 of public key (der format)
    sha256_hash = hashlib.sha256(public_key_pem).digest()
    # Step 2: RIPEMD-160 of the SHA-256 hash
    ripemd160 = hashlib.new('ripemd160')
    ripemd160.update(sha256_hash)
    hashed_public_key = ripemd160.digest()
    # Step 3: Add network byte (0x00 for mainnet)
    network_byte = b'\x00' + hashed_public_key
    # Step 4: Double SHA-256 checksum
    checksum = hashlib.sha256(hashlib.sha256(network_byte).digest()).digest()[:4]
    # Step 5: Concatenate and base58 encode (simplified: use hex for display)
    address_bytes = network_byte + checksum
    # Convert to hex for display (real implementation uses Base58Check)
    return binascii.hexlify(address_bytes).decode('utf-8')

# Generate address from our public key
address_hex = generate_bitcoin_address(public_key_bytes)
print(f"Generated Bitcoin address (hex): {address_hex[:20]}...")
print(f"Length: {len(address_hex)} characters")

# ----------------------------------------------------------------
# PART F: COMPARISON OF HASH FUNCTIONS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: Performance Comparison of Hash Functions")
print("-"*60)

# Benchmark hash functions
import time

data = "Blockchain" * 1000  # 9KB string
iterations = 10000

def benchmark_hash(name, hash_func, data):
    start = time.time()
    for _ in range(iterations):
        hash_func(data)
    elapsed = time.time() - start
    return elapsed

hash_functions = {
    'SHA-256': hashlib.sha256,
    'SHA-512': hashlib.sha512,
    'RIPEMD-160': lambda x: hashlib.new('ripemd160', x.encode('utf-8')).digest(),
    'BLAKE2': hashlib.blake2b,
}

print(f"Benchmarking {iterations} iterations on {len(data)} bytes of data:")
results = {}
for name, func in hash_functions.items():
    try:
        elapsed = benchmark_hash(name, lambda x: func(x.encode('utf-8')).hexdigest(), data)
        results[name] = elapsed
        print(f"{name:12s}: {elapsed:.4f} seconds")
    except Exception as e:
        print(f"{name:12s}: Error - {e}")

# Visualise performance
if results:
    fig, ax = plt.subplots(figsize=(8, 4))
    names = list(results.keys())
    times = list(results.values())
    ax.bar(names, times, color='teal', alpha=0.7)
    ax.set_ylabel('Time (seconds)')
    ax.set_title('Hash Function Performance')
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('hash_performance.png', dpi=300, bbox_inches='tight')
    plt.show()
    print("Performance chart saved as 'hash_performance.png'")

# ----------------------------------------------------------------
# PART G: SUMMARY AND RECOMMENDATIONS
# ----------------------------------------------------------------

print("\n" + "="*70)
print("PART G: Summary and Recommendations")
print("="*70)

print("""
Cryptographic Foundations of Blockchain – Key Takeaways:

1. Cryptographic hash functions provide data integrity and security.
   - SHA-256, Keccak-256, RIPEMD-160 are widely used.
2. Public-key cryptography (ECC) enables secure ownership and signatures.
   - Private key controls assets; public key acts as an address.
3. Digital signatures prove authenticity and integrity without revealing keys.
4. Merkle trees allow efficient verification of transaction inclusion.
5. Address generation involves hashing public keys and adding checksums.
6. Quantum computing poses a future threat to ECC and hashes.

Recommendations:
  - Use established cryptographic libraries (cryptography, hashlib).
  - Never hardcode private keys; use secure key management.
  - Understand the security assumptions and limitations.
  - Stay informed about quantum-resistant algorithms.
  - Validate all signatures and hashes in production systems.
""")

print("="*70)
print("END OF LESSON 2 – MODULE 1")
print("="*70)