#include "StdAfx.h"

PnPControl* PnPControl::Instance = NULL;
PnPControl::Reference = 0;

void PnPControl::FindControllers()
{
    wstring name;
    for (int i = 0; FindInstance(GUID_VIOSERIAL_CONTROLLER, i, name); ++i)
    {
        AddController(name.c_str());
    }
}

void PnPControl::FindPorts()
{
    wstring name;
    for (int i = 0; FindInstance(GUID_VIOSERIAL_PORT, i, name); ++i)
    {
        AddPort(name.c_str());
    }
}


BOOL PnPControl::FindInstance(GUID guid, DWORD idx, wstring& name)
{
    HDEVINFO HardwareDeviceInfo;
    SP_DEVICE_INTERFACE_DATA DeviceInterfaceData;
    HardwareDeviceInfo = SetupDiGetClassDevs(
        &guid,
        NULL,
        NULL,
        (DIGCF_PRESENT | DIGCF_DEVICEINTERFACE)
        );

    if (HardwareDeviceInfo == INVALID_HANDLE_VALUE)
    {
        printf("Cannot get class devices.\n");
        return FALSE;
    }

    DeviceInterfaceData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);
    if (SetupDiEnumDeviceInterfaces(HardwareDeviceInfo, NULL, &guid, idx, &DeviceInterfaceData))
    {
        DWORD RequiredLength = 0;
        PSP_DEVICE_INTERFACE_DETAIL_DATA DeviceInterfaceDetailData = NULL;
        SP_DEVINFO_DATA DevInfoData = {sizeof(SP_DEVINFO_DATA)};

        SetupDiGetDeviceInterfaceDetail(
            HardwareDeviceInfo,
            &DeviceInterfaceData,
            NULL,
            0,
            &RequiredLength,
            NULL
            );

        DeviceInterfaceDetailData = (PSP_DEVICE_INTERFACE_DETAIL_DATA) malloc(RequiredLength);

        if (DeviceInterfaceDetailData == NULL)
        {
            printf("Cannot allocate memory.\n");
            return FALSE;
        }

        DeviceInterfaceDetailData->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA);

        SetupDiGetDeviceInterfaceDetail(
            HardwareDeviceInfo,
            &DeviceInterfaceData,
            DeviceInterfaceDetailData,
            RequiredLength,
            NULL,
            &DevInfoData
            );
        name = DeviceInterfaceDetailData->DevicePath;
        free ((PVOID)DeviceInterfaceDetailData);
        SetupDiDestroyDeviceInfoList(HardwareDeviceInfo);
        return TRUE;
    }
    SetupDiDestroyDeviceInfoList(HardwareDeviceInfo);
    return FALSE;
}

void PnPControl::Init( )
{
    DWORD id;
    Thread = CreateThread(
        NULL,
        0,
        (LPTHREAD_START_ROUTINE) ServiceThread,
        (LPVOID)this,
        0,
        &id);

    if (Thread == NULL)
    {
        printf("Cannot create thread Error = %d.\n", GetLastError());
    }

    if ( !InitializeCriticalSectionAndSpinCount(&PortsCS, 0x4000))
    {
        printf("Cannot initalize critical section Error = %d.\n", GetLastError());
    }

}

void PnPControl::Close( )
{
    if (PortNotify != NULL)
    {
        UnregisterDeviceNotification(PortNotify);
        PortNotify = NULL;
    }
    if (ControllerNotify != NULL)
    {
        UnregisterDeviceNotification(ControllerNotify);
        ControllerNotify = NULL;
    }
    if (Wnd && Thread)
    {
        SendMessage(Wnd, WM_DESTROY, 0, 0);
        if (WAIT_TIMEOUT == WaitForSingleObject(Thread, 1000))
        {
            printf("Cannot close thread after 1 sec\n");
            TerminateThread(Thread, 0);
        }
        Thread = NULL;
    }
    DeleteCriticalSection(&PortsCS);
}

DWORD WINAPI PnPControl::ServiceThread(PnPControl* ptr)
{
    ptr->Run();
    return 0;
}

