#!/usr/bin/env python3
"""
Integrated protocol for DryBox - combines Icing protocol with GSM simulator
Supports encrypted voice communication with 4FSK modulation
"""

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

# Add parent directory to path to import protocol modules
parent_dir = str(Path(__file__).parent.parent)
current_dir = str(Path(__file__).parent)

# Remove current directory from path temporarily to avoid importing local protocol.py
if current_dir in sys.path:
    sys.path.remove(current_dir)
    
# Add parent directory at the beginning
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

# Import from parent directory
from protocol import IcingProtocol
from voice_codec import VoiceProtocol, FSKModem, Codec2Wrapper, Codec2Mode
from encryption import encrypt_message, decrypt_message, generate_iv
import transmission

# Add current directory back
if current_dir not in sys.path:
    sys.path.append(current_dir)

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


class IntegratedDryBoxProtocol:
    """Integrates Icing protocol with DryBox GSM simulator"""
    
    def __init__(self, gsm_host="localhost", gsm_port=12345, mode="sender"):
        """
        Initialize integrated protocol
        
        Args:
            gsm_host: GSM simulator host
            gsm_port: GSM simulator port
            mode: "sender" or "receiver"
        """
        self.gsm_host = gsm_host
        self.gsm_port = gsm_port
        self.mode = mode
        
        # Initialize Icing protocol
        self.protocol = IcingProtocol()
        
        # GSM connection
        self.gsm_socket = None
        self.connected = False
        
        # Voice processing
        self.voice_protocol = None
        self.modem = FSKModem(sample_rate=8000, baud_rate=600)
        self.codec = Codec2Wrapper(Codec2Mode.MODE_1200)
        
        # Audio files
        self.input_file = "input.wav"
        self.output_file = "received.wav"
        
        # Threading
        self.receive_thread = None
        self.running = False
        
        print(f"{GREEN}[DRYBOX]{RESET} Initialized in {mode} mode")
    
    def connect_gsm(self):
        """Connect to GSM simulator"""
        try:
            self.gsm_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.gsm_socket.connect((self.gsm_host, self.gsm_port))
            self.connected = True
            print(f"{GREEN}[GSM]{RESET} Connected to simulator at {self.gsm_host}:{self.gsm_port}")
            
            # Start receive thread
            self.running = True
            self.receive_thread = threading.Thread(target=self._receive_loop)
            self.receive_thread.daemon = True
            self.receive_thread.start()
            
            return True
        except Exception as e:
            print(f"{RED}[ERROR]{RESET} GSM connection failed: {e}")
            return False
    
    def setup_protocol_connection(self, peer_port=None, peer_identity=None):
        """
        Setup Icing protocol connection
        
        Args:
            peer_port: Port to connect to (for initiator)
            peer_identity: Peer's identity public key hex (required)
        """
        if peer_identity:
            self.protocol.set_peer_identity(peer_identity)
        
        if peer_port:
            # Connect to peer
            self.protocol.connect_to_peer(peer_port)
            print(f"{GREEN}[PROTOCOL]{RESET} Connected to peer on port {peer_port}")
        else:
            print(f"{GREEN}[PROTOCOL]{RESET} Listening on port {self.protocol.local_port}")
        
        return self.protocol.local_port
    
    def initiate_key_exchange(self, cipher_type=1):
        """
        Initiate key exchange with ChaCha20-Poly1305 by default
        
        Args:
            cipher_type: 0=AES-GCM, 1=ChaCha20-Poly1305
        """
        print(f"{BLUE}[KEY-EXCHANGE]{RESET} Starting key exchange...")
        
        # Enable auto mode for automatic handshake
        self.protocol.configure_auto_mode(
            ping_response_accept=True,
            preferred_cipher=cipher_type,
            active_mode=True
        )
        self.protocol.start_auto_mode()
        
        # Send initial ping
        self.protocol.send_ping_request(cipher_type)
        
        # Wait for key exchange to complete
        timeout = 10
        start_time = time.time()
        while not self.protocol.state.get("key_exchange_complete") and time.time() - start_time < timeout:
            time.sleep(0.1)
        
        if self.protocol.state.get("key_exchange_complete"):
            print(f"{GREEN}[KEY-EXCHANGE]{RESET} Key exchange completed!")
            print(f"  Cipher: {'ChaCha20-Poly1305' if self.protocol.cipher_type == 1 else 'AES-256-GCM'}")
            print(f"  HKDF Key: {self.protocol.hkdf_key[:16]}...")
            
            # Initialize voice protocol with encryption key
            self.voice_protocol = VoiceProtocol(self.protocol)
            return True
        else:
            print(f"{RED}[ERROR]{RESET} Key exchange timeout")
            return False
    
    def send_voice(self):
        """Send voice data through GSM channel"""
        if not self.connected:
            print(f"{RED}[ERROR]{RESET} Not connected to GSM")
            return
        
        if not self.protocol.hkdf_key:
            print(f"{RED}[ERROR]{RESET} No encryption key available")
            return
        
        # Encode audio with GSM codec
        if os.path.exists(self.input_file):
            print(f"{BLUE}[VOICE]{RESET} Processing {self.input_file}...")
            
            # Convert to 8kHz mono if needed
            input_8k = "input_8k_mono.wav"
            subprocess.run([
                "sox", self.input_file, "-r", "8000", "-c", "1", input_8k
            ], capture_output=True)
            
            # Read PCM audio
            with open(input_8k, 'rb') as f:
                # Skip WAV header (44 bytes)
                f.seek(44)
                pcm_data = f.read()
            
            # Convert to samples
            samples = struct.unpack(f'{len(pcm_data)//2}h', pcm_data)
            
            # Process through voice protocol (compress, encrypt, modulate)
            modulated = self.voice_protocol.process_voice_input(samples)
            
            if modulated is not None:
                # Convert float samples to bytes for transmission
                if hasattr(modulated, 'tobytes'):
                    # numpy array
                    transmit_data = (modulated * 32767).astype('int16').tobytes()
                else:
                    # array.array
                    transmit_data = struct.pack(f'{len(modulated)}h', 
                                              *[int(s * 32767) for s in modulated])
                
                # Send through GSM
                self.gsm_socket.send(transmit_data)
                print(f"{GREEN}[VOICE]{RESET} Sent {len(transmit_data)} bytes")
                
                # Clean up
                os.remove(input_8k)
            else:
                print(f"{RED}[ERROR]{RESET} Voice processing failed")
        else:
            print(f"{RED}[ERROR]{RESET} Input file {self.input_file} not found")
    
    def _receive_loop(self):
        """Background thread to receive data from GSM"""
        self.gsm_socket.settimeout(0.5)
        received_data = b""
        
        while self.running:
            try:
                data = self.gsm_socket.recv(4096)
                if not data:
                    print(f"{YELLOW}[GSM]{RESET} Connection closed")
                    break
                
                received_data += data
                
                # Process when we have enough data (at least 1 second of audio)
                if len(received_data) >= 16000:  # 8000 Hz * 2 bytes * 1 second
                    self._process_received_audio(received_data)
                    received_data = b""
                    
            except socket.timeout:
                # Process any remaining data
                if received_data:
                    self._process_received_audio(received_data)
                    received_data = b""
            except Exception as e:
                print(f"{RED}[ERROR]{RESET} Receive error: {e}")
                break
    
    def _process_received_audio(self, data):
        """Process received audio data"""
        if not self.voice_protocol:
            print(f"{YELLOW}[WARN]{RESET} Voice protocol not initialized, storing raw audio")
            # Just save raw audio
            with open("received_raw.pcm", "wb") as f:
                f.write(data)
            return
        
        print(f"{BLUE}[RECEIVE]{RESET} Processing {len(data)} bytes...")
        
        try:
            # Convert bytes to float samples
            samples = struct.unpack(f'{len(data)//2}h', data)
            float_samples = [s / 32768.0 for s in samples]
            
            # Demodulate, decrypt, decompress
            pcm_output = self.voice_protocol.process_voice_output(float_samples)
            
            if pcm_output is not None:
                # Save as WAV file
                self._save_wav(pcm_output, self.output_file)
                print(f"{GREEN}[VOICE]{RESET} Saved decoded audio to {self.output_file}")
            else:
                print(f"{YELLOW}[WARN]{RESET} Could not decode audio")
                
        except Exception as e:
            print(f"{RED}[ERROR]{RESET} Audio processing failed: {e}")
            import traceback
            traceback.print_exc()
    
    def _save_wav(self, samples, filename):
        """Save PCM samples as WAV file"""
        import wave
        
        with wave.open(filename, 'wb') as wav:
            wav.setnchannels(1)  # Mono
            wav.setsampwidth(2)  # 16-bit
            wav.setframerate(8000)  # 8kHz
            
            if hasattr(samples, 'tobytes'):
                # numpy array
                wav.writeframes(samples.tobytes())
            else:
                # array.array or list
                if hasattr(samples, 'tobytes'):
                    wav.writeframes(samples.tobytes())
                else:
                    # Convert list to bytes
                    wav.writeframes(struct.pack(f'{len(samples)}h', *samples))
    
    def send_encrypted_message(self, message):
        """Send an encrypted text message"""
        if self.protocol.hkdf_key:
            self.protocol.send_encrypted_message(message)
            print(f"{GREEN}[MESSAGE]{RESET} Sent encrypted: {message}")
        else:
            print(f"{RED}[ERROR]{RESET} No encryption key available")
    
    def close(self):
        """Clean up connections"""
        self.running = False
        
        if self.receive_thread:
            self.receive_thread.join(timeout=1)
        
        if self.gsm_socket:
            self.gsm_socket.close()
        
        self.protocol.stop()
        print(f"{RED}[SHUTDOWN]{RESET} Protocol closed")
    
    def get_identity_key(self):
        """Get our identity public key"""
        return self.protocol.identity_pubkey.hex()
    
    def show_status(self):
        """Show protocol status"""
        self.protocol.show_state()


