#!/usr/bin/env python3
"""
Comprehensive Test Suite for PoA Blockchain Implementation

Tests the bootnode and validator services without requiring Docker or pytest.
Verifies core logic, consensus, and data structures.
"""

import sys
import json
import hashlib
import time
from pathlib import Path

# Add paths for imports
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / "bootnode"))
sys.path.insert(0, str(project_root / "validator"))

# ============================================================================
# BOOTNODE TESTS
# ============================================================================

def test_bootnode_initialization():
    """Test bootnode peer registry initialization"""
    import sys
    from pathlib import Path
    project_root = Path(__file__).parent.parent
    sys.path.insert(0, str(project_root / "bootnode"))
    from bootnode import PeerRegistry

    registry = PeerRegistry(peer_timeout_seconds=300)

    assert registry.peers == {}
    assert registry.peer_timeout == 300
    return True

def test_peer_registration():
    """Test registering a peer with bootnode"""
    from bootnode.bootnode import PeerRegistry, PeerInfo

    registry = PeerRegistry()

    peer = PeerInfo(
        node_id="validator-1",
        ip="validator-1",
        p2p_port=30303,
        rpc_port=8001,
        public_key="0x1234567890abcdef"
    )

    registry.register_peer(peer)

    assert "validator-1" in registry.peers
    assert registry.peers["validator-1"]["info"].node_id == "validator-1"
    return True

def test_peer_discovery():
    """Test discovering peers from bootnode"""
    from bootnode.bootnode import PeerRegistry, PeerInfo

    registry = PeerRegistry()

    # Register 3 peers
    for i in range(1, 4):
        peer = PeerInfo(
            node_id=f"validator-{i}",
            ip=f"validator-{i}",
            p2p_port=30300 + i,
            rpc_port=8000 + i,
            public_key=f"0x{'0' * 40}"
        )
        registry.register_peer(peer)

    # Discover peers (exclude validator-1)
    discovered = registry.get_peers_except("validator-1")

    assert len(discovered) == 2
    node_ids = [p.node_id for p in discovered]
    assert "validator-2" in node_ids
    assert "validator-3" in node_ids
    assert "validator-1" not in node_ids
    return True

def test_peer_heartbeat():
    """Test updating peer heartbeat"""
    from bootnode.bootnode import PeerRegistry, PeerInfo

    registry = PeerRegistry(peer_timeout_seconds=5)

    peer = PeerInfo(
        node_id="validator-1",
        ip="validator-1",
        p2p_port=30303,
        rpc_port=8001
    )

    registry.register_peer(peer)
    old_heartbeat = registry.peers["validator-1"]["last_heartbeat"]

    # Update heartbeat
    time.sleep(0.1)
    registry.update_heartbeat("validator-1")
    new_heartbeat = registry.peers["validator-1"]["last_heartbeat"]

    assert new_heartbeat > old_heartbeat
    return True

def test_stale_peer_cleanup():
    """Test cleanup of stale peers"""
    from bootnode.bootnode import PeerRegistry, PeerInfo

    registry = PeerRegistry(peer_timeout_seconds=1)

    # Register peer
    peer = PeerInfo(
        node_id="validator-1",
        ip="validator-1",
        p2p_port=30303,
        rpc_port=8001
    )
    registry.register_peer(peer)
    assert len(registry.get_all_peers()) == 1

    # Manually set old heartbeat time
    registry.peers["validator-1"]["last_heartbeat"] = time.time() - 2

    # Cleanup stale peers
    removed_count = registry.cleanup_stale_peers()

    assert removed_count == 1
    assert len(registry.get_all_peers()) == 0
    return True


# ============================================================================
# VALIDATOR/BLOCKCHAIN TESTS
# ============================================================================

def test_genesis_block_creation():
    """Test genesis block is created correctly"""
    from validator.validator import Blockchain

    blockchain = Blockchain()

    assert blockchain.get_chain_length() == 1
    genesis = blockchain.get_block(0)
    assert genesis.index == 0
    assert genesis.validator == "genesis"
    assert genesis.transactions == []
    return True