void PnPControl::Run()
{
    MSG Msg;
    WNDCLASSEX wc;

    wc.cbSize        = sizeof(WNDCLASSEX);
    wc.style         = 0;
    wc.lpfnWndProc   = GlobalWndProc;
    wc.cbClsExtra    = 0;
    wc.cbWndExtra    = 0;
    wc.hInstance     = reinterpret_cast<HINSTANCE>(GetModuleHandle(0));
    wc.hIcon         = LoadIcon(NULL, IDI_APPLICATION);
    wc.hCursor       = LoadCursor(NULL, IDC_ARROW);
    wc.hbrBackground = (HBRUSH)(COLOR_WINDOW+1);
    wc.lpszMenuName  = NULL;
    wc.lpszClassName = L"VioSerialLib";
    wc.hIconSm       = LoadIcon(NULL, IDI_APPLICATION);

    if(!RegisterClassEx(&wc))
    {
        printf("Window Registration Failed!\n");
    }

    Wnd = CreateWindowEx(
        WS_EX_CLIENTEDGE,
        L"VioSerialLib",
        NULL,
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT, CW_USEDEFAULT, CW_USEDEFAULT,
        NULL, NULL, NULL, NULL
        );

    if(Wnd == NULL)
    {
        printf("Window Creation Failed!\n");
    }
    SetWindowLongPtr(Wnd, GWLP_USERDATA, (LONG_PTR)(this));
    ControllerNotify = RegisterInterfaceNotify(GUID_VIOSERIAL_CONTROLLER);
    PortNotify = RegisterInterfaceNotify(GUID_VIOSERIAL_PORT);
    while(GetMessage(&Msg, NULL, 0, 0) > 0)
    {
        TranslateMessage(&Msg);
        DispatchMessage(&Msg);
    }
}

LRESULT CALLBACK PnPControl::GlobalWndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    switch (msg)
    {
    case WM_CLOSE:
        DestroyWindow(hWnd);
        break;
    case WM_DESTROY:
        PostQuitMessage(0);
        break;
    case WM_DEVICECHANGE:
        ProcessPnPNotification((PnPControl*)(GetWindowLongPtr(hWnd, GWLP_USERDATA)), PnPNotification(msg, wParam, lParam));
        break;
    default:
        return DefWindowProc(hWnd, msg, wParam, lParam);
    }
    return 0;
}

HDEVNOTIFY PnPControl::RegisterInterfaceNotify( 
    GUID InterfaceClassGuid
    )
{
    DEV_BROADCAST_DEVICEINTERFACE NotificationFilter;
    HDEVNOTIFY Notify;

    ZeroMemory( &NotificationFilter, sizeof(NotificationFilter) );
    NotificationFilter.dbcc_size = 
        sizeof(DEV_BROADCAST_DEVICEINTERFACE);
    NotificationFilter.dbcc_devicetype = DBT_DEVTYP_DEVICEINTERFACE;
    NotificationFilter.dbcc_classguid = InterfaceClassGuid;

    Notify = RegisterDeviceNotification( Wnd,
        &NotificationFilter,
        DEVICE_NOTIFY_WINDOW_HANDLE
        );

    if (Notify == NULL)
    {
        printf("RegisterDeviceNotification failed: %d\n",
            GetLastError());
    }
    return Notify;
}

HDEVNOTIFY PnPControl::RegisterHandleNotify(HANDLE handle)
{
    HDEVNOTIFY Notify;
    DEV_BROADCAST_HANDLE NotificationFilter;

    ZeroMemory( &NotificationFilter, sizeof(NotificationFilter) );
    NotificationFilter.dbch_size =
        sizeof(DEV_BROADCAST_HANDLE);
    NotificationFilter.dbch_devicetype = DBT_DEVTYP_HANDLE;
    NotificationFilter.dbch_handle = handle;
    NotificationFilter.dbch_eventguid = GUID_VIOSERIAL_PORT_CHANGE_STATUS;


    Notify = RegisterDeviceNotification( Wnd,
        &NotificationFilter,
        DEVICE_NOTIFY_WINDOW_HANDLE
        );

    if (Notify == NULL)
    {
        printf("RegisterDeviceNotification failed: %d\n",
            GetLastError());
    }
    return Notify;
}

