import sys
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QFrame, QSizePolicy, QStyle
)
from PyQt5.QtCore import Qt, QSize
from PyQt5.QtGui import QFont
from phone_manager import PhoneManager
from waveform_widget import WaveformWidget
from phone_state import PhoneState

class PhoneUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Enhanced Dual Phone Interface")
        self.setGeometry(100, 100, 900, 750)
        self.setStyleSheet("""
            QMainWindow { background-color: #333333; }
            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#settingsButton { background-color: #555555; }
            QPushButton#settingsButton:hover { background-color: #777777; }
            QFrame#phoneDisplay {
                background-color: #1E1E1E; border: 2px solid #0078D4;
                border-radius: 10px;
            }
            QLabel#phoneTitleLabel {
                font-size: 18px; font-weight: bold; padding-bottom: 5px;
                color: #FFFFFF;
            }
            QLabel#mainTitleLabel {
                font-size: 24px; font-weight: bold; color: #00A2E8;
                padding: 15px;
            }
            QWidget#phoneWidget {
                border: 1px solid #4A4A4A; border-radius: 8px;
                padding: 10px; background-color: #3A3A3A;
            }
        """)

        self.manager = PhoneManager()
        self.manager.initialize_phones()

        # Main widget and layout
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QVBoxLayout()
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setAlignment(Qt.AlignCenter)
        main_widget.setLayout(main_layout)

        # App Title
        app_title_label = QLabel("Dual Phone Control Panel")
        app_title_label.setObjectName("mainTitleLabel")
        app_title_label.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(app_title_label)

        # Phone displays layout
        phone_controls_layout = QHBoxLayout()
        phone_controls_layout.setSpacing(50)
        phone_controls_layout.setAlignment(Qt.AlignCenter)
        main_layout.addLayout(phone_controls_layout)

        # Setup UI for phones
        for phone in self.manager.phones:
            phone_container_widget, phone_display_frame, phone_button, waveform_widget, sent_waveform_widget, phone_status_label = self._create_phone_ui(
                f"Phone {phone['id']+1}", lambda checked, pid=phone['id']: self.manager.phone_action(pid, self)
            )
            phone['button'] = phone_button
            phone['waveform'] = waveform_widget
            phone['sent_waveform'] = sent_waveform_widget
            phone['status_label'] = phone_status_label
            phone_controls_layout.addWidget(phone_container_widget)
            phone['client'].data_received.connect(lambda data, cid=phone['id']: self.manager.update_waveform(cid, data))
            phone['client'].state_changed.connect(lambda state, num, cid=phone['id']: self.set_phone_state(cid, state, num))
            phone['client'].start()

        # Spacer
        main_layout.addStretch(1)

        # Settings Button
        self.settings_button = QPushButton("Settings")
        self.settings_button.setObjectName("settingsButton")
        self.settings_button.setFixedWidth(180)
        self.settings_button.setIcon(self.style().standardIcon(QStyle.SP_FileDialogDetailedView))
        self.settings_button.setIconSize(QSize(20, 20))
        self.settings_button.clicked.connect(self.settings_action)
        settings_layout = QHBoxLayout()
        settings_layout.addStretch()
        settings_layout.addWidget(self.settings_button)
        settings_layout.addStretch()
        main_layout.addLayout(settings_layout)

        # Initialize UI
        for phone in self.manager.phones:
            self.update_phone_ui(phone['id'])

    def _create_phone_ui(self, title, action_slot):
        phone_container_widget = QWidget()
        phone_container_widget.setObjectName("phoneWidget")
        phone_layout = QVBoxLayout()
        phone_layout.setAlignment(Qt.AlignCenter)
        phone_layout.setSpacing(15)
        phone_container_widget.setLayout(phone_layout)

        phone_title_label = QLabel(title)
        phone_title_label.setObjectName("phoneTitleLabel")
        phone_title_label.setAlignment(Qt.AlignCenter)
        phone_layout.addWidget(phone_title_label)

        phone_display_frame = QFrame()
        phone_display_frame.setObjectName("phoneDisplay")
        phone_display_frame.setFixedSize(250, 350)
        phone_display_frame.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)

        display_content_layout = QVBoxLayout(phone_display_frame)
        display_content_layout.setAlignment(Qt.AlignCenter)
        phone_status_label = QLabel("Idle")
        phone_status_label.setAlignment(Qt.AlignCenter)
        phone_status_label.setFont(QFont("Arial", 16))
        display_content_layout.addWidget(phone_status_label)
        phone_layout.addWidget(phone_display_frame, alignment=Qt.AlignCenter)

        phone_button = QPushButton()
        phone_button.setFixedWidth(120)
        phone_button.setIconSize(QSize(20, 20))
        phone_button.clicked.connect(action_slot)
        phone_layout.addWidget(phone_button, alignment=Qt.AlignCenter)

        # Received waveform
        waveform_label = QLabel(f"{title} Received Audio")
        waveform_label.setAlignment(Qt.AlignCenter)
        waveform_label.setStyleSheet("font-size: 14px; color: #E0E0E0;")
        phone_layout.addWidget(waveform_label)
        waveform_widget = WaveformWidget(dynamic=False)
        phone_layout.addWidget(waveform_widget, alignment=Qt.AlignCenter)

        # Sent waveform
        sent_waveform_label = QLabel(f"{title} Sent Audio")
        sent_waveform_label.setAlignment(Qt.AlignCenter)
        sent_waveform_label.setStyleSheet("font-size: 14px; color: #E0E0E0;")
        phone_layout.addWidget(sent_waveform_label)
        sent_waveform_widget = WaveformWidget(dynamic=False)
        phone_layout.addWidget(sent_waveform_widget, alignment=Qt.AlignCenter)

        return phone_container_widget, phone_display_frame, phone_button, waveform_widget, sent_waveform_widget, phone_status_label

    def update_phone_ui(self, phone_id):
        phone = self.manager.phones[phone_id]
        other_phone = self.manager.phones[1 - phone_id]
        state = phone['state']
        phone_number = other_phone['number'] if state != PhoneState.IDLE else ""
        button = phone['button']
        status_label = phone['status_label']

        if state == PhoneState.IDLE:
            button.setText("Call")
            button.setIcon(self.style().standardIcon(QStyle.SP_MediaPlay))
            status_label.setText("Idle")
            button.setStyleSheet("background-color: #0078D4;")
        elif state == PhoneState.CALLING:
            button.setText("Cancel")
            button.setIcon(self.style().standardIcon(QStyle.SP_MediaStop))
            status_label.setText(f"Calling {phone_number}...")
            button.setStyleSheet("background-color: #E81123;")
        elif state == PhoneState.IN_CALL:
            button.setText("Hang Up")
            button.setIcon(self.style().standardIcon(QStyle.SP_DialogCancelButton))
            status_label.setText(f"In Call with {phone_number}")
            button.setStyleSheet("background-color: #E81123;")
        elif state == PhoneState.RINGING:
            button.setText("Answer")
            button.setIcon(self.style().standardIcon(QStyle.SP_DialogApplyButton))
            status_label.setText(f"Incoming Call from {phone_number}")
            button.setStyleSheet("background-color: #107C10;")

    def set_phone_state(self, client_id, state_str, number):
        state = self.manager.map_state(state_str)
        phone = self.manager.phones[client_id]
        other_phone = self.manager.phones[1 - client_id]
        print(f"Setting state for Phone {client_id + 1}: {state}, number: {number}, is_initiator: {phone['is_initiator']}")
        phone['state'] = state
        if state == PhoneState.IN_CALL:
            print(f"Phone {client_id + 1} confirmed in IN_CALL state")
            if number == "IN_CALL" and phone['is_initiator']:
                print(f"Phone {client_id + 1} (initiator) starting handshake")
                phone['client'].send("HANDSHAKE")
                phone['client'].start_handshake(initiator=True, keypair=phone['keypair'], peer_pubkey=other_phone['public_key'])
            elif number == "HANDSHAKE" and not phone['is_initiator']:
                print(f"Phone {client_id + 1} (responder) starting handshake")
                phone['client'].start_handshake(initiator=False, keypair=phone['keypair'], peer_pubkey=other_phone['public_key'])
            elif number == "HANDSHAKE_DONE":
                self.manager.start_audio(client_id, parent=self)  # Pass self as parent
        self.update_phone_ui(client_id)

    def settings_action(self):
        print("Settings clicked")

    def closeEvent(self, event):
        for phone in self.manager.phones:
            phone['client'].stop()
        event.accept()

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