import secrets
from PyQt5.QtCore import QTimer
from protocol_phone_client import ProtocolPhoneClient
from session import NoiseXKSession
from phone_state import PhoneState  # Added import
from audio_player import AudioPlayer
from audio_processor import AudioProcessor
import struct
import wave
import os

class PhoneManager:
    def __init__(self):
        self.phones = []
        self.handshake_done_count = 0
        self.ui = None  # Will be set by UI
        self.audio_player = AudioPlayer()
        self.audio_player.set_debug_callback(self.debug)
        self.audio_processor = AudioProcessor()
        self.audio_processor.set_debug_callback(self.debug)
        self.audio_buffer = {}  # client_id -> list of audio chunks for processing

    def debug(self, message):
        """Send debug message to UI if available"""
        if self.ui and hasattr(self.ui, 'debug'):
            self.ui.debug(f"[PhoneManager] {message}")
        else:
            print(f"[PhoneManager] {message}")
    
    def initialize_phones(self):
        for i in range(2):
            client = ProtocolPhoneClient(i)  # Use protocol client
            client.set_debug_callback(self.debug)  # Set debug callback
            client.manager = self  # Set manager reference for handshake lookup
            keypair = NoiseXKSession.generate_keypair()
            phone = {
                'id': i,
                'client': client,
                'state': PhoneState.IDLE,
                'number': "123-4567" if i == 0 else "987-6543",
                'audio_timer': None,
                'keypair': keypair,
                'public_key': keypair.public,
                'is_initiator': False,
                'audio_file': None,  # For test audio
                'frame_counter': 0,
                'playback_enabled': False,
                'recording_enabled': False
            }
            client.keypair = keypair  # Also set keypair on client
            self.phones.append(phone)
            self.debug(f"Initialized Phone {i+1} with public key: {keypair.public.data.hex()[:32]}...")

        self.phones[0]['peer_public_key'] = self.phones[1]['public_key']
        self.phones[1]['peer_public_key'] = self.phones[0]['public_key']

    def phone_action(self, phone_id, ui_manager):
        phone = self.phones[phone_id]
        other_phone = self.phones[1 - phone_id]
        self.debug(f"Phone {phone_id + 1} action triggered, current state: {phone['state'].name}")

        if phone['state'] == PhoneState.IDLE:
            self.debug(f"Phone {phone_id + 1} initiating call to Phone {2-phone_id}")
            phone['state'] = PhoneState.CALLING
            other_phone['state'] = PhoneState.RINGING
            phone['is_initiator'] = True
            other_phone['is_initiator'] = False
            phone['client'].send("RINGING")
        elif phone['state'] == PhoneState.RINGING:
            self.debug(f"Phone {phone_id + 1} answering call from Phone {2-phone_id}")
            phone['state'] = PhoneState.IN_CALL
            # Don't set other_phone state here - let it set when it receives IN_CALL
            phone['client'].send("IN_CALL")
        elif phone['state'] in [PhoneState.IN_CALL, PhoneState.CALLING]:
            if not phone['client'].state.handshake_in_progress and phone['state'] != PhoneState.CALLING:
                phone['state'] = other_phone['state'] = PhoneState.IDLE
                phone['client'].send("CALL_END")
                for p in [phone, other_phone]:
                    if p['audio_timer']:
                        p['audio_timer'].stop()
                    # End voice session
                    if p['client'].voice_active:
                        p['client'].end_voice_session()
                    # Close audio file
                    if p['audio_file']:
                        p['audio_file'].close()
                        p['audio_file'] = None
                    p['frame_counter'] = 0
            else:
                self.debug(f"Phone {phone_id + 1} cannot hang up during handshake or call setup")
        
        ui_manager.update_phone_ui(phone_id)
        ui_manager.update_phone_ui(1 - phone_id)

    def send_audio(self, phone_id):
        phone = self.phones[phone_id]
        if phone['state'] != PhoneState.IN_CALL:
            self.debug(f"Phone {phone_id + 1} not in call, stopping audio timer")
            if phone['audio_timer']:
                phone['audio_timer'].stop()
            return
            
        if not phone['client'].handshake_complete:
            self.debug(f"Phone {phone_id + 1} handshake not complete, skipping audio send")
            return
            
        if not phone['client'].voice_active:
            self.debug(f"Phone {phone_id + 1} voice not active, skipping audio send")
            return
            
        if phone['state'] == PhoneState.IN_CALL and phone['client'].handshake_complete and phone['client'].voice_active:
            # Load test audio file if not loaded
            if phone['audio_file'] is None:
                wav_path = "../wav/input.wav"
                if not os.path.exists(wav_path):
                    wav_path = "wav/input.wav"
                if os.path.exists(wav_path):
                    try:
                        phone['audio_file'] = wave.open(wav_path, 'rb')
                        self.debug(f"Phone {phone_id + 1} loaded test audio file: {wav_path}")
                        # Verify it's 8kHz mono
                        if phone['audio_file'].getframerate() != 8000:
                            self.debug(f"Warning: {wav_path} is {phone['audio_file'].getframerate()}Hz, expected 8000Hz")
                        if phone['audio_file'].getnchannels() != 1:
                            self.debug(f"Warning: {wav_path} has {phone['audio_file'].getnchannels()} channels, expected 1")
                        
                        # Skip initial silence - jump to 1 second in (8000 samples)
                        phone['audio_file'].setpos(8000)
                        self.debug(f"Phone {phone_id + 1} skipped initial silence, starting at 1 second")
                    except Exception as e:
                        self.debug(f"Phone {phone_id + 1} failed to load audio: {e}")
                        # Use mock audio as fallback
                        phone['audio_file'] = None
            
            # Read audio frame (40ms at 8kHz = 320 samples)
            if phone['audio_file']:
                try:
                    frames = phone['audio_file'].readframes(320)
                    if not frames or len(frames) < 640:  # 320 samples * 2 bytes
                        # Loop back to 1 second (skip silence)
                        phone['audio_file'].setpos(8000)
                        frames = phone['audio_file'].readframes(320)
                        self.debug(f"Phone {phone_id + 1} looped audio back to 1 second mark")
                    
                    # Send through protocol (codec + 4FSK + encryption)
                    phone['client'].send_voice_frame(frames)
                    
                    # Update waveform
                    if len(frames) >= 2:
                        samples = struct.unpack(f'{len(frames)//2}h', frames)
                        self.update_sent_waveform(phone_id, frames)
                    
                    # If playback is enabled on the sender, play the original audio
                    if phone['playback_enabled']:
                        self.audio_player.add_audio_data(phone_id, frames)
                        if phone['frame_counter'] % 25 == 0:
                            self.debug(f"Phone {phone_id + 1} playing original audio (sender playback)")
                    
                    phone['frame_counter'] += 1
                    if phone['frame_counter'] % 25 == 0:  # Log every second
                        self.debug(f"Phone {phone_id + 1} sent {phone['frame_counter']} voice frames")
                        
                except Exception as e:
                    self.debug(f"Phone {phone_id + 1} audio send error: {e}")
            else:
                # Fallback: send mock audio
                mock_audio = secrets.token_bytes(320)
                phone['client'].send_voice_frame(mock_audio)
                self.update_sent_waveform(phone_id, mock_audio)

    def start_audio(self, client_id, parent=None):
        self.handshake_done_count += 1
        self.debug(f"HANDSHAKE_DONE received for client {client_id}, count: {self.handshake_done_count}")
        
        # Start voice session for this client
        phone = self.phones[client_id]
        if phone['client'].handshake_complete and not phone['client'].voice_active:
            phone['client'].start_voice_session()
        
        if self.handshake_done_count == 2:
            # Add a small delay to ensure both sides are ready
            def start_audio_timers():
                self.debug("Starting audio timers for both phones")
                for phone in self.phones:
                    if phone['state'] == PhoneState.IN_CALL:
                        if not phone['audio_timer'] or not phone['audio_timer'].isActive():
                            phone['audio_timer'] = QTimer(parent)  # Parent to PhoneUI
                            phone['audio_timer'].timeout.connect(lambda pid=phone['id']: self.send_audio(pid))
                            phone['audio_timer'].start(40)  # 40ms for each voice frame
            
            # Delay audio start by 500ms to ensure both sides are ready
            QTimer.singleShot(500, start_audio_timers)
            self.handshake_done_count = 0

    def update_waveform(self, client_id, data):
        # Only process actual audio data (should be 640 bytes for 320 samples * 2 bytes)
        # Ignore small control messages
        if len(data) < 320:  # Less than 160 samples (too small for audio)
            self.debug(f"Phone {client_id + 1} received non-audio data: {len(data)} bytes (ignoring)")
            return
            
        self.phones[client_id]['waveform'].set_data(data)
        
        # Debug log audio data reception (only occasionally to avoid spam)
        if not hasattr(self, '_audio_frame_count'):
            self._audio_frame_count = {}
        if client_id not in self._audio_frame_count:
            self._audio_frame_count[client_id] = 0
        self._audio_frame_count[client_id] += 1
        
        if self._audio_frame_count[client_id] == 1 or self._audio_frame_count[client_id] % 25 == 0:
            self.debug(f"Phone {client_id + 1} received audio frame #{self._audio_frame_count[client_id]}: {len(data)} bytes")
        
        # Store audio data in buffer for potential processing
        if client_id not in self.audio_buffer:
            self.audio_buffer[client_id] = []
        self.audio_buffer[client_id].append(data)
        
        # Keep buffer size reasonable (last 30 seconds at 8kHz)
        max_chunks = 30 * 25  # 30 seconds * 25 chunks/second
        if len(self.audio_buffer[client_id]) > max_chunks:
            self.audio_buffer[client_id] = self.audio_buffer[client_id][-max_chunks:]
        
        # Forward audio data to player if playback is enabled
        if self.phones[client_id]['playback_enabled']:
            if self._audio_frame_count[client_id] == 1:
                self.debug(f"Phone {client_id + 1} forwarding audio to player (playback enabled)")
            self.audio_player.add_audio_data(client_id, data)

    def update_sent_waveform(self, client_id, data):
        self.phones[client_id]['sent_waveform'].set_data(data)
    
    def toggle_playback(self, client_id):
        """Toggle audio playback for a phone"""
        phone = self.phones[client_id]
        
        if phone['playback_enabled']:
            # Stop playback
            self.audio_player.stop_playback(client_id)
            phone['playback_enabled'] = False
            self.debug(f"Phone {client_id + 1} playback stopped")
        else:
            # Start playback
            if self.audio_player.start_playback(client_id):
                phone['playback_enabled'] = True
                self.debug(f"Phone {client_id + 1} playback started")
                # Removed test beep - we want to hear actual audio
            else:
                self.debug(f"Phone {client_id + 1} failed to start playback")
        
        return phone['playback_enabled']
    
    def toggle_recording(self, client_id):
        """Toggle audio recording for a phone"""
        phone = self.phones[client_id]
        
        if phone['recording_enabled']:
            # Stop recording and save
            save_path = self.audio_player.stop_recording(client_id)
            phone['recording_enabled'] = False
            if save_path:
                self.debug(f"Phone {client_id + 1} recording saved to {save_path}")
            return False, save_path
        else:
            # Start recording
            self.audio_player.start_recording(client_id)
            phone['recording_enabled'] = True
            self.debug(f"Phone {client_id + 1} recording started")
            return True, None
    
    def save_received_audio(self, client_id, filename=None):
        """Save the last received audio to a file"""
        if client_id not in self.phones:
            return None
            
        save_path = self.audio_player.stop_recording(client_id, filename)
        if save_path:
            self.debug(f"Phone {client_id + 1} audio saved to {save_path}")
        return save_path

    def process_audio(self, client_id, processing_type, **kwargs):
        """Process buffered audio with specified processing type"""
        if client_id not in self.audio_buffer or not self.audio_buffer[client_id]:
            self.debug(f"No audio data available for Phone {client_id + 1}")
            return None
        
        # Combine all audio chunks
        combined_audio = b''.join(self.audio_buffer[client_id])
        
        # Apply processing based on type
        processed_audio = combined_audio
        
        if processing_type == "normalize":
            target_db = kwargs.get('target_db', -3)
            processed_audio = self.audio_processor.normalize_audio(combined_audio, target_db)
        
        elif processing_type == "gain":
            gain_db = kwargs.get('gain_db', 0)
            processed_audio = self.audio_processor.apply_gain(combined_audio, gain_db)
        
        elif processing_type == "noise_gate":
            threshold_db = kwargs.get('threshold_db', -40)
            processed_audio = self.audio_processor.apply_noise_gate(combined_audio, threshold_db)
        
        elif processing_type == "low_pass":
            cutoff_hz = kwargs.get('cutoff_hz', 3400)
            processed_audio = self.audio_processor.apply_low_pass_filter(combined_audio, cutoff_hz)
        
        elif processing_type == "high_pass":
            cutoff_hz = kwargs.get('cutoff_hz', 300)
            processed_audio = self.audio_processor.apply_high_pass_filter(combined_audio, cutoff_hz)
        
        elif processing_type == "remove_silence":
            threshold_db = kwargs.get('threshold_db', -40)
            processed_audio = self.audio_processor.remove_silence(combined_audio, threshold_db)
        
        # Save processed audio
        save_path = f"wav/phone{client_id + 1}_received.wav"
        processed_path = self.audio_processor.save_processed_audio(
            processed_audio, save_path, processing_type
        )
        
        return processed_path
    
    def export_buffered_audio(self, client_id, filename=None):
        """Export current audio buffer to file"""
        if client_id not in self.audio_buffer or not self.audio_buffer[client_id]:
            self.debug(f"No audio data available for Phone {client_id + 1}")
            return None
        
        # Combine all audio chunks
        combined_audio = b''.join(self.audio_buffer[client_id])
        
        # Generate filename if not provided
        if not filename:
            from datetime import datetime
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"wav/phone{client_id + 1}_buffer_{timestamp}.wav"
        
        # Ensure directory exists
        os.makedirs(os.path.dirname(filename), exist_ok=True)
        
        try:
            with wave.open(filename, 'wb') as wav_file:
                wav_file.setnchannels(1)
                wav_file.setsampwidth(2)
                wav_file.setframerate(8000)
                wav_file.writeframes(combined_audio)
            
            self.debug(f"Exported audio buffer for Phone {client_id + 1} to {filename}")
            return filename
            
        except Exception as e:
            self.debug(f"Failed to export audio buffer: {e}")
            return None
    
    def clear_audio_buffer(self, client_id):
        """Clear audio buffer for a phone"""
        if client_id in self.audio_buffer:
            self.audio_buffer[client_id] = []
            self.debug(f"Cleared audio buffer for Phone {client_id + 1}")

    def map_state(self, state_str):
        if state_str == "RINGING":
            return PhoneState.RINGING
        elif state_str in ["CALL_END", "CALL_DROPPED"]:
            return PhoneState.IDLE
        elif state_str == "IN_CALL":
            return PhoneState.IN_CALL
        elif state_str == "HANDSHAKE":
            return PhoneState.IN_CALL
        elif state_str == "HANDSHAKE_DONE":
            return PhoneState.IN_CALL
        return PhoneState.IDLE