app.py

import sys
import json
import time
import sounddevice as sd
import numpy as np
from dataclasses import dataclass, asdict
from typing import Optional, List, Dict
from enum import Enum, auto
from pathlib import Path
import ctypes
from ctypes import wintypes
import os
import subprocess
from datetime import datetime
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QComboBox, QPushButton, QPlainTextEdit, QLabel, QTabWidget,
    QGroupBox, QCheckBox, QSpinBox, QMessageBox, QFileDialog,
    QScrollArea, QFrame, QProgressBar, QStyle, QMenu
)
from PyQt5.QtCore import Qt, pyqtSignal, QThread, QTimer
from PyQt5.QtGui import QColor, QPalette, QIcon

AUDIO_STATE_FLAGS = {
    0x00000001: "ACTIVE",
    0x00000002: "DISABLED",
    0x00000004: "NOT_PRESENT",
    0x00000008: "UNPLUGGED",
}

class DeviceType(Enum):
    UNKNOWN = auto()
    MICROPHONE = auto()
    SPEAKER = auto()
    HEADPHONES = auto()
    VIRTUAL = auto()
    LOOPBACK = auto()

class DeviceRole(Enum):
    INPUT = auto()
    OUTPUT = auto()
    BOTH = auto()

@dataclass
class WindowsAudioProperties:
    is_default: bool = False
    is_default_communication: bool = False
    listen_to_device: bool = False
    playback_through: Optional[str] = None
    device_state: str = "Unknown"
    friendly_name: Optional[str] = None
    device_id: Optional[str] = None
    state_flags: int = 0

@dataclass
class AudioDeviceMetadata:
    id: int
    name: str
    type: DeviceType
    role: DeviceRole
    max_input: int
    max_output: int
    sample_rates: List[int]
    default_sample_rate: int
    low_latency: Dict
    windows_props: WindowsAudioProperties
    driver_version: Optional[str] = None
    hardware_id: Optional[str] = None

    def to_dict(self):
        return {
            "id": self.id,
            "name": self.name,
            "type": self.type.name,
            "role": self.role.name,
            "max_input": self.max_input,
            "max_output": self.max_output,
            "sample_rates": self.sample_rates,
            "default_sample_rate": self.default_sample_rate,
            "low_latency": self.low_latency,
            "windows_props": asdict(self.windows_props),
            "driver_version": self.driver_version,
            "hardware_id": self.hardware_id
        }

class DeviceTestResult:
    def __init__(self, device: AudioDeviceMetadata):
        self.device = device
        self.success = False
        self.error = None
        self.tested_sample_rates = []
        self.working_sample_rates = []
        self.tested_at = datetime.now()
        self.latency_ms = None
        self.channels_tested = []
        self.noise_floor = None

    def to_dict(self):
        return {
            "device": self.device.to_dict(),
            "success": self.success,
            "error": str(self.error) if self.error else None,
            "tested_sample_rates": self.tested_sample_rates,
            "working_sample_rates": self.working_sample_rates,
            "tested_at": self.tested_at.isoformat(),
            "latency_ms": self.latency_ms,
            "channels_tested": self.channels_tested,
            "noise_floor": self.noise_floor
        }

