#!/usr/bin/env python3
"""
Test script for DryBox integration with Icing protocol
Tests encrypted voice communication with 4FSK modulation
"""

import time
import subprocess
import sys
import os
import threading
from pathlib import Path

# Add DryBox to path
sys.path.append(str(Path(__file__).parent / "DryBox"))

from integrated_protocol import IntegratedDryBoxProtocol

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


def start_gsm_simulator():
    """Start GSM simulator in background"""
    print(f"{BLUE}[TEST]{RESET} Starting GSM simulator...")
    
    gsm_path = Path(__file__).parent / "DryBox" / "gsm_simulator.py"
    process = subprocess.Popen(
        [sys.executable, str(gsm_path)],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )
    
    time.sleep(2)  # Give it time to start
    print(f"{GREEN}[TEST]{RESET} GSM simulator started")
    return process


def test_key_exchange():
    """Test key exchange between two DryBox instances"""
    print(f"\n{YELLOW}=== Testing Key Exchange ==={RESET}")
    
    # Create sender and receiver
    sender = IntegratedDryBoxProtocol(mode="sender")
    receiver = IntegratedDryBoxProtocol(mode="receiver")
    
    # Connect to GSM
    if not sender.connect_gsm():
        print(f"{RED}[ERROR]{RESET} Sender failed to connect to GSM")
        return False
        
    if not receiver.connect_gsm():
        print(f"{RED}[ERROR]{RESET} Receiver failed to connect to GSM")
        return False
    
    # Exchange identities
    sender_identity = sender.get_identity_key()
    receiver_identity = receiver.get_identity_key()
    
    print(f"{BLUE}[SENDER]{RESET} Identity: {sender_identity[:32]}...")
    print(f"{BLUE}[RECEIVER]{RESET} Identity: {receiver_identity[:32]}...")
    
    # Setup connections
    receiver_port = receiver.setup_protocol_connection(peer_identity=sender_identity)
    print(f"{BLUE}[RECEIVER]{RESET} Listening on port {receiver_port}")
    
    sender.setup_protocol_connection(peer_port=receiver_port, peer_identity=receiver_identity)
    print(f"{BLUE}[SENDER]{RESET} Connected to receiver")
    
    # Initiate key exchange with ChaCha20
    success = sender.initiate_key_exchange(cipher_type=1)
    
    if success:
        print(f"{GREEN}[SUCCESS]{RESET} Key exchange completed!")
        print(f"  Cipher: {'ChaCha20-Poly1305' if sender.protocol.cipher_type == 1 else 'AES-256-GCM'}")
        print(f"  Sender HKDF Key: {sender.protocol.hkdf_key[:16]}...")
        
        # Wait for receiver to complete
        time.sleep(2)
        if receiver.protocol.hkdf_key:
            print(f"  Receiver HKDF Key: {receiver.protocol.hkdf_key[:16]}...")
            
            # Keys should match
            if sender.protocol.hkdf_key == receiver.protocol.hkdf_key:
                print(f"{GREEN}[PASS]{RESET} Keys match!")
            else:
                print(f"{RED}[FAIL]{RESET} Keys don't match!")
                return False
    else:
        print(f"{RED}[FAIL]{RESET} Key exchange failed")
        return False
    
    # Test encrypted message
    print(f"\n{YELLOW}=== Testing Encrypted Messages ==={RESET}")
    sender.send_encrypted_message("Hello from sender!")
    time.sleep(1)
    
    # Check receiver got it
    if receiver.protocol.inbound_messages:
        last_msg = receiver.protocol.inbound_messages[-1]
        if last_msg["type"] == "ENCRYPTED_MESSAGE":
            decrypted = receiver.protocol.decrypt_received_message(len(receiver.protocol.inbound_messages) - 1)
            if decrypted:
                print(f"{GREEN}[PASS]{RESET} Message decrypted: {decrypted}")
            else:
                print(f"{RED}[FAIL]{RESET} Failed to decrypt message")
    
    # Clean up
    sender.close()
    receiver.close()
    
    return True


