#!/usr/bin/env python3
"""
Test script for the Icing protocol.
This script demonstrates the full protocol flow between two peers:
1. Connection establishment
2. Ping exchange
3. Key exchange (ECDH + HKDF)
4. Encrypted messaging
"""

import time
import sys
import threading
from protocol import IcingProtocol

# ANSI colors
RED = "\033[91m"
GREEN = "\033[92m"
YELLOW = "\033[93m"
BLUE = "\033[94m"
RESET = "\033[0m"

def test_manual_protocol():
    """Test the protocol with manual step-by-step progression."""
    print(f"\n{BLUE}=== Manual Protocol Test ==={RESET}")
    print("This test demonstrates manual control of the protocol flow.\n")
    
    # Create two protocol instances
    alice = IcingProtocol()
    bob = IcingProtocol()
    
    print(f"Alice listening on port: {alice.local_port}")
    print(f"Bob listening on port: {bob.local_port}")
    
    # Exchange identity keys
    print(f"\n{YELLOW}1. Exchanging identity keys...{RESET}")
    alice.set_peer_identity(bob.identity_pubkey.hex())
    bob.set_peer_identity(alice.identity_pubkey.hex())
    print("   Identity keys exchanged.")
    
    # Establish connection
    print(f"\n{YELLOW}2. Establishing connection...{RESET}")
    alice.connect_to_peer(bob.local_port)
    time.sleep(1)  # Allow connection to establish
    print("   Connection established.")
    
    # Send ping from Alice
    print(f"\n{YELLOW}3. Sending PING request...{RESET}")
    alice.send_ping_request(cipher_type=0)  # AES-256-GCM
    time.sleep(1)  # Allow ping to be received
    
    # Bob responds to ping
    if bob.inbound_messages:
        print(f"   Bob received PING, responding...")
        bob.respond_to_ping(0, answer=1)  # Accept
        time.sleep(1)
    
    # Generate ephemeral keys
    print(f"\n{YELLOW}4. Generating ephemeral keys...{RESET}")
    alice.generate_ephemeral_keys()
    bob.generate_ephemeral_keys()
    print("   Ephemeral keys generated.")
    
    # Alice sends handshake
    print(f"\n{YELLOW}5. Sending handshake...{RESET}")
    alice.send_handshake()
    time.sleep(1)
    
    # Bob processes handshake and responds
    if bob.inbound_messages:
        for i, msg in enumerate(bob.inbound_messages):
            if msg["type"] == "HANDSHAKE":
                print(f"   Bob processing handshake...")
                bob.generate_ecdhe(i)
                bob.send_handshake()
                break
    time.sleep(1)
    
    # Alice processes Bob's handshake
    if alice.inbound_messages:
        for i, msg in enumerate(alice.inbound_messages):
            if msg["type"] == "HANDSHAKE":
                print(f"   Alice processing handshake...")
                alice.generate_ecdhe(i)
                break
    
    # Derive HKDF keys
    print(f"\n{YELLOW}6. Deriving encryption keys...{RESET}")
    alice.derive_hkdf()
    bob.derive_hkdf()
    print("   HKDF keys derived.")
    
    # Send encrypted messages
    print(f"\n{YELLOW}7. Sending encrypted messages...{RESET}")
    alice.send_encrypted_message("Hello Bob! This is a secure message.")
    time.sleep(1)
    
    # Bob decrypts the message
    if bob.inbound_messages:
        for i, msg in enumerate(bob.inbound_messages):
            if msg["type"] == "ENCRYPTED_MESSAGE":
                print(f"   Bob decrypting message...")
                bob.decrypt_received_message(i)
                break
    
    # Bob sends a reply
    bob.send_encrypted_message("Hi Alice! Message received securely.")
    time.sleep(1)
    
    # Alice decrypts the reply
    if alice.inbound_messages:
        for i, msg in enumerate(alice.inbound_messages):
            if msg["type"] == "ENCRYPTED_MESSAGE":
                print(f"   Alice decrypting message...")
                alice.decrypt_received_message(i)
                break
    
    # Show final state
    print(f"\n{YELLOW}8. Final protocol state:{RESET}")
    print("\nAlice:")
    alice.show_state()
    print("\nBob:")
    bob.show_state()
    
    # Cleanup
    alice.stop()
    bob.stop()
    print(f"\n{GREEN}Manual test completed successfully!{RESET}")


