#!/usr/bin/env python3
"""
Integrated UI for DryBox with Icing Protocol
Supports encrypted voice communication with 4FSK modulation
"""

import sys
import random
import socket
import threading
import time
import subprocess
import os
from pathlib import Path
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QFrame, QSizePolicy, QStyle, QTextEdit,
    QLineEdit, QCheckBox
)
from PyQt5.QtCore import Qt, QTimer, QSize, QPointF, pyqtSignal, QThread
from PyQt5.QtGui import QPainter, QColor, QPen, QLinearGradient, QBrush, QIcon, QFont

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

# Import from DryBox directory
from integrated_protocol import IntegratedDryBoxProtocol

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


class ProtocolThread(QThread):
    """Thread for running the integrated protocol"""
    status_update = pyqtSignal(str)
    key_exchange_complete = pyqtSignal(bool)
    message_received = pyqtSignal(str)
    
    def __init__(self, mode, gsm_host="localhost", gsm_port=12345):
        super().__init__()
        self.mode = mode
        self.gsm_host = gsm_host
        self.gsm_port = gsm_port
        self.protocol = None
        self.running = True
        
    def run(self):
        """Run the protocol in background"""
        try:
            # Create protocol instance
            self.protocol = IntegratedDryBoxProtocol(
                gsm_host=self.gsm_host,
                gsm_port=self.gsm_port,
                mode=self.mode
            )
            
            self.status_update.emit(f"Protocol initialized in {self.mode} mode")
            
            # Connect to GSM
            if self.protocol.connect_gsm():
                self.status_update.emit("Connected to GSM simulator")
            else:
                self.status_update.emit("Failed to connect to GSM")
                return
            
            # Get identity
            identity = self.protocol.get_identity_key()
            self.status_update.emit(f"Identity: {identity[:32]}...")
            
            # Keep running
            while self.running:
                time.sleep(0.1)
                
                # Check for key exchange completion
                if (self.protocol.protocol.state.get("key_exchange_complete") and 
                    not hasattr(self, '_key_exchange_notified')):
                    self._key_exchange_notified = True
                    self.key_exchange_complete.emit(True)
                    
        except Exception as e:
            self.status_update.emit(f"Protocol error: {str(e)}")
            
    def stop(self):
        """Stop the protocol thread"""
        self.running = False
        if self.protocol:
            self.protocol.close()
            
    def setup_connection(self, peer_port=None, peer_identity=None):
        """Setup protocol connection"""
        if self.protocol:
            port = self.protocol.setup_protocol_connection(
                peer_port=peer_port,
                peer_identity=peer_identity
            )
            return port
        return None
        
    def initiate_key_exchange(self, cipher_type=1):
        """Initiate key exchange"""
        if self.protocol:
            return self.protocol.initiate_key_exchange(cipher_type)
        return False
        
    def send_voice(self, audio_file):
        """Send voice through protocol"""
        if self.protocol:
            # Temporarily set input file
            old_input = self.protocol.input_file
            self.protocol.input_file = audio_file
            self.protocol.send_voice()
            self.protocol.input_file = old_input
            
    def send_message(self, message):
        """Send encrypted text message"""
        if self.protocol:
            self.protocol.send_encrypted_message(message)


