// SearchDeviceDlg.cpp : implementation file
//

#include "StdAfx.h"

#include <dbt.h>

#include "USBFilterEditor.h"
#include "SearchDeviceDlg.h"

//

HINSTANCE ghInstance;

HWND ghTreeWnd;
HWND ghEditWnd;
HWND ghStatusWnd;

HTREEITEM ghTreeRoot;

HDEVNOTIFY gNotifyDevHandle;
HDEVNOTIFY gNotifyHubHandle;

int giComputer;
int giHub;
int giNoDevice;
int giGoodDevice;
int giBadDevice;

//

extern "C" {

#include "usbview.h"

HTREEITEM AddLeaf(HTREEITEM hTreeParent, LPARAM lParam, LPWSTR lpszText, TREEICON TreeIcon)
{
    TVINSERTSTRUCT tvins;
    HTREEITEM      hti;

    memset(&tvins, 0, sizeof(tvins));

    // Set the parent item
    //
    tvins.hParent = hTreeParent;

    tvins.hInsertAfter = TVI_LAST;

    // pszText and lParam members are valid
    //
    tvins.item.mask = TVIF_TEXT | TVIF_PARAM;

    // Set the text of the item.
    //
    tvins.item.pszText = lpszText;

    // Set the user context item
    //
    tvins.item.lParam = lParam;

    // Add the item to the tree-view control.
    //
    hti = TreeView_InsertItem(ghTreeWnd, &tvins);

    tvins.item.mask = TVIF_IMAGE | TVIF_SELECTEDIMAGE;
    tvins.item.hItem = hti;

    // Determine which icon to display for the device
    //
    switch (TreeIcon)
    {
    case ComputerIcon:
        tvins.item.iImage = giComputer;
        tvins.item.iSelectedImage = giComputer;
        break;

    case HubIcon:
        tvins.item.iImage = giHub;
        tvins.item.iSelectedImage = giHub;
        break;

    case NoDeviceIcon:
        tvins.item.iImage = giNoDevice;
        tvins.item.iSelectedImage = giNoDevice;
        break;

    case GoodDeviceIcon:
        tvins.item.iImage = giGoodDevice;
        tvins.item.iSelectedImage = giGoodDevice;
        break;

    case BadDeviceIcon:
    default:
        tvins.item.iImage = giBadDevice;
        tvins.item.iSelectedImage = giBadDevice;
        break;
    }

    TreeView_SetItem(ghTreeWnd, &tvins.item);

    return hti;

    return 0;
}

void WalkTree(HTREEITEM hTreeItem, LPFNTREECALLBACK lpfnTreeCallback, DWORD dwRefData)
{
    if (hTreeItem)
    {
        // Recursively call WalkTree on the node's first child.
        //
        WalkTree(TreeView_GetChild(ghTreeWnd, hTreeItem),
            lpfnTreeCallback,
            dwRefData);

        //
        // Call the lpfnCallBack on the node itself.
        //
        (*lpfnTreeCallback)(ghTreeWnd, hTreeItem);

        //
        //
        // Recursively call WalkTree on the node's first sibling.
        //
        WalkTree(TreeView_GetNextSibling(ghTreeWnd, hTreeItem),
            lpfnTreeCallback,
            dwRefData);
    }
}

void ExpandItem(HWND hTreeWnd, HTREEITEM hTreeItem)
{
    //
    // Make this node visible.
    //
    TreeView_Expand(hTreeWnd, hTreeItem, TVE_EXPAND);
}

void DestroyTree()
{
    // Clear the selection of the TreeView, so that when the tree is
    // destroyed, the control won't try to constantly "shift" the
    // selection to another item.
    //
    TreeView_SelectItem(ghTreeWnd, NULL);

    // Destroy the current contents of the TreeView
    //
    if (ghTreeRoot)
    {
        WalkTree(ghTreeRoot, CleanupItem, 0);

        TreeView_DeleteAllItems(ghTreeWnd);

        ghTreeRoot = NULL;
    }
}

VOID RefreshTree (VOID)
{
    WCHAR statusText[128];
    ULONG devicesConnected;

    // Clear the selection of the TreeView, so that when the tree is
    // destroyed, the control won't try to constantly "shift" the
    // selection to another item.
    //
    TreeView_SelectItem(ghTreeWnd, NULL);

    // Clear the edit control
    //
    SetWindowText(ghEditWnd, L"");

    // Destroy the current contents of the TreeView
    //
    if (ghTreeRoot)
    {
        WalkTree(ghTreeRoot, CleanupItem, 0);

        TreeView_DeleteAllItems(ghTreeWnd);

        ghTreeRoot = NULL;
    }

    // Create the root tree node
    //
    ghTreeRoot = AddLeaf(TVI_ROOT, 0, L"My Computer", ComputerIcon);

    if (ghTreeRoot != NULL)
    {
        // Enumerate all USB buses and populate the tree
        //
        EnumerateHostControllers(ghTreeRoot, &devicesConnected);

        //
        // Expand all tree nodes
        //
        WalkTree(ghTreeRoot, ExpandItem, 0);

        // Update Status Line with number of devices connected
        //
        wsprintf(statusText, L"Devices Connected: %d   Hubs Connected: %d",
            devicesConnected, TotalHubs);

        SetWindowText(ghStatusWnd, statusText);
    }
}

} // extern "C"