class DeviceTester(QThread):
    update_signal = pyqtSignal(str)
    progress_signal = pyqtSignal(int, int)
    finished_signal = pyqtSignal(dict)
    
    def __init__(self, devices: List[AudioDeviceMetadata]):
        super().__init__()
        self.devices = devices
        self.results = {}
        self.test_sample_rates = [44100, 48000, 96000]
        self.test_durations = [0.1, 0.5, 1.0]
        
    def measure_noise_floor(self, stream, duration=0.5):
        frames = int(duration * stream.samplerate)
        data, overflow = stream.read(frames)
        if overflow:
            self.update_signal.emit("Warning: Buffer overflow detected")
        return np.mean(np.abs(data))

    def test_device(self, device: AudioDeviceMetadata) -> DeviceTestResult:
        result = DeviceTestResult(device)
        
        try:
            self.update_signal.emit(f"\nTesting {device.name} ({device.type.name})")
            
            if device.role == DeviceRole.OUTPUT:
                result.error = "Output-only device"
                return result

            for rate in self.test_sample_rates:
                for channels in range(1, min(3, device.max_input + 1)):
                    stream = None
                    try:
                        stream = sd.InputStream(
                            device=device.id,
                            channels=channels,
                            samplerate=rate,
                            blocksize=1024
                        )
                        result.tested_sample_rates.append(rate)
                        result.channels_tested.append(channels)
                        
                        with stream:
                            # Measure initial latency
                            latency_start = time.perf_counter()
                            stream.read(1024)
                            latency = (time.perf_counter() - latency_start) * 1000
                            
                            # Measure noise floor
                            noise_floor = self.measure_noise_floor(stream)
                            if result.noise_floor is None or noise_floor < result.noise_floor:
                                result.noise_floor = noise_floor

                        if rate not in result.working_sample_rates:
                            result.working_sample_rates.append(rate)
                        
                        if not result.latency_ms or latency < result.latency_ms:
                            result.latency_ms = latency
                            
                    except Exception as e:
                        self.update_signal.emit(f"  Failed at {rate}Hz, {channels}ch: {str(e)}")
                        continue
                        
            result.success = len(result.working_sample_rates) > 0
            
        except Exception as e:
            result.error = str(e)
            
        return result

    def run(self):
        for i, device in enumerate(self.devices):
            self.progress_signal.emit(i + 1, len(self.devices))
            result = self.test_device(device)
            self.results[device.id] = result
            
            if result.success:
                self.update_signal.emit(
                    f"✓ Success\n"
                    f"  Working sample rates: {result.working_sample_rates}\n"
                    f"  Channels tested: {result.channels_tested}\n"
                    f"  Best latency: {result.latency_ms:.1f}ms\n"
                    f"  Noise floor: {result.noise_floor:.2e}\n"
                )
            else:
                self.update_signal.emit(
                    f"✗ Failed: {result.error}\n"
                )
                
        self.finished_signal.emit(
            {dev_id: result.to_dict() for dev_id, result in self.results.items()}
        )

