/* $Id: VBoxUsbPwr.cpp $ */
/** @file
 * USB Power state Handling
 */
/*
 * Copyright (C) 2011 Oracle Corporation
 *
 * This file is part of VirtualBox Open Source Edition (OSE), as
 * available from http://www.virtualbox.org. This file is free software;
 * you can redistribute it and/or modify it under the terms of the GNU
 * General Public License (GPL) as published by the Free Software
 * Foundation, in version 2 as it comes in the "COPYING" file of the
 * VirtualBox OSE distribution. VirtualBox OSE is distributed in the
 * hope that it will be useful, but WITHOUT ANY WARRANTY of any kind.
 */

#include "VBoxUsbCmn.h"

#include <iprt/assert.h>

DECLHIDDEN(VOID) vboxUsbPwrStateInit(PVBOXUSBDEV_EXT pDevExt)
{
    POWER_STATE PowerState;
    PowerState.SystemState = PowerSystemWorking;
    PowerState.DeviceState = PowerDeviceD0;
    PoSetPowerState(pDevExt->pFDO, DevicePowerState, PowerState);
    pDevExt->DdiState.PwrState.PowerState = PowerState;
    pDevExt->DdiState.PwrState.PowerDownLevel = PowerDeviceUnspecified;
}

static NTSTATUS vboxUsbPwrMnDefault(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    NTSTATUS Status;
    PoStartNextPowerIrp(pIrp);
    IoSkipCurrentIrpStackLocation(pIrp);
    Status = PoCallDriver(pDevExt->pLowerDO, pIrp);
    Assert(NT_SUCCESS(Status) || Status == STATUS_NOT_SUPPORTED);
    vboxUsbDdiStateRelease(pDevExt);
    return Status;
}

static NTSTATUS vboxUsbPwrMnPowerSequence(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    Assert(0);
    return vboxUsbPwrMnDefault(pDevExt, pIrp);
}

typedef struct VBOXUSB_PWRDEV_CTX
{
    PVBOXUSBDEV_EXT pDevExt;
    PIRP pIrp;
} VBOXUSB_PWRDEV_CTX, *PVBOXUSB_PWRDEV_CTX;

static VOID vboxUsbPwrIoDeviceCompletion(IN PDEVICE_OBJECT pDeviceObject,
                    IN UCHAR MinorFunction,
                    IN POWER_STATE PowerState,
                    IN PVOID pvContext,
                    IN PIO_STATUS_BLOCK pIoStatus)
{
    PVBOXUSB_PWRDEV_CTX pDevCtx = (PVBOXUSB_PWRDEV_CTX)pvContext;
    PVBOXUSBDEV_EXT pDevExt = pDevCtx->pDevExt;
    PIRP pIrp = pDevCtx->pIrp;
    pIrp->IoStatus.Status = pIoStatus->Status;
    pIrp->IoStatus.Information = 0;

    PoStartNextPowerIrp(pIrp);
    IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    vboxUsbDdiStateRelease(pDevExt);

    vboxUsbMemFree(pDevCtx);
}

static NTSTATUS vboxUsbPwrIoRequestDev(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);
    POWER_STATE PwrState;
    PwrState.SystemState = pSl->Parameters.Power.State.SystemState;
    PwrState.DeviceState = pDevExt->DdiState.DevCaps.DeviceState[PwrState.SystemState];

    NTSTATUS Status = STATUS_INSUFFICIENT_RESOURCES;
    PVBOXUSB_PWRDEV_CTX pDevCtx = (PVBOXUSB_PWRDEV_CTX)vboxUsbMemAlloc(sizeof (*pDevCtx));
    Assert(pDevCtx);
    if (pDevCtx)
    {
        pDevCtx->pDevExt = pDevExt;
        pDevCtx->pIrp = pIrp;

        Status = PoRequestPowerIrp(pDevExt->pPDO, pSl->MinorFunction, PwrState,
                        vboxUsbPwrIoDeviceCompletion, pDevCtx, NULL);
        Assert(NT_SUCCESS(Status));
        if (NT_SUCCESS(Status))
        {
            return STATUS_MORE_PROCESSING_REQUIRED;
        }

        vboxUsbMemFree(pDevCtx);
    }

    PoStartNextPowerIrp(pIrp);
    pIrp->IoStatus.Status = Status;
    pIrp->IoStatus.Information = 0;
    vboxUsbDdiStateRelease(pDevExt);

    /* the "real" Status is stored in pIrp->IoStatus.Status,
     * return success here to complete the Io */
    return STATUS_SUCCESS;
}

