// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.

#include "precomp.h"

#define BURN_UITHREAD_CLASS_WINDOW L"WixBurnMessageWindow"


// structs

struct UITHREAD_CONTEXT
{
    HANDLE hInitializedEvent;
    HINSTANCE hInstance;
    BURN_ENGINE_STATE* pEngineState;
};

struct UITHREAD_INFO
{
    BOOL fElevatedEngine;
    BURN_USER_EXPERIENCE* pUserExperience;
};


// internal function declarations

static DWORD WINAPI ThreadProc(
    __in LPVOID pvContext
    );

static LRESULT CALLBACK WndProc(
    __in HWND hWnd,
    __in UINT uMsg,
    __in WPARAM wParam,
    __in LPARAM lParam
    );


// function definitions

HRESULT UiCreateMessageWindow(
    __in HINSTANCE hInstance,
    __in BURN_ENGINE_STATE* pEngineState
    )
{
    HRESULT hr = S_OK;
    HANDLE rgWaitHandles[2] = { };
    UITHREAD_CONTEXT context = { };

    // Create event to signal after the UI thread / window is initialized.
    rgWaitHandles[0] = ::CreateEventW(NULL, TRUE, FALSE, NULL);
    ExitOnNullWithLastError(rgWaitHandles[0], hr, "Failed to create initialization event.");

    // Pass necessary information to create the window.
    context.hInitializedEvent = rgWaitHandles[0];
    context.hInstance = hInstance;
    context.pEngineState = pEngineState;

    // Create our separate UI thread.
    rgWaitHandles[1] = ::CreateThread(NULL, 0, ThreadProc, &context, 0, NULL);
    ExitOnNullWithLastError(rgWaitHandles[1], hr, "Failed to create the UI thread.");

    // Wait for either the thread to be initialized or the window to exit / fail prematurely.
    ::WaitForMultipleObjects(countof(rgWaitHandles), rgWaitHandles, FALSE, INFINITE);

    pEngineState->hMessageWindowThread = rgWaitHandles[1];
    rgWaitHandles[1] = NULL;

LExit:
    ReleaseHandle(rgWaitHandles[1]);
    ReleaseHandle(rgWaitHandles[0]);

    return hr;
}

void UiCloseMessageWindow(
    __in BURN_ENGINE_STATE* pEngineState
    )
{
    if (::IsWindow(pEngineState->hMessageWindow))
    {
        ::PostMessageW(pEngineState->hMessageWindow, WM_CLOSE, 0, 0);

        // Give the window 15 seconds to close because if it stays open it can prevent
        // the engine from starting a reboot (should a reboot actually be necessary).
        ::WaitForSingleObject(pEngineState->hMessageWindowThread, 15 * 1000);
    }
}


// internal function definitions

static DWORD WINAPI ThreadProc(
    __in LPVOID pvContext
    )
{
    HRESULT hr = S_OK;
    UITHREAD_CONTEXT* pContext = static_cast<UITHREAD_CONTEXT*>(pvContext);
    UITHREAD_INFO info = { };

    WNDCLASSW wc = { };
    BOOL fRegistered = TRUE;
    HWND hWnd = NULL;

    BOOL fRet = FALSE;
    MSG msg = { };

    BURN_ENGINE_STATE* pEngineState = pContext->pEngineState;
    BOOL fElevatedEngine = BURN_MODE_ELEVATED == pContext->pEngineState->internalCommand.mode;

    // If elevated, set up the thread local storage to store the correct pipe to communicate logging.
    if (fElevatedEngine)
    {
        Assert(TLS_OUT_OF_INDEXES != pEngineState->dwElevatedLoggingTlsId);

        if (!::TlsSetValue(pEngineState->dwElevatedLoggingTlsId, pEngineState->companionConnection.hPipe))
        {
            // If the function failed we cannot write to the pipe so just terminate.
            ExitFunction1(hr = E_INVALIDSTATE);
        }
    }

    wc.lpfnWndProc = WndProc;
    wc.hInstance = pContext->hInstance;
    wc.lpszClassName = BURN_UITHREAD_CLASS_WINDOW;

    if (!::RegisterClassW(&wc))
    {
        ExitWithLastError(hr, "Failed to register window.");
    }

    fRegistered = TRUE;

    info.fElevatedEngine = fElevatedEngine;
    info.pUserExperience = &pEngineState->userExperience;

    // Create the window to handle reboots without activating it.
    hWnd = ::CreateWindowExW(WS_EX_NOACTIVATE, wc.lpszClassName, NULL, WS_POPUP, 0, 0, 0, 0, HWND_DESKTOP, NULL, pContext->hInstance, &info);
    ExitOnNullWithLastError(hWnd, hr, "Failed to create window.");

    ::ShowWindow(hWnd, SW_SHOWNA);

    // Persist the window handle and let the caller know we've initialized.
    pEngineState->hMessageWindow = hWnd;
    ::SetEvent(pContext->hInitializedEvent);

    // Pump messages until the window is closed.
    while (0 != (fRet = ::GetMessageW(&msg, NULL, 0, 0)))
    {
        if (-1 == fRet)
        {
            hr = E_UNEXPECTED;
            ExitOnFailure(hr, "Unexpected return value from message pump.");
        }
        else if (!::IsDialogMessageW(msg.hwnd, &msg))
        {
            ::TranslateMessage(&msg);
            ::DispatchMessageW(&msg);
        }
    }

LExit:
    if (fRegistered)
    {
        ::UnregisterClassW(BURN_UITHREAD_CLASS_WINDOW, pContext->hInstance);
    }

    return hr;
}

static LRESULT CALLBACK WndProc(
    __in HWND hWnd,
    __in UINT uMsg,
    __in WPARAM wParam,
    __in LPARAM lParam
    )
{
    switch (uMsg)
    {
    case WM_NCCREATE:
        {
        LPCREATESTRUCTW lpcs = reinterpret_cast<LPCREATESTRUCTW>(lParam);
        ::SetWindowLongPtrW(hWnd, GWLP_USERDATA, reinterpret_cast<LONG_PTR>(lpcs->lpCreateParams));
        break;
        }

    case WM_NCDESTROY:
        {
        LRESULT lRes = ::DefWindowProcW(hWnd, uMsg, wParam, lParam);
        ::SetWindowLongPtrW(hWnd, GWLP_USERDATA, 0);
        return lRes;
        }

    case WM_QUERYENDSESSION:
        {
        DWORD dwEndSession = static_cast<DWORD>(lParam);
        BOOL fCritical = ENDSESSION_CRITICAL & dwEndSession;
        BOOL fCancel = TRUE;
        BOOL fRet = FALSE;

        // Always block shutdown in the elevated process, but ask the BA in the non-elevated.
        UITHREAD_INFO* pInfo = reinterpret_cast<UITHREAD_INFO*>(::GetWindowLongPtrW(hWnd, GWLP_USERDATA));
        if (!pInfo->fElevatedEngine)
        {
            // TODO: instead of recommending canceling all non-critical shutdowns, maybe we should only recommend cancel
            //       when the engine is doing work?
            fCancel = !fCritical;
            // TODO: There's a race condition here where the BA may not have been loaded, or already was unloaded.
            UserExperienceOnSystemShutdown(pInfo->pUserExperience, dwEndSession, &fCancel);
        }

        fRet = !fCancel;
        LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN, LoggingBoolToString(fCritical), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(fRet));
        return fRet;
        }

    case WM_DESTROY:
        ::PostQuitMessage(0);
        return 0;
    }

    return ::DefWindowProcW(hWnd, uMsg, wParam, lParam);
}