#!/usr/bin/env python3
"""
Test script for the voice-over-GSM protocol integration.
This demonstrates encrypted voice transmission using Codec2 and FSK modulation.
"""

import time
import sys
import array
from protocol import IcingProtocol
from voice_codec import Codec2Mode

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


def generate_test_audio(duration_ms: int, frequency: int = 440) -> array.array:
    """Generate test audio (sine wave)."""
    import math
    sample_rate = 8000
    samples = int(sample_rate * duration_ms / 1000)
    audio = array.array('h')  # 16-bit signed integers
    
    for i in range(samples):
        t = i / sample_rate
        value = int(math.sin(2 * math.pi * frequency * t) * 16384)
        audio.append(value)
    
    return audio


def test_voice_protocol():
    """Test voice protocol with two peers."""
    print(f"\n{BLUE}=== Voice Protocol Test ==={RESET}")
    print("This test demonstrates encrypted voice communication.\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 secure channel...{RESET}")
    alice.set_peer_identity(bob.identity_pubkey.hex())
    bob.set_peer_identity(alice.identity_pubkey.hex())
    
    # Establish connection
    alice.connect_to_peer(bob.local_port)
    time.sleep(0.5)
    
    # Perform key exchange
    print(f"\n{YELLOW}2. Performing key exchange...{RESET}")
    
    # Send ping
    alice.send_ping_request(cipher_type=1)  # Use ChaCha20
    time.sleep(0.5)
    
    # Bob responds
    if bob.inbound_messages:
        bob.respond_to_ping(0, answer=1)
        time.sleep(0.5)
    
    # Generate ephemeral keys
    alice.generate_ephemeral_keys()
    bob.generate_ephemeral_keys()
    
    # Exchange handshakes
    alice.send_handshake()
    time.sleep(0.5)
    
    # Bob processes and responds
    if bob.inbound_messages:
        for i, msg in enumerate(bob.inbound_messages):
            if msg["type"] == "HANDSHAKE":
                bob.generate_ecdhe(i)
                bob.send_handshake()
                break
    time.sleep(0.5)
    
    # Alice processes Bob's handshake
    if alice.inbound_messages:
        for i, msg in enumerate(alice.inbound_messages):
            if msg["type"] == "HANDSHAKE":
                alice.generate_ecdhe(i)
                break
    
    # Derive keys
    alice.derive_hkdf()
    bob.derive_hkdf()
    
    print(f"{GREEN}   Secure channel established!{RESET}")
    
    # Start voice call
    print(f"\n{YELLOW}3. Initiating voice call...{RESET}")
    alice.start_voice_call(codec_mode=5, fec_type=0)  # 1200bps, repetition FEC
    time.sleep(0.5)
    
    # Check if Bob received the call and accept it manually
    voice_active = False
    for i, msg in enumerate(bob.inbound_messages):
        if msg["type"] == "VOICE_START":
            voice_start = msg["parsed"]
            print(f"   Bob accepting voice call...")
            bob.accept_voice_call(voice_start.session_id, voice_start.codec_mode, voice_start.fec_type)
            time.sleep(0.5)
            voice_active = True
            break
    
    if voice_active and alice.voice_session_active:
        print(f"{GREEN}   Voice call established!{RESET}")
        print(f"   Session ID: {alice.voice_session_id:016x}")
    else:
        print(f"{RED}   Voice call failed to establish{RESET}")
    
    if voice_active:
        # Test voice transmission
        print(f"\n{YELLOW}4. Testing voice transmission...{RESET}")
        
        # Generate test audio (440Hz tone for 200ms)
        test_audio = generate_test_audio(200, 440)
        print(f"   Generated {len(test_audio)} audio samples")
        
        # Alice sends audio
        print(f"\n   Alice sending audio...")
        # Convert array to numpy array if needed
        import array
        if isinstance(test_audio, array.array):
            # Voice protocol expects raw array or list
            audio_data = test_audio
        else:
            audio_data = test_audio
            
        success = alice.send_voice_audio(audio_data)
        if success:
            print(f"{GREEN}   Audio processed and modulated{RESET}")
        else:
            print(f"{RED}   Failed to process audio{RESET}")
        
        # Test voice codec directly
        print(f"\n{YELLOW}5. Testing voice codec components...{RESET}")
        
        if alice.voice_protocol:
            # Test Codec2
            print(f"\n   Testing Codec2 compression...")
            # Get one frame worth of samples
            if hasattr(test_audio, '__getitem__'):
                frame_audio = test_audio[:320] if len(test_audio) >= 320 else test_audio
            else:
                frame_audio = list(test_audio)[:320]
            
            codec_frame = alice.voice_protocol.codec.encode(frame_audio)
            if codec_frame:
                print(f"   Compressed to {len(codec_frame.bits)} bytes")
                
                # Test decompression
                decoded = alice.voice_protocol.codec.decode(codec_frame)
                print(f"   Decompressed to {len(decoded)} samples")
            
            # Test FSK modulation
            print(f"\n   Testing FSK modulation...")
            test_data = b"Voice test data"
            modulated = alice.voice_protocol.modem.modulate(test_data)
            print(f"   Modulated {len(test_data)} bytes to {len(modulated)} audio samples")
            
            # Test demodulation
            demodulated, confidence = alice.voice_protocol.modem.demodulate(modulated)
            print(f"   Demodulated with {confidence:.1%} confidence")
            print(f"   Data match: {demodulated == test_data}")
        
        # Send sync frame
        print(f"\n{YELLOW}6. Testing synchronization...{RESET}")
        from messages import VoiceSync
        sync_msg = VoiceSync(
            session_id=alice.voice_session_id,
            sequence=1,
            timestamp=100
        )
        alice._send_packet(alice.connections[0], sync_msg.serialize(), "VOICE_SYNC")
        time.sleep(0.5)
        
        # End voice call
        print(f"\n{YELLOW}7. Ending voice call...{RESET}")
        alice.end_voice_call()
        time.sleep(0.5)
    
    # Show final state
    print(f"\n{YELLOW}8. Final state:{RESET}")
    print("\nAlice voice status:")
    print(f"  Active: {alice.voice_session_active}")
    print(f"  Voice codec initialized: {alice.voice_protocol is not None}")
    
    print("\nBob voice status:")
    print(f"  Active: {bob.voice_session_active}")
    print(f"  Voice codec initialized: {bob.voice_protocol is not None}")
    
    # Cleanup
    alice.stop()
    bob.stop()
    
    print(f"\n{GREEN}Voice protocol test completed!{RESET}")