static NTSTATUS vboxUsbPwrIoPostSysCompletion(IN PDEVICE_OBJECT pDevObj, IN PIRP pIrp, IN PVOID pvContext)
{
    PVBOXUSBDEV_EXT pDevExt = (PVBOXUSBDEV_EXT)pvContext;
    NTSTATUS Status = pIrp->IoStatus.Status;
    Assert(Status == STATUS_SUCCESS);
    if (NT_SUCCESS(Status))
    {
        PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);
        switch (pSl->MinorFunction)
        {
            case IRP_MN_SET_POWER:
            {
                pDevExt->DdiState.PwrState.PowerState.SystemState = pSl->Parameters.Power.State.SystemState;
                break;
            }
            default:
            {
                break;
            }
        }

        return vboxUsbPwrIoRequestDev(pDevExt, pIrp);
    }

    PoStartNextPowerIrp(pIrp);
    vboxUsbDdiStateRelease(pDevExt);
    return STATUS_SUCCESS;
}

static NTSTATUS vboxUsbPwrIoPostSys(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    IoMarkIrpPending(pIrp);
    IoCopyCurrentIrpStackLocationToNext(pIrp);
    IoSetCompletionRoutine(pIrp, vboxUsbPwrIoPostSysCompletion, pDevExt, TRUE, TRUE, TRUE);
    NTSTATUS Status = PoCallDriver(pDevExt->pLowerDO, pIrp);
    Assert(NT_SUCCESS(Status));
    return STATUS_PENDING;
}

static NTSTATUS vboxUsbPwrQueryPowerSys(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);
    SYSTEM_POWER_STATE enmSysPState = pSl->Parameters.Power.State.SystemState;

    return vboxUsbPwrIoPostSys(pDevExt, pIrp);
}

static NTSTATUS vboxUsbPwrIoPostDevCompletion(IN PDEVICE_OBJECT pDevObj, IN PIRP pIrp, IN PVOID pvContext)
{
    PVBOXUSBDEV_EXT pDevExt = (PVBOXUSBDEV_EXT)pvContext;

    if (pIrp->PendingReturned)
    {
        IoMarkIrpPending(pIrp);
    }

    NTSTATUS Status = pIrp->IoStatus.Status;
    Assert(Status == STATUS_SUCCESS);
    if (NT_SUCCESS(Status))
    {
        PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);
        switch (pSl->MinorFunction)
        {
            case IRP_MN_SET_POWER:
            {
                pDevExt->DdiState.PwrState.PowerState.DeviceState = pSl->Parameters.Power.State.DeviceState;
                PoSetPowerState(pDevExt->pFDO, DevicePowerState, pSl->Parameters.Power.State);
                break;
            }
            default:
            {
                break;
            }
        }
    }

    PoStartNextPowerIrp(pIrp);
    vboxUsbDdiStateRelease(pDevExt);
    return STATUS_SUCCESS;
}

static NTSTATUS vboxUsbPwrIoPostDev(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    IoMarkIrpPending(pIrp);
    IoCopyCurrentIrpStackLocationToNext(pIrp);
    IoSetCompletionRoutine(pIrp, vboxUsbPwrIoPostDevCompletion, pDevExt, TRUE, TRUE, TRUE);
    NTSTATUS Status = PoCallDriver(pDevExt->pLowerDO, pIrp);
    Assert(NT_SUCCESS(Status));
    return STATUS_PENDING;
}

typedef struct VBOXUSB_IOASYNC_CTX
{
    PIO_WORKITEM pWrkItem;
    PIRP pIrp;
} VBOXUSB_IOASYNC_CTX, *PVBOXUSB_IOASYNC_CTX;

