import secrets
from PyQt5.QtCore import QTimer
from phone_client import PhoneClient
from session import NoiseXKSession
from phone_state import PhoneState  # Added import

class PhoneManager:
    def __init__(self):
        self.phones = []
        self.handshake_done_count = 0

    def initialize_phones(self):
        for i in range(2):
            client = PhoneClient(i)
            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
            }
            self.phones.append(phone)

        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]
        print(f"Phone {phone_id + 1} Action, current state: {phone['state']}, is_initiator: {phone['is_initiator']}")

        if phone['state'] == PhoneState.IDLE:
            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:
            phone['state'] = other_phone['state'] = PhoneState.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()
            else:
                print(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 and phone['client'].state.session and phone['client'].sock:
            mock_audio = secrets.token_bytes(16)
            try:
                self.update_sent_waveform(phone_id, mock_audio)
                phone['client'].state.session.send(phone['client'].sock, mock_audio)
                print(f"Client {phone_id} sent encrypted audio packet, length=32")
            except Exception as e:
                print(f"Client {phone_id} failed to send audio: {e}")

    def start_audio(self, client_id, parent=None):
        self.handshake_done_count += 1
        print(f"HANDSHAKE_DONE received for client {client_id}, count: {self.handshake_done_count}")
        if self.handshake_done_count == 2:
            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(100)
            self.handshake_done_count = 0

    def update_waveform(self, client_id, data):
        self.phones[client_id]['waveform'].set_data(data)

    def update_sent_waveform(self, client_id, data):
        self.phones[client_id]['sent_waveform'].set_data(data)

    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