def test_voice_transmission():
    """Test voice transmission with encryption and FSK"""
    print(f"\n{YELLOW}=== Testing Voice Transmission ==={RESET}")
    
    # Check if input.wav exists
    input_file = Path(__file__).parent / "DryBox" / "input.wav"
    if not input_file.exists():
        print(f"{YELLOW}[SKIP]{RESET} input.wav not found, creating test file...")
        
        # Create a test tone
        subprocess.run([
            "sox", "-n", str(input_file),
            "synth", "1", "sine", "440",
            "rate", "8000"
        ], capture_output=True)
    
    # Create sender and receiver
    sender = IntegratedDryBoxProtocol(mode="sender")
    receiver = IntegratedDryBoxProtocol(mode="receiver")
    
    # Connect and exchange keys
    if not sender.connect_gsm() or not receiver.connect_gsm():
        print(f"{RED}[ERROR]{RESET} Failed to connect to GSM")
        return False
    
    # Setup protocol
    receiver_port = receiver.setup_protocol_connection()
    sender.setup_protocol_connection(
        peer_port=receiver_port,
        peer_identity=receiver.get_identity_key()
    )
    
    # Key exchange
    if not sender.initiate_key_exchange(cipher_type=1):
        print(f"{RED}[ERROR]{RESET} Key exchange failed")
        return False
    
    print(f"{BLUE}[TEST]{RESET} Sending voice with 4FSK modulation...")
    
    # Send voice
    sender.send_voice()
    
    # Wait for transmission
    time.sleep(5)
    
    # Check if receiver created output file
    output_file = Path(__file__).parent / "DryBox" / "received.wav"
    if output_file.exists():
        print(f"{GREEN}[PASS]{RESET} Voice received and decoded!")
        print(f"  Output file: {output_file}")
        
        # Clean up output
        os.remove(output_file)
    else:
        print(f"{RED}[FAIL]{RESET} Voice not received")
        return False
    
    # Clean up
    sender.close()
    receiver.close()
    
    return True


def test_full_integration():
    """Run all integration tests"""
    print(f"{BLUE}{'='*60}{RESET}")
    print(f"{BLUE}DryBox Integration Test Suite{RESET}")
    print(f"{BLUE}{'='*60}{RESET}")
    
    # Start GSM simulator
    gsm_process = start_gsm_simulator()
    
    try:
        # Run tests
        tests_passed = 0
        tests_total = 0
        
        # Test 1: Key Exchange
        tests_total += 1
        if test_key_exchange():
            tests_passed += 1
        
        # Test 2: Voice Transmission
        tests_total += 1
        if test_voice_transmission():
            tests_passed += 1
        
        # Summary
        print(f"\n{BLUE}{'='*60}{RESET}")
        print(f"{BLUE}Test Summary:{RESET}")
        print(f"  Passed: {tests_passed}/{tests_total}")
        
        if tests_passed == tests_total:
            print(f"{GREEN}  All tests passed!{RESET}")
        else:
            print(f"{RED}  Some tests failed{RESET}")
            
    finally:
        # Stop GSM simulator
        print(f"\n{BLUE}[CLEANUP]{RESET} Stopping GSM simulator...")
        gsm_process.terminate()
        gsm_process.wait()


def manual_test():
    """Interactive manual test"""
    print(f"{BLUE}{'='*60}{RESET}")
    print(f"{BLUE}DryBox Manual Test Mode{RESET}")
    print(f"{BLUE}{'='*60}{RESET}")
    
    # Start GSM simulator
    gsm_process = start_gsm_simulator()
    
    try:
        mode = input("\nEnter mode (sender/receiver): ").strip().lower()
        
        # Create protocol instance
        protocol = IntegratedDryBoxProtocol(mode=mode)
        
        if not protocol.connect_gsm():
            print(f"{RED}[ERROR]{RESET} Failed to connect to GSM")
            return
        
        print(f"\n{YELLOW}Protocol Information:{RESET}")
        print(f"  Mode: {mode}")
        print(f"  Identity: {protocol.get_identity_key()}")
        print(f"  Protocol port: {protocol.protocol.local_port}")
        
        if mode == "sender":
            peer_port = input("\nEnter receiver's protocol port: ")
            peer_identity = input("Enter receiver's identity key: ")
            
            protocol.setup_protocol_connection(
                peer_port=int(peer_port),
                peer_identity=peer_identity
            )
            
            print("\nInitiating key exchange...")
            if protocol.initiate_key_exchange(cipher_type=1):
                print(f"{GREEN}Key exchange successful!{RESET}")
                
                while True:
                    print("\nOptions:")
                    print("  1. Send encrypted message")
                    print("  2. Send voice")
                    print("  3. Show status")
                    print("  4. Exit")
                    
                    choice = input("\nChoice: ")
                    
                    if choice == "1":
                        msg = input("Enter message: ")
                        protocol.send_encrypted_message(msg)
                    elif choice == "2":
                        protocol.send_voice()
                    elif choice == "3":
                        protocol.show_status()
                    elif choice == "4":
                        break
        else:
            # Receiver mode
            port = protocol.setup_protocol_connection()
            print(f"\nTell sender to connect to port: {port}")
            print("Waiting for connection...")
            
            try:
                while True:
                    time.sleep(1)
                    if protocol.protocol.state.get("key_exchange_complete"):
                        print(f"{GREEN}Key exchange completed!{RESET}")
                        print("Listening for messages...")
                        
                        while True:
                            time.sleep(1)
            except KeyboardInterrupt:
                pass
                
    finally:
        protocol.close()
        gsm_process.terminate()


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Test DryBox Integration")
    parser.add_argument("--manual", action="store_true", help="Run manual test mode")
    args = parser.parse_args()
    
    if args.manual:
        manual_test()
    else:
        test_full_integration()