static VOID vboxUsbPwrIoWaitCompletionAndPostAsyncWorker(IN PDEVICE_OBJECT pDeviceObject, IN PVOID pvContext)
{
    PVBOXUSBDEV_EXT pDevExt = (PVBOXUSBDEV_EXT)pDeviceObject->DeviceExtension;
    PVBOXUSB_IOASYNC_CTX pCtx = (PVBOXUSB_IOASYNC_CTX)pvContext;
    PIRP pIrp = pCtx->pIrp;

    vboxUsbPwrIoPostDev(pDevExt, pIrp);

    IoFreeWorkItem(pCtx->pWrkItem);
    vboxUsbMemFree(pCtx);
}

static NTSTATUS vboxUsbPwrIoWaitCompletionAndPostAsync(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    NTSTATUS Status = STATUS_INSUFFICIENT_RESOURCES;
    PVBOXUSB_IOASYNC_CTX pCtx = (PVBOXUSB_IOASYNC_CTX)vboxUsbMemAlloc(sizeof (*pCtx));
    Assert(pCtx);
    if (pCtx)
    {
        PIO_WORKITEM pWrkItem = IoAllocateWorkItem(pDevExt->pFDO);
        Assert(pWrkItem);
        if (pWrkItem)
        {
            pCtx->pWrkItem = pWrkItem;
            pCtx->pIrp = pIrp;
            IoMarkIrpPending(pIrp);
            IoQueueWorkItem(pWrkItem, vboxUsbPwrIoWaitCompletionAndPostAsyncWorker, DelayedWorkQueue, pCtx);
            return STATUS_PENDING;
        }
        vboxUsbMemFree(pCtx);
    }
    return Status;
}

static NTSTATUS vboxUsbPwrQueryPowerDev(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);
    DEVICE_POWER_STATE enmDevPState = pSl->Parameters.Power.State.DeviceState;
    NTSTATUS Status = STATUS_SUCCESS;

    if (enmDevPState >= pDevExt->DdiState.PwrState.PowerState.DeviceState)
    {
        Status = vboxUsbPwrIoWaitCompletionAndPostAsync(pDevExt, pIrp);
        Assert(NT_SUCCESS(Status));
        if (NT_SUCCESS(Status))
            return Status;
    }

    pIrp->IoStatus.Status = Status;
    pIrp->IoStatus.Information = 0;

    PoStartNextPowerIrp(pIrp);

    if (NT_SUCCESS(Status))
    {
        IoSkipCurrentIrpStackLocation(pIrp);
        Status = PoCallDriver(pDevExt->pLowerDO, pIrp);
    }
    else
    {
        IoCompleteRequest(pIrp, IO_NO_INCREMENT);
    }

    vboxUsbDdiStateRelease(pDevExt);

    return Status;
}

static NTSTATUS vboxUsbPwrMnQueryPower(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);

    switch (pSl->Parameters.Power.Type)
    {
        case SystemPowerState:
        {
            return vboxUsbPwrQueryPowerSys(pDevExt, pIrp);
        }
        case DevicePowerState:
        {
            return vboxUsbPwrQueryPowerDev(pDevExt, pIrp);
        }
        default:
        {
            Assert(0);
            return vboxUsbPwrMnDefault(pDevExt, pIrp);
        }

    }
    return vboxUsbPwrMnDefault(pDevExt, pIrp);
}

static NTSTATUS vboxUsbPwrSetPowerSys(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);
    SYSTEM_POWER_STATE enmSysPState = pSl->Parameters.Power.State.SystemState;

    return vboxUsbPwrIoPostSys(pDevExt, pIrp);
}

