main.cpp

#include "logger.h"
#include "usb_monitor.h"
#include "audio_monitor.h"
#include "usb_power.h"
#include <Windows.h>
#include <atomic>
#include <string>

// Global monitors
USBMonitor g_usbMonitor;
AudioMonitor g_audioMonitor;
USBPowerMonitor g_powerMonitor;
std::atomic<bool> g_running{true};

// Hidden window class name
const wchar_t* WINDOW_CLASS_NAME = L"DeviceMonitorWindow";

// Window procedure for receiving device change messages
LRESULT CALLBACK WindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {
    switch (uMsg) {
        case WM_DEVICECHANGE:
            g_usbMonitor.processDeviceChange(wParam, lParam);
            return TRUE;

        case WM_CLOSE:
            g_running = false;
            PostQuitMessage(0);
            return 0;

        default:
            return DefWindowProcW(hwnd, uMsg, wParam, lParam);
    }
}

// Console control handler for Ctrl+C
BOOL WINAPI ConsoleHandler(DWORD ctrlType) {
    switch (ctrlType) {
        case CTRL_C_EVENT:
        case CTRL_BREAK_EVENT:
        case CTRL_CLOSE_EVENT:
            LOG_INFO(L"Shutting down...");
            g_running = false;
            return TRUE;
        default:
        return FALSE;
    }
}

HWND createHiddenWindow() {
    WNDCLASSEXW wc = {};
    wc.cbSize = sizeof(WNDCLASSEXW);
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = GetModuleHandleW(nullptr);
    wc.lpszClassName = WINDOW_CLASS_NAME;

    if (!RegisterClassExW(&wc)) {
        return nullptr;
    }

    return CreateWindowExW(
        0,
        WINDOW_CLASS_NAME,
        L"Device Monitor",
        0,
        0, 0, 0, 0,
        HWND_MESSAGE,  // Message-only window
        nullptr,
        GetModuleHandleW(nullptr),
        nullptr);
}

void writeConsole(const std::wstring& text) {
    HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
    DWORD written = 0;
    WriteConsoleW(hConsole, text.c_str(), static_cast<DWORD>(text.length()), &written, nullptr);
}

void printBanner() {
    writeConsole(L"\n");
    writeConsole(L"  ========================================\n");
    writeConsole(L"    USB & Audio Device Monitor\n");
    writeConsole(L"    Press Ctrl+C to exit\n");
    writeConsole(L"  ========================================\n");
    writeConsole(L"\n");
}

int main() {
    // Initialize COM for audio monitoring
    HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
    if (FAILED(hr)) {
        writeConsole(L"Failed to initialize COM\n");
        return 1;
    }

    // Initialize logger
    Logger::instance().init();

    printBanner();

    // Create hidden window for USB notifications
    HWND hwnd = createHiddenWindow();
    if (!hwnd) {
        LOG_ERROR(L"Failed to create message window");
        CoUninitialize();
        return 1;
    }

    // Setup console control handler
    SetConsoleCtrlHandler(ConsoleHandler, TRUE);

    // Start USB monitoring
    if (!g_usbMonitor.start(hwnd)) {
        LOG_ERROR(L"Failed to start USB monitor");
    }

    // Start audio monitoring
    if (!g_audioMonitor.start()) {
        LOG_ERROR(L"Failed to start audio monitor");
    }

    // List current devices
    g_usbMonitor.listCurrentDevices();
    writeConsole(L"\n");

    // Show USB power topology
    g_powerMonitor.printTopology();
    writeConsole(L"\n");
    g_powerMonitor.printPowerSummary();
    writeConsole(L"\n");

    g_audioMonitor.listCurrentDevices();
    writeConsole(L"\n");

    LOG_SECTION(L"MONITORING STARTED");
    LOG_INFO(L"Waiting for device changes...");
    writeConsole(L"\n");

    // Message loop
    MSG msg;
    while (g_running) {
        // Use PeekMessage to allow checking g_running
        while (PeekMessageW(&msg, nullptr, 0, 0, PM_REMOVE)) {
            if (msg.message == WM_QUIT) {
                g_running = false;
                break;
            }
            TranslateMessage(&msg);
            DispatchMessageW(&msg);
        }

        // Small sleep to reduce CPU usage
        Sleep(10);
    }

    // Cleanup
    LOG_INFO(L"Stopping monitors...");
    g_usbMonitor.stop();
    g_audioMonitor.stop();

    DestroyWindow(hwnd);
    UnregisterClassW(WINDOW_CLASS_NAME, GetModuleHandleW(nullptr));

    Logger::instance().shutdown();
    CoUninitialize();

    return 0;
}