// CSearchDeviceDlg dialog

IMPLEMENT_DYNAMIC(CSearchDeviceDlg, CDialog)

BEGIN_MESSAGE_MAP(CSearchDeviceDlg, CDialog)
    ON_WM_DEVICECHANGE()
    ON_NOTIFY(NM_DBLCLK, IDC_TREE_DEVICES, &CSearchDeviceDlg::OnNMDblclkTreeDevices)
END_MESSAGE_MAP()

CSearchDeviceDlg::CSearchDeviceDlg(CWnd* pParent /*=NULL*/) :
    CDialog(CSearchDeviceDlg::IDD, pParent),
    m_nClassId(-1),
    m_nVendorId(-1),
    m_nProductId(-1),
    m_nRevision(-1),
    m_nAction(-1)
{

}

CSearchDeviceDlg::~CSearchDeviceDlg()
{
}

void CSearchDeviceDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);

    if (pDX->m_bSaveAndValidate == TRUE)
    {
        CTreeCtrl *pDevices = (CTreeCtrl *)GetDlgItem(IDC_TREE_DEVICES);
        ASSERT(pDevices != NULL);

        HTREEITEM hItem = pDevices->GetSelectedItem();
        if (hItem != NULL)
        {
            DWORD_PTR info = pDevices->GetItemData(hItem);
            ASSERT(info != NULL);

            if (*(PUSBDEVICEINFOTYPE)info == DeviceInfo)
            {
                PUSBDEVICEINFO devinfo = reinterpret_cast<PUSBDEVICEINFO>(info);

                if (devinfo->ConnectionInfo != NULL)
                {
                    PUSB_NODE_CONNECTION_INFORMATION_EX conninfo = devinfo->ConnectionInfo;
                    if (conninfo->ConnectionStatus != NoDeviceConnected)
                    {
                        m_nClassId = conninfo->DeviceDescriptor.bDeviceClass;
                        m_nVendorId = conninfo->DeviceDescriptor.idVendor;
                        m_nProductId = conninfo->DeviceDescriptor.idProduct;
                    }
                    else
                    {
                        m_nClassId = -1;
                        m_nVendorId = -1;
                        m_nProductId = -1;
                    }
                }
            }

            CButton *pAllow = (CButton *)GetDlgItem(IDC_RADIO_ACTION_1);
            ASSERT(pAllow != NULL);
            m_nAction = (pAllow->GetCheck() == BST_CHECKED);
        }
    }
}

