import socket
import time
import select
import struct
import array
from PyQt5.QtCore import QThread, pyqtSignal
from protocol_client_state import ProtocolClientState
from session import NoiseXKSession
from noise_wrapper import NoiseXKWrapper
from dissononce.dh.keypair import KeyPair
from dissononce.dh.x25519.public import PublicKey
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from voice_codec import Codec2Wrapper, FSKModem, Codec2Mode
# ChaCha20 removed - using only Noise XK encryption

class ProtocolPhoneClient(QThread):
    """Integrated phone client with Noise XK, Codec2, 4FSK, and ChaCha20"""
    data_received = pyqtSignal(bytes, int)
    state_changed = pyqtSignal(str, str, int)
    
    def __init__(self, client_id):
        super().__init__()
        self.host = "localhost"
        self.port = 12345
        self.client_id = client_id
        self.sock = None
        self.running = True
        self.state = ProtocolClientState(client_id)
        
        # Noise XK session
        self.noise_session = None
        self.noise_wrapper = None
        self.handshake_complete = False
        self.handshake_initiated = False
        
        # No buffer needed with larger frame size
        
        # Voice codec components
        self.codec = Codec2Wrapper(mode=Codec2Mode.MODE_1200)
        self.modem = FSKModem()
        
        # Voice encryption handled by Noise XK
        # No separate voice key needed
        
        # Voice state
        self.voice_active = False
        self.voice_frame_counter = 0
        
        # Message buffer for fragmented messages
        self.recv_buffer = bytearray()
        
        # Debug callback
        self.debug_callback = None
    
    def set_debug_callback(self, callback):
        """Set debug callback function"""
        self.debug_callback = callback
        self.state.debug_callback = callback
    
    def debug(self, message):
        """Send debug message"""
        if self.debug_callback:
            self.debug_callback(f"[Phone{self.client_id+1}] {message}")
        else:
            print(f"[Phone{self.client_id+1}] {message}")
        
    def connect_socket(self):
        retries = 3
        for attempt in range(retries):
            try:
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
                self.sock.settimeout(120)
                self.sock.connect((self.host, self.port))
                self.debug(f"Connected to GSM simulator at {self.host}:{self.port}")
                return True
            except Exception as e:
                self.debug(f"Connection attempt {attempt + 1} failed: {e}")
                if attempt < retries - 1:
                    time.sleep(1)
                self.sock = None
        return False
    
    def run(self):
        while self.running:
            if not self.sock:
                if not self.connect_socket():
                    self.debug("Failed to connect after retries")
                    self.state_changed.emit("CALL_END", "", self.client_id)
                    break
            try:
                while self.running:
                    self.state.process_command(self)
                    self.state.check_handshake_timeout(self)
                    
                    if self.handshake_complete and self.voice_active:
                        # Process voice data if active
                        self._process_voice_data()
                    
                    # Always check for incoming data, even during handshake
                    if self.sock is None:
                        break
                    readable, _, _ = select.select([self.sock], [], [], 0.01)
                    if readable:
                        try:
                            if self.sock is None:
                                break
                            chunk = self.sock.recv(4096)
                            if not chunk:
                                self.debug("Disconnected from server")
                                self.state_changed.emit("CALL_END", "", self.client_id)
                                break
                            
                            # Add to buffer
                            self.recv_buffer.extend(chunk)
                            
                            # Process complete messages
                            while len(self.recv_buffer) >= 4:
                                # Read message length
                                msg_len = struct.unpack('>I', self.recv_buffer[:4])[0]
                                
                                # Check if we have the complete message
                                if len(self.recv_buffer) >= 4 + msg_len:
                                    # Extract message
                                    data = bytes(self.recv_buffer[4:4+msg_len])
                                    # Remove from buffer
                                    self.recv_buffer = self.recv_buffer[4+msg_len:]
                                    # Pass to state handler
                                    self.state.handle_data(self, data)
                                else:
                                    # Wait for more data
                                    break
                                    
                        except socket.error as e:
                            self.debug(f"Socket error: {e}")
                            self.state_changed.emit("CALL_END", "", self.client_id)
                            break
                    
                    self.msleep(1)
            except Exception as e:
                self.debug(f"Unexpected error in run loop: {e}")
                self.state_changed.emit("CALL_END", "", self.client_id)
                break
            finally:
                if self.sock:
                    try:
                        self.sock.close()
                    except Exception as e:
                        self.debug(f"Error closing socket: {e}")
                    self.sock = None
    
    def _handle_encrypted_data(self, data):
        """Handle encrypted data after handshake"""
        if not self.handshake_complete or not self.noise_wrapper:
            self.debug(f"Cannot decrypt - handshake not complete")
            return
            
        # All data after handshake is encrypted, decrypt it first
        try:
            plaintext = self.noise_wrapper.decrypt(data)
            
            # Check if it's a text message
            try:
                text_msg = plaintext.decode('utf-8').strip()
                if text_msg == "HANDSHAKE_DONE":
                    self.debug(f"Received encrypted HANDSHAKE_DONE")
                    self.state_changed.emit("HANDSHAKE_DONE", "HANDSHAKE_DONE", self.client_id)
                    return
            except:
                pass
                
            # Otherwise handle as protocol message
            self._handle_protocol_message(plaintext)
        except Exception as e:
            # Suppress common decryption errors
            pass
    
    def _handle_protocol_message(self, plaintext):
        """Handle decrypted protocol messages"""
        if len(plaintext) < 1:
            return
            
        msg_type = plaintext[0]
        msg_data = plaintext[1:]
        
        if msg_type == 0x10:  # Voice start
            self.debug("Received VOICE_START message")
            self._handle_voice_start(msg_data)
        elif msg_type == 0x11:  # Voice data
            self._handle_voice_data(msg_data)
        elif msg_type == 0x12:  # Voice end
            self.debug("Received VOICE_END message")
            self._handle_voice_end(msg_data)
        elif msg_type == 0x20:  # Noise handshake
            self.debug("Received NOISE_HS message")
            self._handle_noise_handshake(msg_data)
        else:
            self.debug(f"Received unknown protocol message type: 0x{msg_type:02x}")
            # Don't emit control messages to data_received - that's only for audio
            # Control messages should be handled via state_changed signal
    
    def _handle_voice_start(self, data):
        """Handle voice session start"""
        self.debug("Voice session started by peer")
        self.voice_active = True
        self.voice_frame_counter = 0
        self.state_changed.emit("VOICE_START", "", self.client_id)
    
    def _handle_voice_data(self, data):
        """Handle voice frame (already decrypted by Noise)"""
        if len(data) < 4:
            return
            
        try:
            # Data is float array packed as bytes
            # Unpack the float array
            num_floats = len(data) // 4
            modulated_signal = struct.unpack(f'{num_floats}f', data)
            
            # Demodulate FSK
            demodulated_data, confidence = self.modem.demodulate(modulated_signal)
            
            if confidence > 0.5:  # Only decode if confidence is good
                # Create Codec2Frame from demodulated data
                from voice_codec import Codec2Frame, Codec2Mode
                frame = Codec2Frame(
                    mode=Codec2Mode.MODE_1200,
                    bits=demodulated_data,
                    timestamp=time.time(),
                    frame_number=self.voice_frame_counter
                )
                
                # Decode with Codec2
                pcm_samples = self.codec.decode(frame)
                
                if self.voice_frame_counter == 0:
                    self.debug(f"First voice frame demodulated with confidence {confidence:.2f}")
                
                # Send PCM to UI for playback
                if pcm_samples is not None and len(pcm_samples) > 0:
                    # Only log details for first frame and every 25th frame
                    if self.voice_frame_counter == 0 or self.voice_frame_counter % 25 == 0:
                        self.debug(f"Decoded PCM samples: type={type(pcm_samples)}, len={len(pcm_samples)}")
                    
                    # Convert to bytes if needed
                    if hasattr(pcm_samples, 'tobytes'):
                        pcm_bytes = pcm_samples.tobytes()
                    elif isinstance(pcm_samples, (list, array.array)):
                        # Convert array to bytes
                        import array
                        if isinstance(pcm_samples, list):
                            pcm_array = array.array('h', pcm_samples)
                            pcm_bytes = pcm_array.tobytes()
                        else:
                            pcm_bytes = pcm_samples.tobytes()
                    else:
                        pcm_bytes = bytes(pcm_samples)
                    
                    if self.voice_frame_counter == 0:
                        self.debug(f"Emitting first PCM frame: {len(pcm_bytes)} bytes")
                    
                    self.data_received.emit(pcm_bytes, self.client_id)
                    self.voice_frame_counter += 1
                    # Log frame reception periodically
                    if self.voice_frame_counter == 1 or self.voice_frame_counter % 25 == 0:
                        self.debug(f"Received voice data frame #{self.voice_frame_counter}")
                else:
                    self.debug(f"Codec decode returned None or empty")
            else:
                if self.voice_frame_counter % 10 == 0:
                    self.debug(f"Low confidence demodulation: {confidence:.2f}")
                    
        except Exception as e:
            self.debug(f"Voice decode error: {e}")
    
    def _handle_voice_end(self, data):
        """Handle voice session end"""
        self.debug("Voice session ended by peer")
        self.voice_active = False
        self.state_changed.emit("VOICE_END", "", self.client_id)
    
    def _handle_noise_handshake(self, data):
        """Handle Noise handshake message"""
        if not self.noise_wrapper:
            self.debug("Received handshake message but no wrapper initialized")
            return
            
        try:
            # Process the handshake message
            self.noise_wrapper.process_handshake_message(data)
            
            # Check if we need to send a response
            response = self.noise_wrapper.get_next_handshake_message()
            if response:
                self.send(b'\x20' + response)
                
            # Check if handshake is complete
            if self.noise_wrapper.handshake_complete and not self.handshake_complete:
                self.debug("Noise wrapper handshake complete, calling complete_handshake()")
                self.complete_handshake()
                
        except Exception as e:
            self.debug(f"Handshake processing error: {e}")
            self.state_changed.emit("CALL_END", "", self.client_id)
    
    def _process_voice_data(self):
        """Process outgoing voice data"""
        # This would be called when we have voice input to send
        # For now, this is a placeholder
        pass
    
    def send_voice_frame(self, pcm_samples):
        """Send a voice frame through the protocol"""
        if not self.handshake_complete:
            self.debug("Cannot send voice - handshake not complete")
            return
        if not self.voice_active:
            self.debug("Cannot send voice - voice session not active")
            return
            
        try:
            # Encode with Codec2
            codec_frame = self.codec.encode(pcm_samples)
            if not codec_frame:
                return
            
            if self.voice_frame_counter % 25 == 0:  # Log every 25 frames (1 second)
                self.debug(f"Encoding voice frame #{self.voice_frame_counter}: {len(pcm_samples)} bytes PCM → {len(codec_frame.bits)} bytes compressed")
                
            # Modulate with FSK
            modulated_data = self.modem.modulate(codec_frame.bits)
            
            # Convert modulated float array to bytes
            modulated_bytes = struct.pack(f'{len(modulated_data)}f', *modulated_data)
            
            if self.voice_frame_counter % 25 == 0:
                self.debug(f"Voice frame size: {len(modulated_bytes)} bytes")
            
            # Build voice data message (no ChaCha20, will be encrypted by Noise)
            msg = bytes([0x11]) + modulated_bytes
            
            # Send through Noise encrypted channel
            self.send(msg)
            
            self.voice_frame_counter += 1
            
        except Exception as e:
            self.debug(f"Voice encode error: {e}")
    
    def send(self, message):
        """Send data through Noise encrypted channel with proper framing"""
        if self.sock and self.running:
            try:
                # Handshake messages (0x20) bypass Noise encryption
                if isinstance(message, bytes) and len(message) > 0 and message[0] == 0x20:
                    # Add length prefix for framing
                    framed = struct.pack('>I', len(message)) + message
                    self.sock.send(framed)
                    return
                
                if self.handshake_complete and self.noise_wrapper:
                    # Encrypt everything with Noise after handshake
                    # Convert string to bytes if needed
                    if isinstance(message, str):
                        message = message.encode('utf-8')
                    encrypted = self.noise_wrapper.encrypt(message)
                    # Add length prefix for framing
                    framed = struct.pack('>I', len(encrypted)) + encrypted
                    self.sock.send(framed)
                else:
                    # During handshake, send raw with framing
                    if isinstance(message, str):
                        data = message.encode('utf-8')
                        framed = struct.pack('>I', len(data)) + data
                        self.sock.send(framed)
                        self.debug(f"Sent control message: {message}")
                    else:
                        framed = struct.pack('>I', len(message)) + message
                        self.sock.send(framed)
            except socket.error as e:
                self.debug(f"Send error: {e}")
                self.state_changed.emit("CALL_END", "", self.client_id)
    
    def stop(self):
        self.running = False
        self.voice_active = False
        if self.sock:
            try:
                self.sock.close()
            except Exception as e:
                self.debug(f"Error closing socket in stop: {e}")
            self.sock = None
        self.quit()
        self.wait(1000)
    
    def start_handshake(self, initiator, keypair, peer_pubkey):
        """Start Noise XK handshake"""
        self.debug(f"Starting Noise XK handshake as {'initiator' if initiator else 'responder'}")
        self.debug(f"Our public key: {keypair.public.data.hex()[:32]}...")
        self.debug(f"Peer public key: {peer_pubkey.data.hex()[:32]}...")
        
        # Create noise wrapper
        self.noise_wrapper = NoiseXKWrapper(keypair, peer_pubkey, self.debug)
        self.noise_wrapper.start_handshake(initiator)
        self.handshake_initiated = True
        
        # Send first handshake message if initiator
        if initiator:
            msg = self.noise_wrapper.get_next_handshake_message()
            if msg:
                # Send as NOISE_HS message type
                self.send(b'\x20' + msg)  # 0x20 = Noise handshake message
    
    def complete_handshake(self):
        """Called when Noise handshake completes"""
        self.handshake_complete = True
        
        self.debug("Noise XK handshake complete!")
        self.debug("Secure channel established")
        
        # Send HANDSHAKE_DONE message
        self.send("HANDSHAKE_DONE")
        
        self.state_changed.emit("HANDSHAKE_COMPLETE", "", self.client_id)
    
    def start_voice_session(self):
        """Start a voice session"""
        if not self.handshake_complete:
            self.debug("Cannot start voice - handshake not complete")
            return
            
        self.voice_active = True
        self.voice_frame_counter = 0
        
        # Send voice start message
        msg = bytes([0x10])  # Voice start message type
        self.send(msg)
        
        self.debug("Voice session started")
        self.state_changed.emit("VOICE_START", "", self.client_id)
    
    def end_voice_session(self):
        """End a voice session"""
        if not self.voice_active:
            return
            
        self.voice_active = False
        
        # Send voice end message
        msg = bytes([0x12])  # Voice end message type
        self.send(msg)
        
        self.debug("Voice session ended")
        self.state_changed.emit("VOICE_END", "", self.client_id)