def test_block_hash_calculation():
    """Test block hash is calculated correctly"""
    from validator.validator import Blockchain, Block, Transaction

    blockchain = Blockchain()

    # Create a block
    tx = Transaction(
        voter_id="test-voter",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    block = Block(
        index=1,
        prev_hash=blockchain.get_latest_block().block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )

    # Calculate hash
    hash1 = Blockchain.calculate_block_hash(block)

    # Same block should produce same hash
    hash2 = Blockchain.calculate_block_hash(block)

    assert hash1 == hash2
    assert hash1.startswith("0x")
    assert len(hash1) == 66  # "0x" + 64 hex chars
    return True

def test_block_validation():
    """Test block validation logic"""
    from validator.validator import Blockchain, Block, Transaction

    blockchain = Blockchain()

    tx = Transaction(
        voter_id="test-voter",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    # Create valid block
    prev_block = blockchain.get_latest_block()
    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)

    assert blockchain.validate_block(block) == True

    # Test invalid block (wrong index)
    invalid_block = Block(
        index=99,  # Wrong index
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    invalid_block.block_hash = Blockchain.calculate_block_hash(invalid_block)

    assert blockchain.validate_block(invalid_block) == False

    # Test invalid block (wrong prev_hash)
    invalid_block2 = Block(
        index=1,
        prev_hash="0x" + "0" * 64,  # Wrong prev_hash
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    invalid_block2.block_hash = Blockchain.calculate_block_hash(invalid_block2)

    assert blockchain.validate_block(invalid_block2) == False

    # Test invalid block (unauthorized validator)
    invalid_block3 = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="unauthorized-node"  # Not in authorized list
    )
    invalid_block3.block_hash = Blockchain.calculate_block_hash(invalid_block3)

    assert blockchain.validate_block(invalid_block3) == False
    return True

def test_add_block_to_chain():
    """Test adding valid blocks to blockchain"""
    from validator.validator import Blockchain, Block, Transaction

    blockchain = Blockchain()

    tx = Transaction(
        voter_id="test-voter",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    prev_block = blockchain.get_latest_block()
    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)

    result = blockchain.add_block(block)

    assert result == True
    assert blockchain.get_chain_length() == 2
    assert blockchain.get_block(1).index == 1
    return True

def test_blockchain_integrity_verification():
    """Test blockchain integrity verification"""
    from validator.validator import Blockchain, Block, Transaction

    blockchain = Blockchain()

    # Add valid blocks
    for i in range(1, 4):
        prev_block = blockchain.get_latest_block()
        tx = Transaction(
            voter_id=f"voter-{i}",
            election_id=1,
            encrypted_vote=f"0x{i:04x}",
            ballot_hash=f"0x{i:04x}",
            timestamp=1699360000 + i
        )

        block = Block(
            index=i,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010 + i,
            transactions=[tx],
            validator=f"validator-{(i % 3) + 1}"
        )
        block.block_hash = Blockchain.calculate_block_hash(block)
        blockchain.add_block(block)

    # Verify integrity
    assert blockchain.verify_integrity() == True
    return True

def test_chain_immutability():
    """Test that modifying past blocks breaks chain"""
    from validator.validator import Blockchain, Block, Transaction

    blockchain = Blockchain()

    # Add a block
    prev_block = blockchain.get_latest_block()
    tx = Transaction(
        voter_id="voter-1",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)
    blockchain.add_block(block)

    # Verify chain is valid
    assert blockchain.verify_integrity() == True

    # Try to modify the transaction in block 1
    blockchain.chain[1].transactions[0].encrypted_vote = "0x9999"

    # Chain should now be invalid (hash mismatch)
    assert blockchain.verify_integrity() == False
    return True


# ============================================================================
# POA CONSENSUS TESTS
# ============================================================================

def test_round_robin_block_creation():
    """Test round-robin block creation eligibility"""
    from validator.validator import PoAValidator

    # Create 3 validators
    validators = [
        PoAValidator("validator-1", "0x1234", "http://bootnode:8546", 8001, 30303),
        PoAValidator("validator-2", "0x5678", "http://bootnode:8546", 8002, 30304),
        PoAValidator("validator-3", "0xabcd", "http://bootnode:8546", 8003, 30305),
    ]

    # Test round-robin eligibility
    validator1 = validators[0]
    assert validator1.should_create_block() == True  # Block 1

    validator2 = validators[1]
    assert validator2.should_create_block() == False  # Not eligible for block 1

    validator3 = validators[2]
    assert validator3.should_create_block() == False  # Not eligible for block 1

    return True

def test_authorized_validators():
    """Test that only authorized validators can create blocks"""
    from validator.validator import Blockchain, Block, Transaction

    blockchain = Blockchain()

    # Test authorized validators
    authorized = Blockchain.AUTHORIZED_VALIDATORS
    assert "validator-1" in authorized
    assert "validator-2" in authorized
    assert "validator-3" in authorized

    # Create block with authorized validator
    tx = Transaction(
        voter_id="voter-1",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    prev_block = blockchain.get_latest_block()
    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)

    assert blockchain.validate_block(block) == True

    # Create block with unauthorized validator
    block.validator = "unauthorized-node"
    block.block_hash = Blockchain.calculate_block_hash(block)

    assert blockchain.validate_block(block) == False
    return True


# ============================================================================
# DATA STRUCTURE TESTS
# ============================================================================

def test_transaction_model():
    """Test Transaction data model"""
    from validator.validator import Transaction

    tx = Transaction(
        voter_id="voter-123",
        election_id=1,
        encrypted_vote="0x1234567890abcdef",
        ballot_hash="0xabcdef1234567890",
        proof="0x...",
        timestamp=1699360000
    )

    assert tx.voter_id == "voter-123"
    assert tx.election_id == 1
    assert tx.encrypted_vote == "0x1234567890abcdef"
    return True

def test_block_serialization():
    """Test block to/from dictionary conversion"""
    from validator.validator import Block, Transaction

    tx = Transaction(
        voter_id="voter-1",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    block = Block(
        index=1,
        prev_hash="0x" + "0" * 64,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1",
        block_hash="0x" + "1" * 64,
        signature="0x" + "2" * 64
    )

    # Serialize to dict
    block_dict = block.to_dict()

    assert block_dict["index"] == 1
    assert block_dict["validator"] == "validator-1"
    assert len(block_dict["transactions"]) == 1
    assert block_dict["transactions"][0]["voter_id"] == "voter-1"

    # Deserialize from dict
    restored_block = Block.from_dict(block_dict)

    assert restored_block.index == block.index
    assert restored_block.validator == block.validator
    assert len(restored_block.transactions) == 1
    return True


# ============================================================================
# INTEGRATION TESTS
# ============================================================================

def test_multi_validator_consensus():
    """Test that multiple validators can reach consensus"""
    from validator.validator import Blockchain, Block, Transaction

    # Create 3 blockchains (one per validator)
    blockchains = [Blockchain(), Blockchain(), Blockchain()]

    # Simulate block creation and consensus
    for block_num in range(1, 4):
        # Determine which validator creates this block (round-robin)
        creator_idx = (block_num - 1) % 3
        validator_name = f"validator-{creator_idx + 1}"

        # Create block
        prev_block = blockchains[creator_idx].get_latest_block()
        tx = Transaction(
            voter_id=f"voter-{block_num}",
            election_id=1,
            encrypted_vote=f"0x{block_num:04x}",
            ballot_hash=f"0x{block_num:04x}",
            timestamp=1699360000 + block_num
        )

        block = Block(
            index=block_num,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010 + block_num,
            transactions=[tx],
            validator=validator_name
        )
        block.block_hash = Blockchain.calculate_block_hash(block)

        # Add to creator's chain
        blockchains[creator_idx].add_block(block)

        # Broadcast to other validators
        for i in range(3):
            if i != creator_idx:
                blockchains[i].add_block(block)

    # Verify all validators have same blockchain
    chain_hashes = []
    for i, bc in enumerate(blockchains):
        last_block_hash = bc.get_latest_block().block_hash
        chain_hashes.append(last_block_hash)
        assert bc.verify_integrity() == True

    # All should have same last block
    assert chain_hashes[0] == chain_hashes[1] == chain_hashes[2]
    return True

def test_vote_immutability_across_validators():
    """Test that votes are immutable once recorded on blockchain"""
    from validator.validator import Blockchain, Block, Transaction

    blockchain = Blockchain()

    # Record a vote
    original_vote = Transaction(
        voter_id="voter-immutable",
        election_id=1,
        encrypted_vote="0x1234567890abcdef",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    prev_block = blockchain.get_latest_block()
    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[original_vote],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)
    blockchain.add_block(block)

    # Verify vote is recorded
    assert blockchain.chain[1].transactions[0].encrypted_vote == "0x1234567890abcdef"
    assert blockchain.verify_integrity() == True

    # Try to modify the vote (simulating attack)
    blockchain.chain[1].transactions[0].encrypted_vote = "0xhackedvalue"

    # Blockchain integrity should now fail
    assert blockchain.verify_integrity() == False
    return True


# ============================================================================
# JSON-RPC TESTS
# ============================================================================

def test_json_rpc_structure():
    """Test JSON-RPC request/response format"""
    # Valid JSON-RPC request
    request = {
        "jsonrpc": "2.0",
        "method": "eth_sendTransaction",
        "params": [{"data": "0x1234"}],
        "id": 1
    }

    # Verify structure
    assert request["jsonrpc"] == "2.0"
    assert request["method"] in ["eth_sendTransaction", "eth_getTransactionReceipt", "eth_blockNumber", "eth_getBlockByNumber"]
    assert "id" in request

    # Valid JSON-RPC response
    response = {
        "jsonrpc": "2.0",
        "result": "0xabc123",
        "id": 1
    }

    assert response["jsonrpc"] == "2.0"
    assert "result" in response or "error" in response
    assert response["id"] == 1
    return True


# ============================================================================
# RUN ALL TESTS
# ============================================================================

def run_all_tests():
    """Run all tests and report results"""
    tests = [
        # Bootnode tests
        ("Bootnode initialization", test_bootnode_initialization),
        ("Peer registration", test_peer_registration),
        ("Peer discovery", test_peer_discovery),
        ("Peer heartbeat", test_peer_heartbeat),
        ("Stale peer cleanup", test_stale_peer_cleanup),

        # Blockchain tests
        ("Genesis block creation", test_genesis_block_creation),
        ("Block hash calculation", test_block_hash_calculation),
        ("Block validation", test_block_validation),
        ("Add block to chain", test_add_block_to_chain),
        ("Blockchain integrity verification", test_blockchain_integrity_verification),
        ("Chain immutability", test_chain_immutability),

        # PoA Consensus tests
        ("Round-robin block creation", test_round_robin_block_creation),
        ("Authorized validators", test_authorized_validators),

        # Data structure tests
        ("Transaction model", test_transaction_model),
        ("Block serialization", test_block_serialization),

        # Integration tests
        ("Multi-validator consensus", test_multi_validator_consensus),
        ("Vote immutability across validators", test_vote_immutability_across_validators),

        # JSON-RPC tests
        ("JSON-RPC structure", test_json_rpc_structure),
    ]

    print("\n" + "=" * 80)
    print("PoA BLOCKCHAIN IMPLEMENTATION TEST SUITE")
    print("=" * 80 + "\n")

    passed = 0
    failed = 0
    errors = []

    for test_name, test_func in tests:
        try:
            result = test_func()
            if result:
                print(f"✅ {test_name}")
                passed += 1
            else:
                print(f"❌ {test_name}")
                failed += 1
                errors.append((test_name, "Assertion failed"))
        except Exception as e:
            print(f"❌ {test_name}")
            failed += 1
            errors.append((test_name, str(e)))

    # Print summary
    print("\n" + "=" * 80)
    print("TEST SUMMARY")
    print("=" * 80 + "\n")

    print(f"✅ Passed: {passed}/{len(tests)}")
    print(f"❌ Failed: {failed}/{len(tests)}")

    if errors:
        print("\nErrors:")
        for test_name, error in errors:
            print(f"  - {test_name}: {error}")

    if failed == 0:
        print("\n✅ ALL TESTS PASSED!")
        print("\nTest Categories:")
        print("  ✅ Bootnode: 5/5 tests passed")
        print("  ✅ Blockchain: 6/6 tests passed")
        print("  ✅ PoA Consensus: 2/2 tests passed")
        print("  ✅ Data Structures: 2/2 tests passed")
        print("  ✅ Integration: 2/2 tests passed")
        print("  ✅ JSON-RPC: 1/1 tests passed")
        print("  ───────────────────────")
        print("  ✅ TOTAL: 18/18 tests passed\n")
        return True
    else:
        print(f"\n❌ {failed} TESTS FAILED\n")
        return False


if __name__ == "__main__":
    success = run_all_tests()
    sys.exit(0 if success else 1)