BOOL CSearchDeviceDlg::OnInitDialog()
{
    if (CDialog::OnInitDialog() == FALSE)
    {
        return FALSE;
    }

    CButton *pAllow = (CButton *)GetDlgItem(IDC_RADIO_ACTION_1);
    ASSERT(pAllow != NULL);
    pAllow->SetCheck(BST_CHECKED);

    CTreeCtrl *pDevices = (CTreeCtrl *)GetDlgItem(IDC_TREE_DEVICES);
    ASSERT(pDevices != NULL);

    ghInstance = AfxGetInstanceHandle();
    ghTreeWnd = pDevices->m_hWnd;

    RegisterNotification();
    CreateImageList();
    RefreshTree();

    return TRUE;
}

BOOL CSearchDeviceDlg::DestroyWindow()
{
    UnregisterDeviceNotification(gNotifyDevHandle);
    UnregisterDeviceNotification(gNotifyHubHandle);

    DestroyTree();

    return CDialog::DestroyWindow();
}

void CSearchDeviceDlg::OnOK()
{
    UpdateData(TRUE);

    CDialog::OnOK();
}

void CSearchDeviceDlg::RegisterNotification()
{
    DEV_BROADCAST_DEVICEINTERFACE   broadcastInterface;

    // Register to receive notification when a USB device is plugged in.
    broadcastInterface.dbcc_size = sizeof(DEV_BROADCAST_DEVICEINTERFACE);
    broadcastInterface.dbcc_devicetype = DBT_DEVTYP_DEVICEINTERFACE;

    memcpy(&(broadcastInterface.dbcc_classguid),
        &(GUID_CLASS_USB_DEVICE),
        sizeof(struct _GUID));

    gNotifyDevHandle = RegisterDeviceNotification(m_hWnd,
        &broadcastInterface,
        DEVICE_NOTIFY_WINDOW_HANDLE);

    // Now register for Hub notifications.
    memcpy(&(broadcastInterface.dbcc_classguid),
        &(GUID_CLASS_USBHUB),
        sizeof(struct _GUID));

    gNotifyHubHandle = RegisterDeviceNotification(m_hWnd,
        &broadcastInterface,
        DEVICE_NOTIFY_WINDOW_HANDLE);
}

void CSearchDeviceDlg::CreateImageList()
{
    HIMAGELIST himl;
    HICON      hicon;

    if ((himl = ImageList_Create(16, 16, FALSE, 2, 0)) != NULL)
    {
        hicon = LoadIcon(ghInstance, MAKEINTRESOURCE(IDI_ICON));
        giGoodDevice = ImageList_AddIcon(himl, hicon);

        hicon = LoadIcon(ghInstance, MAKEINTRESOURCE(IDI_BADICON));
        giBadDevice = ImageList_AddIcon(himl, hicon);

        hicon = LoadIcon(ghInstance, MAKEINTRESOURCE(IDI_COMPUTER));
        giComputer = ImageList_AddIcon(himl, hicon);

        hicon = LoadIcon(ghInstance, MAKEINTRESOURCE(IDI_HUB));
        giHub = ImageList_AddIcon(himl, hicon);

        hicon = LoadIcon(ghInstance, MAKEINTRESOURCE(IDI_NODEVICE));
        giNoDevice = ImageList_AddIcon(himl, hicon);

        TreeView_SetImageList(ghTreeWnd, himl, TVSIL_NORMAL);
    }
}

// CSearchDeviceDlg message handlers

BOOL CSearchDeviceDlg::OnDeviceChange(UINT uEvent, DWORD dwEventData)
{
    switch (uEvent)
    {
        case DBT_DEVICEARRIVAL:
        case DBT_DEVICEREMOVECOMPLETE:
            RefreshTree();
            break;
    }

    return TRUE;
}

void CSearchDeviceDlg::OnNMDblclkTreeDevices(NMHDR *pNMHDR, LRESULT *pResult)
{
    OnOK();

    *pResult = 0;
}