def test_auto_mode_protocol():
    """Test the protocol using automatic mode."""
    print(f"\n{BLUE}=== Automatic Mode Protocol Test ==={RESET}")
    print("This test demonstrates the automatic protocol flow.\n")
    
    # Create two protocol instances
    alice = IcingProtocol()
    bob = IcingProtocol()
    
    print(f"Alice listening on port: {alice.local_port}")
    print(f"Bob listening on port: {bob.local_port}")
    
    # Exchange identity keys
    print(f"\n{YELLOW}1. Setting up peers...{RESET}")
    alice.set_peer_identity(bob.identity_pubkey.hex())
    bob.set_peer_identity(alice.identity_pubkey.hex())
    
    # Configure auto mode for Alice (initiator)
    print(f"\n{YELLOW}2. Configuring auto mode...{RESET}")
    alice.configure_auto_mode(
        active_mode=True,
        ping_auto_initiate=True,
        preferred_cipher=0,  # AES-256-GCM
        auto_message_enabled=True,
        message_interval=2.0,
        message_content="Auto-generated secure message from Alice"
    )
    
    # Configure auto mode for Bob (responder)
    bob.configure_auto_mode(
        ping_response_accept=True,
        auto_message_enabled=True,
        message_interval=2.0,
        message_content="Auto-generated secure reply from Bob"
    )
    
    # Start auto mode
    print(f"   Starting auto mode for both peers...")
    alice.start_auto_mode()
    bob.start_auto_mode()
    
    # Establish connection (this will trigger the auto protocol)
    print(f"\n{YELLOW}3. Establishing connection...{RESET}")
    alice.connect_to_peer(bob.local_port)
    
    # Let the protocol run automatically
    print(f"\n{YELLOW}4. Running automatic protocol exchange...{RESET}")
    print("   Waiting for automatic protocol completion...")
    
    # Monitor progress
    for i in range(10):
        time.sleep(2)
        print(f"\n   Progress check {i+1}/10:")
        print(f"   Alice state: {alice.auto_mode.state}")
        print(f"   Bob state: {bob.auto_mode.state}")
        
        # Check if key exchange is complete
        if alice.state.get("key_exchange_complete") and bob.state.get("key_exchange_complete"):
            print(f"\n{GREEN}   Key exchange completed!{RESET}")
            break
    
    # Queue some additional messages
    print(f"\n{YELLOW}5. Queueing additional messages...{RESET}")
    alice.queue_auto_message("Custom message 1 from Alice")
    alice.queue_auto_message("Custom message 2 from Alice")
    bob.queue_auto_message("Custom reply from Bob")
    
    # Let messages be exchanged
    time.sleep(5)
    
    # Show final state
    print(f"\n{YELLOW}6. Final protocol state:{RESET}")
    print("\nAlice:")
    alice.show_state()
    print("\nBob:")
    bob.show_state()
    
    # Stop auto mode
    alice.stop_auto_mode()
    bob.stop_auto_mode()
    
    # Cleanup
    alice.stop()
    bob.stop()
    print(f"\n{GREEN}Automatic mode test completed successfully!{RESET}")


def main():
    """Main function to run protocol tests."""
    print(f"{BLUE}{'='*60}{RESET}")
    print(f"{BLUE}    Icing Protocol Test Suite{RESET}")
    print(f"{BLUE}{'='*60}{RESET}")
    
    print("\nSelect test mode:")
    print("1. Manual protocol test (step-by-step)")
    print("2. Automatic mode test (auto protocol flow)")
    print("3. Run both tests")
    print("0. Exit")
    
    try:
        choice = input("\nEnter your choice (0-3): ").strip()
        
        if choice == "1":
            test_manual_protocol()
        elif choice == "2":
            test_auto_mode_protocol()
        elif choice == "3":
            test_manual_protocol()
            print(f"\n{YELLOW}{'='*60}{RESET}\n")
            test_auto_mode_protocol()
        elif choice == "0":
            print("Exiting...")
            return 0
        else:
            print(f"{RED}Invalid choice. Please enter 0-3.{RESET}")
            return 1
            
    except KeyboardInterrupt:
        print(f"\n\n{YELLOW}Test interrupted by user.{RESET}")
        return 0
    except Exception as e:
        print(f"\n{RED}ERROR: {e}{RESET}")
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())