class DeviceManager:
    @staticmethod
    def detect_device_type(name: str, props: Dict) -> DeviceType:
        name_lower = name.lower()
        
        virtual_keywords = ['virtual', 'vb-audio', 'voicemeeter', 'cable']
        if any(x in name_lower for x in virtual_keywords):
            return DeviceType.VIRTUAL
            
        if any(x in name_lower for x in ['mic', 'input', 'array']):
            return DeviceType.MICROPHONE
            
        if 'headphone' in name_lower or 'headset' in name_lower:
            return DeviceType.HEADPHONES
            
        if any(x in name_lower for x in ['speaker', 'output']):
            return DeviceType.SPEAKER
            
        if 'loopback' in name_lower:
            return DeviceType.LOOPBACK
            
        return DeviceType.UNKNOWN

    @staticmethod
    def get_windows_properties(device_id: int) -> WindowsAudioProperties:
        props = WindowsAudioProperties()
        try:
            # Here we would use the Windows Core Audio API
            # This is a simplified implementation
            default_id = sd.default.device[1]  # Default output device
            props.is_default = (device_id == default_id)
            
            device_info = sd.query_devices(device_id)
            state_flags = 0x1  # Assume ACTIVE
            if device_info.get('max_input_channels', 0) == 0 and \
               device_info.get('max_output_channels', 0) == 0:
                state_flags |= 0x2  # DISABLED
            
            props.state_flags = state_flags
            props.device_state = " | ".join(
                flag for mask, flag in AUDIO_STATE_FLAGS.items()
                if state_flags & mask
            )
            
        except Exception as e:
            print(f"Error getting Windows properties: {e}")
            
        return props

    @staticmethod
    def get_device_sample_rates(device_id: int) -> List[int]:
        try:
            # Common sample rates to test
            test_rates = [44100, 48000, 88200, 96000, 192000]
            supported_rates = []
            
            for rate in test_rates:
                try:
                    sd.check_input_settings(
                        device=device_id,
                        samplerate=rate,
                        channels=1
                    )
                    supported_rates.append(rate)
                except:
                    continue
                    
            return supported_rates
            
        except Exception as e:
            print(f"Error testing sample rates: {e}")
            return [48000]  # Default fallback

    @staticmethod
    def enumerate_devices() -> List[AudioDeviceMetadata]:
        devices = []
        try:
            device_list = sd.query_devices()
            default_input = sd.default.device[0]
            default_output = sd.default.device[1]
            
            for idx, device in enumerate(device_list):
                try:
                    windows_props = DeviceManager.get_windows_properties(idx)
                    
                    if device['max_input_channels'] > 0 and device['max_output_channels'] > 0:
                        role = DeviceRole.BOTH
                    elif device['max_input_channels'] > 0:
                        role = DeviceRole.INPUT
                    else:
                        role = DeviceRole.OUTPUT
                    
                    device_type = DeviceManager.detect_device_type(
                        device['name'],
                        windows_props.__dict__
                    )
                    
                    sample_rates = DeviceManager.get_device_sample_rates(idx)
                    
                    metadata = AudioDeviceMetadata(
                        id=idx,
                        name=device['name'],
                        type=device_type,
                        role=role,
                        max_input=device['max_input_channels'],
                        max_output=device['max_output_channels'],
                        sample_rates=sample_rates,
                        default_sample_rate=int(device.get('default_samplerate', 48000)),
                        low_latency=device.get('default_low_latency', {}),
                        windows_props=windows_props
                    )
                    
                    devices.append(metadata)
                    
                except Exception as e:
                    print(f"Error processing device {idx}: {e}")
                
        except Exception as e:
            print(f"Error enumerating devices: {e}")
            
        return devices