void PnPControl::ProcessPnPNotification(PnPControl* ptr, PnPNotification Notification)
{
    switch (Notification.wParam) 
    {
    case DBT_DEVICEARRIVAL: {
        PDEV_BROADCAST_HDR pHdr = (PDEV_BROADCAST_HDR)Notification.lParam;
        if (pHdr->dbch_devicetype == DBT_DEVTYP_DEVICEINTERFACE) {
            PDEV_BROADCAST_DEVICEINTERFACE pDevInf = (PDEV_BROADCAST_DEVICEINTERFACE)pHdr;
            if (IsEqualGUID(GUID_VIOSERIAL_CONTROLLER, pDevInf->dbcc_classguid))
            {
                ptr->AddController(pDevInf->dbcc_name);
            }
            else if (IsEqualGUID(GUID_VIOSERIAL_PORT, pDevInf->dbcc_classguid))
            {
                ptr->AddPort(pDevInf->dbcc_name);
            }
        }
                            }
                            break;
    case DBT_DEVICEREMOVECOMPLETE: {
        PDEV_BROADCAST_HDR pHdr = (PDEV_BROADCAST_HDR)Notification.lParam;
        if (pHdr->dbch_devicetype == DBT_DEVTYP_DEVICEINTERFACE) {
            PDEV_BROADCAST_DEVICEINTERFACE pDevInf = (PDEV_BROADCAST_DEVICEINTERFACE)pHdr;
            if (IsEqualGUID(GUID_VIOSERIAL_CONTROLLER, pDevInf->dbcc_classguid))
            {
                ptr->RemoveController(pDevInf->dbcc_name);
            }
            else if (IsEqualGUID(GUID_VIOSERIAL_PORT, pDevInf->dbcc_classguid))
            {
                ptr->RemovePort(pDevInf->dbcc_name);
            }
        }
                                   }
                                   break;
    default:
        ptr->DispatchPnpMessage(Notification);
        break;
    }
}

void PnPControl::AddController(const wchar_t* name)
{
    printf ("Add Controller %ws\n", name);
    Controllers.push_back((new SerialController(name)));
}

void PnPControl::RemoveController(wchar_t* name)
{
    printf ("Remove Controller %ws\n", name);
    for(Iterator it = Controllers.begin(); it != Controllers.end(); it++)
    {
        if (_wcsnicmp((*it)->Name.c_str(), name, (*it)->Name.size()) == 0)
        {
            delete *it;
            Controllers.remove(*it);
            break;
        }
    }
}

void PnPControl::AddPort(const wchar_t* name)
{
    printf ("Add Port %ws\n", name);
    EnterCriticalSection(&PortsCS);
    Ports.push_back(new SerialPort(name, this));
    LeaveCriticalSection(&PortsCS);
}

void PnPControl::RemovePort(wchar_t* name)
{
    printf ("Remove Port %ws\n", name);
    EnterCriticalSection(&PortsCS);
    for(Iterator it = Ports.begin(); it != Ports.end(); it++)
    {
        if (_wcsnicmp((*it)->Name.c_str(), name, (*it)->Name.size()) == 0)
        {
            delete *it;
            Ports.remove(*it);
            break;
        }
    }
    LeaveCriticalSection(&PortsCS);
}