class WaveformWidget(QWidget):
    """Widget for displaying audio waveform"""
    def __init__(self, parent=None, dynamic=False):
        super().__init__(parent)
        self.dynamic = dynamic
        self.setMinimumSize(200, 80)
        self.setMaximumHeight(100)
        self.waveform_data = [random.randint(10, 90) for _ in range(50)]
        if self.dynamic:
            self.timer = QTimer(self)
            self.timer.timeout.connect(self.update_waveform)
            self.timer.start(100)

    def update_waveform(self):
        self.waveform_data = self.waveform_data[1:] + [random.randint(10, 90)]
        self.update()

    def set_data(self, data):
        amplitude = sum(byte for byte in data) % 90 + 10
        self.waveform_data = self.waveform_data[1:] + [amplitude]
        self.update()

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.fillRect(self.rect(), QColor("#2D2D2D"))
        gradient = QLinearGradient(0, 0, 0, self.height())
        gradient.setColorAt(0.0, QColor("#0078D4"))
        gradient.setColorAt(1.0, QColor("#50E6A4"))
        pen = QPen(QBrush(gradient), 2)
        painter.setPen(pen)
        bar_width = self.width() / len(self.waveform_data)
        max_h = self.height() - 10
        for i, val in enumerate(self.waveform_data):
            bar_height = (val / 100.0) * max_h
            x = i * bar_width
            y = (self.height() - bar_height) / 2
            painter.drawLine(QPointF(x + bar_width / 2, y), 
                           QPointF(x + bar_width / 2, y + bar_height))


class IntegratedPhoneUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("DryBox Integrated Protocol UI")
        self.setGeometry(100, 100, 1000, 800)
        self.setStyleSheet("""
            QMainWindow { background-color: #1e1e1e; }
            QLabel { color: #E0E0E0; font-size: 14px; }
            QPushButton {
                background-color: #0078D4; color: white; border: none;
                padding: 10px 15px; border-radius: 5px; font-size: 14px;
                min-height: 30px;
            }
            QPushButton:hover { background-color: #005A9E; }
            QPushButton:pressed { background-color: #003C6B; }
            QPushButton:disabled { background-color: #555555; }
            QPushButton#dangerButton { background-color: #E81123; }
            QPushButton#dangerButton:hover { background-color: #C50E1F; }
            QPushButton#successButton { background-color: #107C10; }
            QPushButton#successButton:hover { background-color: #0E6E0E; }
            QFrame {
                background-color: #2D2D2D; border: 1px solid #3D3D3D;
                border-radius: 8px;
            }
            QTextEdit {
                background-color: #1E1E1E; color: #E0E0E0;
                border: 1px solid #3D3D3D; border-radius: 4px;
                font-family: 'Consolas', 'Monaco', monospace;
                padding: 5px;
            }
            QLineEdit {
                background-color: #2D2D2D; color: #E0E0E0;
                border: 1px solid #3D3D3D; border-radius: 4px;
                padding: 5px;
            }
            QCheckBox { color: #E0E0E0; }
            QLabel#titleLabel {
                font-size: 24px; font-weight: bold; color: #00A2E8;
                padding: 15px;
            }
            QLabel#sectionLabel {
                font-size: 16px; font-weight: bold; color: #FFFFFF;
                padding: 5px;
            }
        """)
        
        # Protocol threads
        self.phone1_protocol = None
        self.phone2_protocol = None
        
        # Setup UI
        self.setup_ui()
        
    def setup_ui(self):
        """Setup the user interface"""
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QVBoxLayout()
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_widget.setLayout(main_layout)
        
        # Title
        title = QLabel("DryBox Encrypted Voice Protocol")
        title.setObjectName("titleLabel")
        title.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(title)
        
        # Horizontal layout for phones
        phones_layout = QHBoxLayout()
        phones_layout.setSpacing(20)
        main_layout.addLayout(phones_layout)
        
        # Phone 1
        self.phone1_frame = self.create_phone_frame("Phone 1", 1)
        phones_layout.addWidget(self.phone1_frame)
        
        # Phone 2
        self.phone2_frame = self.create_phone_frame("Phone 2", 2)
        phones_layout.addWidget(self.phone2_frame)
        
        # Protocol status
        status_frame = QFrame()
        status_layout = QVBoxLayout(status_frame)
        
        status_label = QLabel("Protocol Status")
        status_label.setObjectName("sectionLabel")
        status_layout.addWidget(status_label)
        
        self.status_text = QTextEdit()
        self.status_text.setMaximumHeight(150)
        self.status_text.setReadOnly(True)
        status_layout.addWidget(self.status_text)
        
        main_layout.addWidget(status_frame)
        
        # Control buttons
        controls_layout = QHBoxLayout()
        controls_layout.setSpacing(10)
        
        self.start_gsm_btn = QPushButton("Start GSM Simulator")
        self.start_gsm_btn.clicked.connect(self.start_gsm_simulator)
        controls_layout.addWidget(self.start_gsm_btn)
        
        self.test_voice_btn = QPushButton("Test Voice Transmission")
        self.test_voice_btn.clicked.connect(self.test_voice_transmission)
        self.test_voice_btn.setEnabled(False)
        controls_layout.addWidget(self.test_voice_btn)
        
        self.auto_test_btn = QPushButton("Run Auto Test")
        self.auto_test_btn.clicked.connect(self.run_auto_test)
        self.auto_test_btn.setEnabled(False)
        self.auto_test_btn.setObjectName("successButton")
        controls_layout.addWidget(self.auto_test_btn)
        
        controls_layout.addStretch()
        main_layout.addLayout(controls_layout)
        
    def create_phone_frame(self, title, phone_id):
        """Create a phone control frame"""
        frame = QFrame()
        layout = QVBoxLayout(frame)
        
        # Title
        title_label = QLabel(title)
        title_label.setObjectName("sectionLabel")
        title_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(title_label)
        
        # Status
        status_label = QLabel("Disconnected")
        status_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(status_label)
        
        # Identity
        identity_label = QLabel("Identity: Not initialized")
        identity_label.setWordWrap(True)
        identity_label.setStyleSheet("font-size: 10px;")
        layout.addWidget(identity_label)
        
        # Connection controls
        conn_layout = QHBoxLayout()
        
        port_input = QLineEdit()
        port_input.setPlaceholderText("Peer port")
        port_input.setMaximumWidth(100)
        conn_layout.addWidget(port_input)
        
        connect_btn = QPushButton("Connect")
        connect_btn.clicked.connect(lambda: self.connect_phone(phone_id))
        conn_layout.addWidget(connect_btn)
        
        layout.addLayout(conn_layout)
        
        # Key exchange
        key_btn = QPushButton("Start Key Exchange")
        key_btn.clicked.connect(lambda: self.start_key_exchange(phone_id))
        key_btn.setEnabled(False)
        layout.addWidget(key_btn)
        
        # Cipher selection
        cipher_layout = QHBoxLayout()
        aes_radio = QCheckBox("AES-GCM")
        chacha_radio = QCheckBox("ChaCha20")
        chacha_radio.setChecked(True)
        cipher_layout.addWidget(aes_radio)
        cipher_layout.addWidget(chacha_radio)
        layout.addLayout(cipher_layout)
        
        # Message input
        msg_input = QLineEdit()
        msg_input.setPlaceholderText("Enter message")
        layout.addWidget(msg_input)
        
        send_btn = QPushButton("Send Encrypted Message")
        send_btn.clicked.connect(lambda: self.send_message(phone_id))
        send_btn.setEnabled(False)
        layout.addWidget(send_btn)
        
        # Voice controls
        voice_btn = QPushButton("Send Voice")
        voice_btn.clicked.connect(lambda: self.send_voice(phone_id))
        voice_btn.setEnabled(False)
        voice_btn.setObjectName("successButton")
        layout.addWidget(voice_btn)
        
        # Waveform
        waveform = WaveformWidget()
        layout.addWidget(waveform)
        
        # Store references
        frame.status_label = status_label
        frame.identity_label = identity_label
        frame.port_input = port_input
        frame.connect_btn = connect_btn
        frame.key_btn = key_btn
        frame.aes_radio = aes_radio
        frame.chacha_radio = chacha_radio
        frame.msg_input = msg_input
        frame.send_btn = send_btn
        frame.voice_btn = voice_btn
        frame.waveform = waveform
        
        return frame
        
    def start_gsm_simulator(self):
        """Start the GSM simulator in background"""
        self.log_status("Starting GSM simulator...")
        
        # Check if simulator is already running
        try:
            test_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            test_sock.connect(("localhost", 12345))
            test_sock.close()
            self.log_status("GSM simulator already running")
            self.enable_phones()
            return
        except:
            pass
        
        # Start simulator
        gsm_path = Path(__file__).parent.parent / "gsm_simulator.py"
        self.gsm_process = subprocess.Popen(
            [sys.executable, str(gsm_path)],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        
        time.sleep(1)  # Give it time to start
        self.log_status("GSM simulator started")
        self.enable_phones()
        
    def enable_phones(self):
        """Enable phone controls"""
        self.phone1_frame.connect_btn.setEnabled(True)
        self.phone2_frame.connect_btn.setEnabled(True)
        self.auto_test_btn.setEnabled(True)
        
        # Start protocol threads
        self.phone1_protocol = ProtocolThread("sender")
        self.phone1_protocol.status_update.connect(
            lambda msg: self.update_phone_status(1, msg))
        self.phone1_protocol.key_exchange_complete.connect(
            lambda: self.on_key_exchange_complete(1))
        self.phone1_protocol.start()
        
        self.phone2_protocol = ProtocolThread("receiver")
        self.phone2_protocol.status_update.connect(
            lambda msg: self.update_phone_status(2, msg))
        self.phone2_protocol.key_exchange_complete.connect(
            lambda: self.on_key_exchange_complete(2))
        self.phone2_protocol.start()
        
        # Update identities
        time.sleep(0.5)
        if self.phone1_protocol.protocol:
            identity = self.phone1_protocol.protocol.get_identity_key()
            self.phone1_frame.identity_label.setText(f"Identity: {identity[:32]}...")
            
        if self.phone2_protocol.protocol:
            identity = self.phone2_protocol.protocol.get_identity_key()
            self.phone2_frame.identity_label.setText(f"Identity: {identity[:32]}...")
            
    def connect_phone(self, phone_id):
        """Connect phone to peer"""
        if phone_id == 1:
            frame = self.phone1_frame
            protocol = self.phone1_protocol
            peer_protocol = self.phone2_protocol
        else:
            frame = self.phone2_frame
            protocol = self.phone2_protocol
            peer_protocol = self.phone1_protocol
            
        try:
            # Get peer port
            peer_port = frame.port_input.text()
            if not peer_port:
                # Use other phone's port
                if peer_protocol and peer_protocol.protocol:
                    peer_port = peer_protocol.protocol.protocol.local_port
                else:
                    self.log_status(f"Phone {phone_id}: Enter peer port")
                    return
            else:
                peer_port = int(peer_port)
                
            # Get peer identity
            if peer_protocol and peer_protocol.protocol:
                peer_identity = peer_protocol.protocol.get_identity_key()
            else:
                peer_identity = None
                
            # Setup connection
            port = protocol.setup_connection(
                peer_port=peer_port,
                peer_identity=peer_identity
            )
            
            self.log_status(f"Phone {phone_id}: Connected to port {peer_port}")
            frame.status_label.setText("Connected")
            frame.key_btn.setEnabled(True)
            
        except Exception as e:
            self.log_status(f"Phone {phone_id} connection error: {str(e)}")
            
    def start_key_exchange(self, phone_id):
        """Start key exchange for phone"""
        if phone_id == 1:
            frame = self.phone1_frame
            protocol = self.phone1_protocol
        else:
            frame = self.phone2_frame
            protocol = self.phone2_protocol
            
        # Get cipher preference
        cipher_type = 1 if frame.chacha_radio.isChecked() else 0
        
        self.log_status(f"Phone {phone_id}: Starting key exchange...")
        
        # Start key exchange in thread
        threading.Thread(
            target=lambda: protocol.initiate_key_exchange(cipher_type),
            daemon=True
        ).start()
        
    def on_key_exchange_complete(self, phone_id):
        """Handle key exchange completion"""
        if phone_id == 1:
            frame = self.phone1_frame
        else:
            frame = self.phone2_frame
            
        self.log_status(f"Phone {phone_id}: Key exchange completed!")
        frame.status_label.setText("Secure - Key Exchanged")
        frame.send_btn.setEnabled(True)
        frame.voice_btn.setEnabled(True)
        self.test_voice_btn.setEnabled(True)
        
    def send_message(self, phone_id):
        """Send encrypted message"""
        if phone_id == 1:
            frame = self.phone1_frame
            protocol = self.phone1_protocol
        else:
            frame = self.phone2_frame
            protocol = self.phone2_protocol
            
        message = frame.msg_input.text()
        if message:
            protocol.send_message(message)
            self.log_status(f"Phone {phone_id}: Sent encrypted: {message}")
            frame.msg_input.clear()
            
    def send_voice(self, phone_id):
        """Send voice from phone"""
        if phone_id == 1:
            protocol = self.phone1_protocol
        else:
            protocol = self.phone2_protocol
            
        # Check if input.wav exists
        audio_file = Path(__file__).parent.parent / "input.wav"
        if not audio_file.exists():
            self.log_status(f"Phone {phone_id}: input.wav not found")
            return
            
        self.log_status(f"Phone {phone_id}: Sending voice...")
        
        # Send in thread
        threading.Thread(
            target=lambda: protocol.send_voice(str(audio_file)),
            daemon=True
        ).start()
        
    def test_voice_transmission(self):
        """Test full voice transmission"""
        self.log_status("Testing voice transmission from Phone 1 to Phone 2...")
        self.send_voice(1)
        
    def run_auto_test(self):
        """Run automated test sequence"""
        self.log_status("="*50)
        self.log_status("Starting Automated Test Sequence")
        self.log_status("="*50)
        
        # Disable auto test button during test
        self.auto_test_btn.setEnabled(False)
        
        # Run test in a separate thread to avoid blocking UI
        threading.Thread(target=self._run_auto_test_sequence, daemon=True).start()
        
    def _run_auto_test_sequence(self):
        """Execute the automated test sequence"""
        try:
            # Test 1: Auto-connect phones
            self.log_status("\n[TEST 1] Auto-connecting phones...")
            time.sleep(0.5)
            
            # Wait for protocols to be ready
            if not self.phone1_protocol or not self.phone2_protocol:
                self.log_status("❌ Protocols not initialized")
                self.auto_test_btn.setEnabled(True)
                return
                
            # Wait a bit for protocols to fully initialize
            max_wait = 5
            wait_time = 0
            while wait_time < max_wait:
                if (hasattr(self.phone1_protocol, 'protocol') and 
                    hasattr(self.phone2_protocol, 'protocol') and
                    self.phone1_protocol.protocol and 
                    self.phone2_protocol.protocol):
                    break
                time.sleep(0.5)
                wait_time += 0.5
                
            if wait_time >= max_wait:
                self.log_status("❌ Protocols failed to initialize")
                self.auto_test_btn.setEnabled(True)
                return
                
            # Get ports
            phone1_port = self.phone1_protocol.protocol.protocol.local_port
            phone2_port = self.phone2_protocol.protocol.protocol.local_port
            
            # Auto-fill peer ports
            self.phone1_frame.port_input.setText(str(phone2_port))
            self.phone2_frame.port_input.setText(str(phone1_port))
            
            self.log_status(f"✓ Phone 1 port: {phone1_port}")
            self.log_status(f"✓ Phone 2 port: {phone2_port}")
            
            # Connect phones
            self.connect_phone(1)
            time.sleep(1)
            self.connect_phone(2)
            time.sleep(2)  # Give more time for connections to establish
            
            # Test 2: Key exchange with AES
            self.log_status("\n[TEST 2] Testing AES-256-GCM encryption...")
            self.phone1_frame.aes_radio.setChecked(True)
            self.phone1_frame.chacha_radio.setChecked(False)
            
            # Only phone 1 initiates key exchange to avoid race condition
            self.start_key_exchange(1)
            
            # Wait for key exchange with proper timeout
            timeout = 10
            start_time = time.time()
            while (not self.phone1_protocol.protocol.protocol.state.get("key_exchange_complete") and 
                   time.time() - start_time < timeout):
                time.sleep(0.2)
            
            if self.phone1_protocol.protocol.protocol.state.get("key_exchange_complete"):
                self.log_status("✓ AES key exchange successful")
                time.sleep(1)  # Let the key exchange settle
                
                # Send test message
                test_msg = "Test message with AES encryption"
                self.phone1_frame.msg_input.setText(test_msg)
                self.send_message(1)
                self.log_status(f"✓ Sent encrypted message: {test_msg}")
                time.sleep(2)  # Wait for message to be received
            else:
                self.log_status("❌ AES key exchange failed")
            
            # Test 3: Test ChaCha20 (skip reset to avoid segfault)
            self.log_status("\n[TEST 3] Testing ChaCha20-Poly1305 encryption...")
            self.log_status("Note: Using same connection with different cipher")
            
            # Set ChaCha20
            self.phone1_frame.aes_radio.setChecked(False)
            self.phone1_frame.chacha_radio.setChecked(True)
            
            # Only phone 1 initiates key exchange
            self.start_key_exchange(1)
            
            # Wait for key exchange with proper timeout
            timeout = 10
            start_time = time.time()
            while (not self.phone1_protocol.protocol.protocol.state.get("key_exchange_complete") and 
                   time.time() - start_time < timeout):
                time.sleep(0.2)
            
            if self.phone1_protocol.protocol.protocol.state.get("key_exchange_complete"):
                self.log_status("✓ ChaCha20 key exchange successful")
                time.sleep(1)  # Let the key exchange settle
                
                # Send test message
                test_msg = "Test message with ChaCha20 encryption"
                self.phone1_frame.msg_input.setText(test_msg)
                self.send_message(1)
                self.log_status(f"✓ Sent encrypted message: {test_msg}")
                time.sleep(2)  # Wait for message to be received
                
                # Test 4: Voice transmission
                self.log_status("\n[TEST 4] Testing voice transmission...")
                
                # Check if input.wav exists
                audio_file = Path(__file__).parent.parent / "input.wav"
                if audio_file.exists():
                    self.test_voice_transmission()
                    self.log_status("✓ Voice transmission initiated")
                else:
                    self.log_status("❌ input.wav not found, skipping voice test")
            else:
                self.log_status("❌ ChaCha20 key exchange failed")
            
            # Summary
            self.log_status("\n" + "="*50)
            self.log_status("Automated Test Sequence Completed")
            self.log_status("✓ Auto-connection successful")
            self.log_status("✓ Encryption tests completed")
            self.log_status("✓ Message transmission tested")
            if (Path(__file__).parent.parent / "input.wav").exists():
                self.log_status("✓ Voice transmission tested")
            self.log_status("="*50)
            
        except Exception as e:
            self.log_status(f"\n❌ Auto test error: {str(e)}")
            import traceback
            self.log_status(traceback.format_exc())
        finally:
            # Re-enable auto test button
            self.auto_test_btn.setEnabled(True)
        
    def update_phone_status(self, phone_id, message):
        """Update phone status display"""
        self.log_status(f"Phone {phone_id}: {message}")
        
    def log_status(self, message):
        """Log status message"""
        timestamp = time.strftime("%H:%M:%S")
        self.status_text.append(f"[{timestamp}] {message}")
        
    def closeEvent(self, event):
        """Clean up on close"""
        if self.phone1_protocol:
            self.phone1_protocol.stop()
        if self.phone2_protocol:
            self.phone2_protocol.stop()
            
        if hasattr(self, 'gsm_process'):
            self.gsm_process.terminate()
            
        event.accept()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = IntegratedPhoneUI()
    window.show()
    sys.exit(app.exec_())