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

# Add parent directory to path
parent_dir = str(Path(__file__).parent.parent)
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

# Import the integrated protocol
try:
    # Try importing from same directory first
    from .integrated_protocol import IntegratedDryBoxProtocol
    HAS_INTEGRATED = True
except ImportError:
    try:
        # Try absolute import
        from integrated_protocol import IntegratedDryBoxProtocol
        HAS_INTEGRATED = True
    except ImportError:
        HAS_INTEGRATED = False
        print("Warning: Integrated protocol not available, using basic mode")

# Configuration
HOST = "localhost"
PORT = 12345
INPUT_FILE = "input.wav"
OUTPUT_FILE = "received.wav"

# Global protocol instance
protocol_instance = None


def encrypt_data(data):
    """Encrypt data using the integrated protocol if available"""
    global protocol_instance
    
    if HAS_INTEGRATED and protocol_instance and protocol_instance.protocol.hkdf_key:
        # Use ChaCha20 encryption from protocol
        from encryption import encrypt_message, generate_iv
        key = bytes.fromhex(protocol_instance.protocol.hkdf_key)
        
        # Generate IV
        if protocol_instance.protocol.last_iv is None:
            iv = generate_iv(initial=True)
        else:
            iv = generate_iv(initial=False, previous_iv=protocol_instance.protocol.last_iv)
        
        protocol_instance.protocol.last_iv = iv
        
        # Encrypt with minimal header
        encrypted = encrypt_message(
            plaintext=data,
            key=key,
            flag=0xABCD,
            retry=0,
            connection_status=0,
            iv=iv,
            cipher_type=protocol_instance.protocol.cipher_type
        )
        return encrypted
    else:
        return data  # Fallback to no encryption


def decrypt_data(data):
    """Decrypt data using the integrated protocol if available"""
    global protocol_instance
    
    if HAS_INTEGRATED and protocol_instance and protocol_instance.protocol.hkdf_key:
        # Use decryption from protocol
        from encryption import decrypt_message
        key = bytes.fromhex(protocol_instance.protocol.hkdf_key)
        
        try:
            decrypted = decrypt_message(data, key, protocol_instance.protocol.cipher_type)
            return decrypted
        except Exception as e:
            print(f"Decryption failed: {e}")
            return data
    else:
        return data  # Fallback to no decryption


def run_protocol(send_mode=True):
    """Connect to the simulator and send/receive data."""
    global protocol_instance
    
    # Initialize integrated protocol if available
    if HAS_INTEGRATED:
        mode = "sender" if send_mode else "receiver"
        protocol_instance = IntegratedDryBoxProtocol(gsm_host=HOST, gsm_port=PORT, mode=mode)
        
        # For testing, use predefined keys
        if send_mode:
            # Sender needs receiver's identity
            receiver_identity = "b472a6f5707d4e5e9c6f7e8d9a0b1c2d3e4f5a6b7c8d9e0f1a2b3c4d5e6f7a8b9c0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f5a6b7c8d9e0f1a2b3c4d5e6f7a8b"
            protocol_instance.setup_protocol_connection(peer_port=40000, peer_identity=receiver_identity)
            
            # Try to establish key exchange
            if protocol_instance.initiate_key_exchange(cipher_type=1):
                print("Key exchange successful, using encrypted communication")
            else:
                print("Key exchange failed, falling back to unencrypted")
        else:
            # Receiver listens
            port = protocol_instance.setup_protocol_connection()
            print(f"Protocol listening on port {port}")
    
    # Original GSM connection
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((HOST, PORT))
    print(f"Connected to simulator at {HOST}:{PORT}")

    if send_mode:
        # Check if we should use integrated voice processing
        if HAS_INTEGRATED and protocol_instance and protocol_instance.protocol.hkdf_key:
            # Use integrated voice processing with encryption and FSK
            print("Using integrated voice protocol with encryption and 4FSK modulation")
            protocol_instance.gsm_socket = sock
            protocol_instance.connected = True
            protocol_instance.send_voice()
        else:
            # Fallback to original GSM-only mode
            print("Using basic GSM mode (no encryption)")
            # Sender mode: Encode audio with toast
            os.system(f"toast -p -l {INPUT_FILE}")  # Creates input.wav.gsm
            input_gsm_file = f"{INPUT_FILE}.gsm"
            if not os.path.exists(input_gsm_file):
                print(f"Error: {input_gsm_file} not created")
                sock.close()
                return
            with open(input_gsm_file, "rb") as f:
                voice_data = f.read()

            encrypted_data = encrypt_data(voice_data)
            sock.send(encrypted_data)
            print(f"Sent {len(encrypted_data)} bytes")
            os.remove(input_gsm_file)  # Clean up
    else:
        # Receiver mode
        if HAS_INTEGRATED and protocol_instance:
            # Use integrated receiver with decryption
            print("Using integrated voice protocol receiver")
            protocol_instance.gsm_socket = sock
            protocol_instance.connected = True
            protocol_instance.running = True
            
            # Start receive thread
            import threading
            receive_thread = threading.Thread(target=protocol_instance._receive_loop)
            receive_thread.daemon = True
            receive_thread.start()
            
            # Wait for data
            try:
                time.sleep(30)  # Wait up to 30 seconds
            except KeyboardInterrupt:
                pass
        else:
            # Fallback to original receiver
            print("Using basic GSM receiver (no decryption)")
            print("Waiting for data from sender...")
            received_data = b""
            sock.settimeout(5.0)
            try:
                while True:
                    print("Calling recv()...")
                    data = sock.recv(1024)
                    print(f"Received {len(data)} bytes")
                    if not data:
                        print("Connection closed by sender or simulator")
                        break
                    received_data += data
            except socket.timeout:
                print("Timed out waiting for data")

            if received_data:
                with open("received.gsm", "wb") as f:
                    f.write(decrypt_data(received_data))
                print(f"Wrote {len(received_data)} bytes to received.gsm")
                # Decode with untoast, then convert to WAV with sox
                result = subprocess.run(["untoast", "received.gsm"], capture_output=True, text=True)
                print(f"untoast return code: {result.returncode}")
                print(f"untoast stderr: {result.stderr}")
                if result.returncode == 0:
                    if os.path.exists("received"):
                        # Convert raw PCM to WAV (8 kHz, mono, 16-bit)
                        subprocess.run(["sox", "-t", "raw", "-r", "8000", "-e", "signed", "-b", "16", "-c", "1", "received",
                                        OUTPUT_FILE])
                        os.remove("received")
                        print(f"Received and saved {len(received_data)} bytes to {OUTPUT_FILE}")
                    else:
                        print("Error: 'received' file not created by untoast")
                else:
                    print(f"untoast failed: {result.stderr}")
            else:
                print("No data received from simulator")

    sock.close()
    
    # Clean up protocol instance
    if protocol_instance:
        protocol_instance.close()


if __name__ == "__main__":
    mode = input("Enter 'send' to send data or 'receive' to receive data: ").strip().lower()
    run_protocol(send_mode=(mode == "send"))