class AudioDeviceWidget(QGroupBox):
    device_selected = pyqtSignal(AudioDeviceMetadata)
    
    def __init__(self, device: AudioDeviceMetadata, parent=None):
        super().__init__(parent)
        self.device = device
        self.test_results = None
        self.setupUi()
        
    def setupUi(self):
        self.setTitle(f"{self.device.name}")
        layout = QVBoxLayout(self)
        
        # Status indicator and type layout
        status_layout = QHBoxLayout()
        status_icon = self.style().standardIcon(
            QStyle.SP_DialogYesButton if "ACTIVE" in self.device.windows_props.device_state
            else QStyle.SP_DialogNoButton
        )
        status_label = QLabel()
        status_label.setPixmap(status_icon.pixmap(16, 16))
        status_layout.addWidget(status_label)
        
        type_label = QLabel(f"Type: {self.device.type.name}")
        status_layout.addWidget(type_label)
        status_layout.addStretch()
        
        if self.device.windows_props.is_default:
            default_label = QLabel("Default")
            default_label.setStyleSheet("color: green; font-weight: bold;")
            status_layout.addWidget(default_label)
            
        layout.addLayout(status_layout)
        
        # Create collapsible sections
        self.basic_info = QWidget()
        self.basic_layout = QVBoxLayout(self.basic_info)
        
        # Basic device information
        info_text = (
            f"Device ID: {self.device.id}\n"
            f"Role: {self.device.role.name}\n"
            f"Channels: In={self.device.max_input}, Out={self.device.max_output}\n"
            f"Sample Rate: {self.device.default_sample_rate}Hz\n"
            f"State: {self.device.windows_props.device_state}\n"
            f"Hardware ID: {self.device.hardware_id or 'N/A'}\n"
            f"Driver Version: {self.device.driver_version or 'N/A'}"
        )
        
        info_label = QLabel(info_text)
        info_label.setWordWrap(True)
        self.basic_layout.addWidget(info_label)
        
        # Create test results section (initially hidden)
        self.test_results_widget = QWidget()
        self.test_results_layout = QVBoxLayout(self.test_results_widget)
        self.test_results_widget.hide()
        
        # Add sections to main layout
        layout.addWidget(self.basic_info)
        layout.addWidget(self.test_results_widget)
        
        # Action buttons
        button_layout = QHBoxLayout()
        
        self.test_button = QPushButton("Test")
        self.test_button.clicked.connect(self.test_device)
        button_layout.addWidget(self.test_button)
        
        menu_button = QPushButton("Actions")
        menu = QMenu(self)
        menu.addAction("Windows Settings", self.open_windows_settings)
        menu.addAction("Set as Default", lambda: self.set_as_default())
        menu.addAction("Advanced Properties", self.open_advanced_properties)
        menu_button.setMenu(menu)
        button_layout.addWidget(menu_button)
        
        layout.addLayout(button_layout)
        
        # Style based on device type
        self.styleWidget()
        
    def update_test_results(self, results: Dict):
        self.test_results = results
        
        # Clear previous test results
        while self.test_results_layout.count():
            item = self.test_results_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()
                
        if not results:
            self.test_results_widget.hide()
            return
            
        # Create detailed test results display
        results_text = QLabel()
        results_text.setWordWrap(True)
        
        test_info = [
            f"Test Status: {'✓ Success' if results['success'] else '✗ Failed'}",
            f"Tested at: {results['tested_at']}",
            "\nSample Rates:",
            f"  Tested: {', '.join(map(str, results['tested_sample_rates']))}",
            f"  Working: {', '.join(map(str, results['working_sample_rates']))}",
            "\nPerformance Metrics:",
            f"  Latency: {results['latency_ms']:.1f}ms" if results.get('latency_ms') is not None else "  Latency: N/A",
            f"  Noise Floor: {results['noise_floor']:.2e}" if results.get('noise_floor') is not None else "  Noise Floor: N/A",
            f"  Channels Tested: {', '.join(map(str, results['channels_tested']))}"
        ]
        
        if not results['success'] and results['error']:
            test_info.append(f"\nError: {results['error']}")
            
        results_text.setText("\n".join(test_info))
        
        # Style based on test success
        if results['success']:
            results_text.setStyleSheet("QLabel { background-color: #e8f5e9; padding: 8px; border-radius: 4px; }")
        else:
            results_text.setStyleSheet("QLabel { background-color: #ffebee; padding: 8px; border-radius: 4px; }")
            
        self.test_results_layout.addWidget(results_text)
        self.test_results_widget.show()
        
    def test_device(self):
        self.device_selected.emit(self.device)
        
    def styleWidget(self):
        colors = {
            DeviceType.MICROPHONE: "#e3f2fd",
            DeviceType.SPEAKER: "#f3e5f5",
            DeviceType.HEADPHONES: "#e8f5e9",
            DeviceType.VIRTUAL: "#fff3e0",
            DeviceType.LOOPBACK: "#fafafa",
            DeviceType.UNKNOWN: "#f5f5f5"
        }
        
        self.setStyleSheet(f"""
            QGroupBox {{
                background-color: {colors.get(self.device.type, "#ffffff")};
                border: 1px solid #cccccc;
                border-radius: 4px;
                margin-top: 8px;
                padding: 8px;
            }}
            QGroupBox::title {{
                subcontrol-origin: margin;
                left: 7px;
                padding: 0 3px;
            }}
        """)
        
    def open_windows_settings(self):
        subprocess.run(['control', 'mmsys.cpl'])
        
    def set_as_default(self):
        try:
            # This would use Windows API to set default device
            # Simplified implementation
            pass
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Could not set as default: {str(e)}")
            
    def open_advanced_properties(self):
        # This would open Windows advanced properties dialog
        # Simplified implementation
        subprocess.run(['control', 'mmsys.cpl'])

class AudioTestWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.device_manager = DeviceManager()
        self.devices = []
        self.test_results = {}
        self.device_tester = None
        self.last_save_path = None
        self.setupUi()
        self.loadDevices()

    def setupUi(self):
        self.setWindowTitle("Audio Device Manager")
        self.setMinimumSize(800, 900)

        central = QWidget()
        self.setCentralWidget(central)
        main_layout = QVBoxLayout(central)

        # Tabs
        self.tab_widget = QTabWidget()
        
        # Devices Tab
        devices_tab = QWidget()
        devices_layout = QVBoxLayout(devices_tab)
        
        # Toolbar
        toolbar = QHBoxLayout()
        
        refresh_button = QPushButton("Refresh")
        refresh_button.clicked.connect(self.loadDevices)
        toolbar.addWidget(refresh_button)
        
        save_button = QPushButton("Save Results")
        save_button.clicked.connect(self.saveResults)
        toolbar.addWidget(save_button)
        
        toolbar.addStretch()
        devices_layout.addLayout(toolbar)
        
        # Filters
        filter_group = QGroupBox("Filters")
        filter_layout = QHBoxLayout(filter_group)
        
        self.show_inputs = QCheckBox("Input Devices")
        self.show_outputs = QCheckBox("Output Devices")
        self.show_virtual = QCheckBox("Virtual Devices")
        self.show_disconnected = QCheckBox("Show Disconnected")
        
        for cb in [self.show_inputs, self.show_outputs, self.show_virtual, self.show_disconnected]:
            cb.setChecked(True)
            cb.stateChanged.connect(self.filterDevices)
            filter_layout.addWidget(cb)
            
        devices_layout.addWidget(filter_group)
        
        # Device list
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        self.device_list = QWidget()
        self.device_layout = QVBoxLayout(self.device_list)
        scroll.setWidget(self.device_list)
        devices_layout.addWidget(scroll)
        
        # Testing Tab
        testing_tab = QWidget()
        testing_layout = QVBoxLayout(testing_tab)
        
        test_controls = QHBoxLayout()
        
        self.test_all_button = QPushButton("Test All Devices")
        self.test_all_button.clicked.connect(self.startDeviceTesting)
        test_controls.addWidget(self.test_all_button)
        
        self.stop_test_button = QPushButton("Stop Testing")
        self.stop_test_button.clicked.connect(self.stopTesting)
        self.stop_test_button.setEnabled(False)
        test_controls.addWidget(self.stop_test_button)
        
        testing_layout.addLayout(test_controls)
        
        self.test_progress = QProgressBar()
        testing_layout.addWidget(self.test_progress)
        
        self.info_display = QPlainTextEdit()
        self.info_display.setReadOnly(True)
        testing_layout.addWidget(self.info_display)
        
        # Add tabs
        self.tab_widget.addTab(devices_tab, "Devices")
        self.tab_widget.addTab(testing_tab, "Testing")
        
        main_layout.addWidget(self.tab_widget)
        
        # Status bar
        self.statusBar().showMessage("Ready")

    def filterDevices(self):
        self.updateDeviceList()
        
    def loadDevices(self):
        self.devices = self.device_manager.enumerate_devices()
        self.updateDeviceList()
        self.statusBar().showMessage(f"Loaded {len(self.devices)} devices")
        
    def updateDeviceList(self):
        while self.device_layout.count():
            item = self.device_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()
                
        for device in self.devices:
            if not self.shouldShowDevice(device):
                continue
                
            widget = AudioDeviceWidget(device)
            widget.device_selected.connect(self.onDeviceSelected)
            self.device_layout.addWidget(widget)
            
        self.device_layout.addStretch()
        
    def shouldShowDevice(self, device: AudioDeviceMetadata) -> bool:
        if not self.show_virtual.isChecked() and device.type == DeviceType.VIRTUAL:
            return False
            
        if not self.show_inputs.isChecked() and device.role in [DeviceRole.INPUT, DeviceRole.BOTH]:
            return False
            
        if not self.show_outputs.isChecked() and device.role in [DeviceRole.OUTPUT, DeviceRole.BOTH]:
            return False
            
        if not self.show_disconnected.isChecked() and "NOT_PRESENT" in device.windows_props.device_state:
            return False
            
        return True
        
    def onDeviceSelected(self, device: AudioDeviceMetadata):
        self.test_single_device(device)
        
    def test_single_device(self, device: AudioDeviceMetadata):
        self.tab_widget.setCurrentIndex(1)  # Switch to testing tab
        self.device_tester = DeviceTester([device])
        self.startTesting()
        
    def startDeviceTesting(self):
        filtered_devices = [d for d in self.devices if self.shouldShowDevice(d)]
        self.device_tester = DeviceTester(filtered_devices)
        self.startTesting()
        
    def startTesting(self):
        if not self.device_tester:
            return
            
        self.test_all_button.setEnabled(False)
        self.stop_test_button.setEnabled(True)
        self.test_progress.setValue(0)
        self.info_display.clear()
        
        self.device_tester.update_signal.connect(self.appendTestInfo)
        self.device_tester.progress_signal.connect(self.updateProgress)
        self.device_tester.finished_signal.connect(self.testingFinished)
        self.device_tester.start()
        
    def stopTesting(self):
        if self.device_tester and self.device_tester.isRunning():
            self.device_tester.terminate()
            self.device_tester.wait()
            self.testingFinished({})
            
    def appendTestInfo(self, text: str):
        self.info_display.appendPlainText(text)
        
    def updateProgress(self, current: int, total: int):
        self.test_progress.setMaximum(total)
        self.test_progress.setValue(current)
        self.statusBar().showMessage(f"Testing device {current} of {total}")
        
    def testingFinished(self, results: Dict):
        self.test_results.update(results)
        self.test_all_button.setEnabled(True)
        self.stop_test_button.setEnabled(False)
        self.statusBar().showMessage("Testing completed")
        self.summarizeResults()
        
        # Update device widgets with test results
        for i in range(self.device_layout.count()):
            item = self.device_layout.itemAt(i)
            if item and item.widget():
                widget = item.widget()
                if isinstance(widget, AudioDeviceWidget):
                    device_id = widget.device.id
                    if device_id in results:
                        widget.update_test_results(results[device_id])
        
    def summarizeResults(self):
        self.info_display.appendPlainText("\n" + "=" * 50)
        self.info_display.appendPlainText("TEST RESULTS SUMMARY")
        self.info_display.appendPlainText("=" * 50 + "\n")
        
        total = len(self.test_results)
        success = sum(1 for r in self.test_results.values() if r['success'])
        
        summary = (
            f"Total Devices Tested: {total}\n"
            f"Successfully Tested: {success}\n"
            f"Failed Tests: {total - success}\n\n"
            "Working Devices:\n"
            "-" * 20 + "\n"
        )
        
        self.info_display.appendPlainText(summary)
        
        for result in self.test_results.values():
            if result['success']:
                device = result['device']
                self.info_display.appendPlainText(
                    f"{device['name']}\n"
                    f"  Sample Rates: {result['working_sample_rates']}\n"
                    f"  Latency: {result['latency_ms']:.1f}ms\n"
                    f"  Noise Floor: {result.get('noise_floor', 'N/A')}\n"
                )
                
    def saveResults(self):
        if not self.test_results:
            QMessageBox.warning(self, "Warning", "No test results to save.")
            return
            
        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "Save Test Results",
            self.last_save_path or str(Path.home() / "audio_test_results.json"),
            "JSON Files (*.json)"
        )
        
        if not file_path:
            return
            
        try:
            self.last_save_path = file_path
            with open(file_path, 'w') as f:
                json.dump({
                    'timestamp': datetime.now().isoformat(),
                    'results': self.test_results
                }, f, indent=2)
            
            self.statusBar().showMessage(f"Results saved to {file_path}")
            
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to save results: {str(e)}")
            
    def closeEvent(self, event):
        self.stopTesting()
        event.accept()

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