BOOL PnPControl::FindPort(const wchar_t* name)
{
    wstring tmp = name;
    EnterCriticalSection(&PortsCS);
    BOOL ret = FALSE;
    for(Iterator it = Ports.begin(); it != Ports.end(); it++)
    {
        if (_wcsnicmp((*it)->SymbolicName.c_str(), name, (*it)->Name.size()) == 0)
        {
            ret = TRUE;
            break;
        }
    }
    LeaveCriticalSection(&PortsCS);
    return ret;
}
PVOID PnPControl::OpenPortByName(const wchar_t* name)
{
    wstring tmp = name;
    PVOID ret = NULL;
    EnterCriticalSection(&PortsCS);
    for(Iterator it = Ports.begin(); it != Ports.end(); it++)
    {
        if (_wcsnicmp((*it)->SymbolicName.c_str(), name, (*it)->Name.size()) == 0)
        {
            if(((SerialPort*)(*it))->OpenPort() == TRUE)
            {
                ret = *it;
                break;
            }
        }
    }
    LeaveCriticalSection(&PortsCS);
    return ret;
}
PVOID PnPControl::OpenPortById(UINT id)
{
    if ((size_t)(id) > NumPorts())
    {
        return NULL;
    }
    PVOID ret = NULL;
    EnterCriticalSection(&PortsCS);
    Iterator it = Ports.begin();
    advance(it, id);
    if (((SerialPort*)(*it))->OpenPort() == TRUE)
    {
        ret = *it;
    }
    LeaveCriticalSection(&PortsCS);
    return ret;
}
BOOL PnPControl::ReadPort(PVOID port, PVOID buf, PULONG size)
{
    Iterator it;
    EnterCriticalSection(&PortsCS);
    for(it = Ports.begin(); it != Ports.end(); it++)
    {
        if (*it == port)
        {
            break;
        }
    }
    LeaveCriticalSection(&PortsCS);
    return (*it != port) ? FALSE : ((SerialPort*)(*it))->ReadPort(buf, (size_t*)(size));
}
BOOL PnPControl::WritePort(PVOID port, PVOID buf, ULONG size)
{
    Iterator it;
    EnterCriticalSection(&PortsCS);
    for(it = Ports.begin(); it != Ports.end(); it++)
    {
        if (*it == port)
        {
            break;
        }
    }
    LeaveCriticalSection(&PortsCS);
    return (*it != port) ? FALSE : ((SerialPort*)(*it))->WritePort(buf, (size_t*)(&size));
}
VOID PnPControl::ClosePort(PVOID port)
{
    EnterCriticalSection(&PortsCS);
    for(Iterator it = Ports.begin(); it != Ports.end(); it++)
    {
        if (*it == port)
        {
            ((SerialPort*)(*it))->ClosePort();
            break;
        }
    }
    LeaveCriticalSection(&PortsCS);
}

wchar_t* PnPControl::PortSymbolicName(int index)
{
    if ((size_t)(index) > NumPorts())
    {
        return NULL;
    }
    EnterCriticalSection(&PortsCS);
    Iterator it = Ports.begin();
    advance(it, index);
    return (wchar_t*)((*it)->SymbolicName.c_str());
}

VOID PnPControl::RegisterNotification(PVOID port, VIOSERIALNOTIFYCALLBACK pfn, PVOID ptr)
{
    for(Iterator it = Ports.begin(); it != Ports.end(); it++)
    {
        if (*it == port)
        {
            ((SerialPort*)(*it))->NotificationPair.first = pfn;
            ((SerialPort*)(*it))->NotificationPair.second = ptr;
            return;
        }
    }
}

BOOL PnPControl::IsRunningAsService()
{
    wchar_t wstrPath[_MAX_FNAME];
    if (!GetModuleFileName( NULL, wstrPath, _MAX_FNAME))
    {
        printf ("Error getting module file name (%d)\n", GetLastError());
        return FALSE;
    }
    printf ("Module File Name is %ws\n", wstrPath);

    wstring fullname(wstrPath);
    size_t pos_begin = fullname.rfind(L'\\') + 1;
    size_t pos_end   = fullname.rfind(L'.');
    wstring filename(fullname, pos_begin, pos_end-pos_begin);
    printf ("File name = %ws\n", filename.c_str());
    SERVICE_STATUS_PROCESS status;
    DWORD                  needed;

    SC_HANDLE scm = OpenSCManager (NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (scm == NULL)
    {
        printf ("OpenSCManager failed (%d)\n", GetLastError());
        return FALSE;
    }
    SC_HANDLE svc = OpenService (scm, filename.c_str(), SERVICE_ALL_ACCESS);
    if (svc == NULL)
    {
        printf ("OpenService failed (%d)\n", GetLastError());
        return FALSE;
    }
    if(!QueryServiceStatusEx(svc, SC_STATUS_PROCESS_INFO, (LPBYTE) &status, sizeof(SERVICE_STATUS_PROCESS), &needed))
    {
        printf ("QueryServiceStatusEx failed (%d)\n", GetLastError());
        CloseServiceHandle (svc);
        CloseServiceHandle (scm);
        return FALSE;
    }

    CloseServiceHandle (svc);
    CloseServiceHandle (scm);
    return (status.dwProcessId == GetCurrentProcessId());
}