//
// Aspia Project
// Copyright (C) 2016-2024 Dmitry Chapyshev <dmitry@aspia.ru>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.
//

#ifndef HOST_WIN_RDP_CLIENT_WINDOW_H
#define HOST_WIN_RDP_CLIENT_WINDOW_H

#include "base/desktop/geometry.h"
#include "base/waitable_timer.h"
// The following header was generated by Visual Studio. We had to check it in due to a bug in VS2013.
// See crbug.com/318952 for details.
#include "host/win/com_imported_mstscax.h"

#include <memory>
#include <string>

#include <atlbase.h>
#include <atlcom.h>
#include <atlctl.h>
#include <atldef.h>
#include <atlhost.h>
#include <atlsecurity.h>
#include <atlwin.h>
#include <atlapp.h>
#include <atlcrack.h>
#include <wrl/client.h>

#include <asio/ip/tcp.hpp>

namespace host {

// RdpClientWindow is used to establish a connection to the given RDP endpoint. It is a GUI window
// class that hosts Microsoft RDP ActiveX control, which takes care of handling RDP properly.
// RdpClientWindow must be used only on a UI thread.
class RdpClientWindow final
    : public CWindowImpl<RdpClientWindow, CWindow, CFrameWinTraits>,
      public IDispEventImpl<1, RdpClientWindow,
                            &__uuidof(mstsc::IMsTscAxEvents),
                            &__uuidof(mstsc::__MSTSCLib), 1, 0>
{
public:
    // Receives connect/disconnect notifications. The notifications can be delivered after
    // RdpClientWindow::Connect() returned success.
    //
    // RdpClientWindow guarantees that OnDisconnected() is the last notification the event handler
    // receives. OnDisconnected() is guaranteed to be called only once.
    class EventHandler
    {
    public:
        virtual ~EventHandler() = default;

        // Invoked when the RDP control has established a connection.
        virtual void onConnected() = 0;

        // Invoked when the RDP control has been disconnected from the RDP server. This includes
        // both graceful shutdown and any fatal error condition.
        //
        // Once RdpClientWindow::Connect() returns success the owner of the |RdpClientWindow| object
        // must keep it alive until OnDisconnected() is called.
        //
        // OnDisconnected() should not delete |RdpClientWindow| object directly. Instead it should
        // post a task to delete the object. The ActiveX code expects the window be alive until the
        // currently handled window message is completely processed.
        virtual void onDisconnected() = 0;
    };

    DECLARE_WND_CLASS(L"RdpClientWindow")

    // Specifies the endpoint to connect to and passes the event handler pointer to be notified
    // about connection events.
    RdpClientWindow(const asio::ip::tcp::endpoint& server_endpoint,
                    const std::string& terminal_id,
                    std::shared_ptr<base::TaskRunner> ui_task_runner,
                    EventHandler* event_handler);
    ~RdpClientWindow() final;

    // Creates the window along with the ActiveX control and initiates the connection. |resolution|
    // specifies resolution of the screen. Returns false if an error occurs.
    bool connect(const base::Size& resolution);

    // Initiates shutdown of the connection. The caller must not delete |this| until it receives
    // OnDisconnected() notification.
    void disconnect();

    // Emulates pressing Ctrl+Alt+End combination that is translated to Secure Attention Sequence by
    // the ActiveX control.
    void injectSas();

    // Change the resolution of the desktop.
    void changeResolution(const base::Size& resolution);

private:
    typedef IDispEventImpl<1, RdpClientWindow,
                           &__uuidof(mstsc::IMsTscAxEvents),
                           &__uuidof(mstsc::__MSTSCLib), 1, 0> RdpEventsSink;

    // Handled window messages.
    BEGIN_MSG_MAP_EX(RdpClientWindow)
        MSG_WM_CLOSE(onClose)
        MSG_WM_CREATE(onCreate)
        MSG_WM_DESTROY(onDestroy)
    END_MSG_MAP()

    // Requests the RDP ActiveX control to close the connection gracefully.
    void onClose();

    // Creates the RDP ActiveX control, configures it, and initiates an RDP connection to
    // |server_endpoint_|.
    LRESULT onCreate(CREATESTRUCT* create_struct);

    // Releases the RDP ActiveX control interfaces.
    void onDestroy();

    BEGIN_SINK_MAP(RdpClientWindow)
        SINK_ENTRY_EX(1, __uuidof(mstsc::IMsTscAxEvents), 2,
                      &RdpClientWindow::onConnected)
        SINK_ENTRY_EX(1, __uuidof(mstsc::IMsTscAxEvents), 3,
                      &RdpClientWindow::onLoginComplete)
        SINK_ENTRY_EX(1, __uuidof(mstsc::IMsTscAxEvents), 4,
                      &RdpClientWindow::onDisconnected)
        SINK_ENTRY_EX(1, __uuidof(mstsc::IMsTscAxEvents), 10,
                      &RdpClientWindow::onFatalError)
        SINK_ENTRY_EX(1, __uuidof(mstsc::IMsTscAxEvents), 15,
                      &RdpClientWindow::onConfirmClose)
        SINK_ENTRY_EX(1, __uuidof(mstsc::IMsTscAxEvents), 18,
                      &RdpClientWindow::onAuthenticationWarningDisplayed)
        SINK_ENTRY_EX(1, __uuidof(mstsc::IMsTscAxEvents), 19,
                      &RdpClientWindow::onAuthenticationWarningDismissed)
    END_SINK_MAP()

    // mstsc::IMsTscAxEvents notifications.
    STDMETHOD(onAuthenticationWarningDisplayed)();
    STDMETHOD(onAuthenticationWarningDismissed)();
    STDMETHOD(onConnected)();
    STDMETHOD(onLoginComplete)();
    STDMETHOD(onDisconnected)(long reason);
    STDMETHOD(onFatalError)(long error_code);
    STDMETHOD(onConfirmClose)(VARIANT_BOOL* allow_close);

    int logOnCreateError(HRESULT error);

    // Wrappers for the event handler's methods that make sure that OnDisconnected() is the last
    // notification delivered and is delivered only once.
    void notifyConnected();
    void notifyDisconnected();

    // Updates the desktop using |screen_resolution_| if resizing is supported.
    HRESULT updateDesktopResolution();

    // Attempts to reapply the requested screen resolution. This method is used to workaround the
    // inconsistent behavior seen by the RDP API which can fail for an indeterminate amount of time
    // (i.e. a few seconds) after the user has logged into the session.  In practice retrying after
    // ~100ms will succeed, but we will use |apply_resolution_timer_| to try a few times if the machine
    // is under load.
    void reapplyDesktopResolution();

    // Invoked to report connect/disconnect events.
    EventHandler* event_handler_;

    // Contains the requested dimensions of the screen.
    base::Size screen_resolution_;

    // Used for applying resolution changes after a timeout.
    base::WaitableTimer apply_resolution_timer_;

    // Tracks the number of resolution change retries.
    int apply_resolution_attempts_ = 0;

    // The endpoint to connect to.
    asio::ip::tcp::endpoint server_endpoint_;

    // The terminal ID assigned to this connection.
    std::string terminal_id_;

    // Our RDP session always starts logged out (i.e. at the logon screen), some functions do not
    // work until the user has logged in so we want to track that event.  We don't support logging
    // out so this value should never revert from true to false.
    bool user_logged_in_ = false;

    // Interfaces exposed by the RDP ActiveX control.
    Microsoft::WRL::ComPtr<mstsc::IMsRdpClient> client_;
    Microsoft::WRL::ComPtr<mstsc::IMsRdpClient9> client_9_;
    Microsoft::WRL::ComPtr<mstsc::IMsRdpClientAdvancedSettings> client_settings_;

    // Used to cancel modal dialog boxes shown by the RDP control.
    class WindowHook;
    std::shared_ptr<WindowHook> window_activate_hook_;
};

} // namespace host

#endif // HOST_WIN_RDP_CLIENT_WINDOW_H