def test_codec_modes():
    """Test different Codec2 modes."""
    print(f"\n{BLUE}=== Codec2 Mode Comparison ==={RESET}")
    
    from voice_codec import Codec2Wrapper, Codec2Mode
    
    modes = [
        (Codec2Mode.MODE_3200, "3200 bps"),
        (Codec2Mode.MODE_2400, "2400 bps"),
        (Codec2Mode.MODE_1600, "1600 bps"),
        (Codec2Mode.MODE_1400, "1400 bps"),
        (Codec2Mode.MODE_1300, "1300 bps"),
        (Codec2Mode.MODE_1200, "1200 bps (recommended)"),
        (Codec2Mode.MODE_700C, "700 bps")
    ]
    
    # Generate test audio
    test_audio = generate_test_audio(100, 440)
    
    print("\nMode comparison:")
    print("-" * 50)
    
    for mode, description in modes:
        try:
            codec = Codec2Wrapper(mode)
            
            # Process one frame
            if hasattr(test_audio, '__getitem__'):
                frame_audio = test_audio[:codec.frame_samples]
            else:
                frame_audio = list(test_audio)[:codec.frame_samples]
                
            if len(frame_audio) < codec.frame_samples:
                # Pad if necessary
                frame_audio = frame_audio + [0] * (codec.frame_samples - len(frame_audio))
            
            frame = codec.encode(frame_audio)
            
            if frame:
                efficiency = (codec.frame_bits / 8) / (codec.frame_ms / 1000) / 1000  # KB/s
                print(f"{description:20} | {codec.frame_bits:3} bits/frame | "
                      f"{codec.frame_ms:2}ms | {efficiency:.2f} KB/s")
            
        except Exception as e:
            print(f"{description:20} | Error: {e}")
    
    print("-" * 50)
    print(f"\n{YELLOW}Note: Lower bitrates provide better GSM vocoder survival{RESET}")
    print(f"{YELLOW}      but reduced voice quality. 1200 bps is recommended.{RESET}")


def main():
    """Main test function."""
    print(f"{BLUE}{'='*60}{RESET}")
    print(f"{BLUE}    Voice-over-GSM Protocol Test Suite{RESET}")
    print(f"{BLUE}{'='*60}{RESET}")
    
    print("\nSelect test:")
    print("1. Full voice protocol test")
    print("2. Codec2 mode comparison")
    print("3. Run both tests")
    print("0. Exit")
    
    try:
        choice = input("\nEnter your choice (0-3): ").strip()
        
        if choice == "1":
            test_voice_protocol()
        elif choice == "2":
            test_codec_modes()
        elif choice == "3":
            test_voice_protocol()
            print(f"\n{YELLOW}{'='*60}{RESET}\n")
            test_codec_modes()
        elif choice == "0":
            print("Exiting...")
            return 0
        else:
            print(f"{RED}Invalid choice.{RESET}")
            return 1
            
    except KeyboardInterrupt:
        print(f"\n\n{YELLOW}Test interrupted.{RESET}")
        return 0
    except Exception as e:
        print(f"\n{RED}ERROR: {e}{RESET}")
        return 1
    
    return 0


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