static NTSTATUS vboxUsbPwrSetPowerDev(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);
    DEVICE_POWER_STATE enmDevPState = pSl->Parameters.Power.State.DeviceState;
    DEVICE_POWER_STATE enmCurDevPState = pDevExt->DdiState.PwrState.PowerState.DeviceState;
    NTSTATUS Status = STATUS_SUCCESS;

    if (enmDevPState > enmCurDevPState && enmCurDevPState == PowerDeviceD0)
    {
        Status = vboxUsbPwrIoWaitCompletionAndPostAsync(pDevExt, pIrp);
        Assert(NT_SUCCESS(Status));
        if (NT_SUCCESS(Status))
            return Status;
    }

    PoStartNextPowerIrp(pIrp);

    if (NT_SUCCESS(Status))
    {
        IoCopyCurrentIrpStackLocationToNext(pIrp);
        IoSetCompletionRoutine(pIrp, vboxUsbPwrIoPostDevCompletion, pDevExt, TRUE, TRUE, TRUE);
        Status = PoCallDriver(pDevExt->pLowerDO, pIrp);
    }
    else
    {
        pIrp->IoStatus.Status = Status;
        pIrp->IoStatus.Information = 0;

        IoCompleteRequest(pIrp, IO_NO_INCREMENT);
        vboxUsbDdiStateRelease(pDevExt);
    }

    return Status;
}


static NTSTATUS vboxUsbPwrMnSetPower(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);

    switch (pSl->Parameters.Power.Type)
    {
        case SystemPowerState:
        {
            return vboxUsbPwrSetPowerSys(pDevExt, pIrp);
        }
        case DevicePowerState:
        {
            return vboxUsbPwrSetPowerDev(pDevExt, pIrp);
        }
        default:
        {
            Assert(0);
            return vboxUsbPwrMnDefault(pDevExt, pIrp);
        }

    }
    return vboxUsbPwrMnDefault(pDevExt, pIrp);
}

static NTSTATUS vboxUsbPwrMnWaitWake(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    AssertFailed();
    return vboxUsbPwrMnDefault(pDevExt, pIrp);
}


static NTSTATUS vboxUsbPwrDispatch(IN PVBOXUSBDEV_EXT pDevExt, IN PIRP pIrp)
{
    PIO_STACK_LOCATION pSl = IoGetCurrentIrpStackLocation(pIrp);

    switch (pSl->MinorFunction)
    {
        case IRP_MN_POWER_SEQUENCE:
        {
            return vboxUsbPwrMnPowerSequence(pDevExt, pIrp);
        }
        case IRP_MN_QUERY_POWER:
        {
            return vboxUsbPwrMnQueryPower(pDevExt, pIrp);
        }
        case IRP_MN_SET_POWER:
        {
            return vboxUsbPwrMnSetPower(pDevExt, pIrp);
        }
        case IRP_MN_WAIT_WAKE:
        {
            return vboxUsbPwrMnWaitWake(pDevExt, pIrp);
        }
        default:
        {
//            Assert(0);
            return vboxUsbPwrMnDefault(pDevExt, pIrp);
        }
    }
}

DECLHIDDEN(NTSTATUS) vboxUsbDispatchPower(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
    PVBOXUSBDEV_EXT pDevExt = (PVBOXUSBDEV_EXT)pDeviceObject->DeviceExtension;
    ENMVBOXUSB_PNPSTATE enmState = vboxUsbDdiStateRetainIfNotRemoved(pDevExt);
    switch (enmState)
    {
        case ENMVBOXUSB_PNPSTATE_REMOVED:
        {
            PoStartNextPowerIrp(pIrp);

            pIrp->IoStatus.Status = STATUS_DELETE_PENDING;
            pIrp->IoStatus.Information = 0;

            IoCompleteRequest(pIrp, IO_NO_INCREMENT);

            vboxUsbDdiStateRelease(pDevExt);

            return STATUS_DELETE_PENDING;
        }
        case ENMVBOXUSB_PNPSTATE_START_PENDING:
        {
            PoStartNextPowerIrp(pIrp);
            IoSkipCurrentIrpStackLocation(pIrp);

            vboxUsbDdiStateRelease(pDevExt);

            return PoCallDriver(pDevExt->pLowerDO, pIrp);
        }
        default:
        {
            return vboxUsbPwrDispatch(pDevExt, pIrp);
        }
    }
}