def test_integrated_protocol():
    """Test the integrated protocol"""
    import sys
    
    mode = sys.argv[1] if len(sys.argv) > 1 else "sender"
    
    # Create protocol instance
    drybox = IntegratedDryBoxProtocol(mode=mode)
    
    # Connect to GSM simulator
    if not drybox.connect_gsm():
        return
    
    print(f"\n{YELLOW}=== DryBox Protocol Test ==={RESET}")
    print(f"Mode: {mode}")
    print(f"Identity key: {drybox.get_identity_key()[:32]}...")
    
    if mode == "sender":
        # Get receiver's identity (in real app, this would be exchanged out-of-band)
        receiver_identity = input("\nEnter receiver's identity key (or press Enter to use test key): ").strip()
        if not receiver_identity:
            # Use a test key
            receiver_identity = "b472a6f5707d4e5e9c6f7e8d9a0b1c2d3e4f5a6b7c8d9e0f1a2b3c4d5e6f7a8b9c0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f5a6b7c8d9e0f1a2b3c4d5e6f7a8b"
        
        # Setup protocol connection
        peer_port = int(input("Enter peer's protocol port: "))
        drybox.setup_protocol_connection(peer_port=peer_port, peer_identity=receiver_identity)
        
        # Initiate key exchange
        if drybox.initiate_key_exchange(cipher_type=1):  # Use ChaCha20
            # Send test message
            drybox.send_encrypted_message("Hello from DryBox!")
            
            # Send voice
            time.sleep(1)
            drybox.send_voice()
        
    else:  # receiver
        # Setup protocol listener
        port = drybox.setup_protocol_connection()
        print(f"\nTell sender to connect to port: {port}")
        print(f"Your identity key: {drybox.get_identity_key()}")
        
        # Wait for connection
        print("\nWaiting for connection...")
        
    # Keep running
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n\nShutting down...")
        drybox.close()


if __name__ == "__main__":
    test_integrated_protocol()