#pragma once

#include "ProcessLauncher.h"
#include <windows.h>
#include <tsvirtualchannels.h>

class ScriptHookPlugin final : public IWTSPlugin, public IWTSListenerCallback
{
public:
    ScriptHookPlugin();

    HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void** ppvObject) override;
    ULONG STDMETHODCALLTYPE AddRef() override;
    ULONG STDMETHODCALLTYPE Release() override;

    // IWTSPlugin
    HRESULT STDMETHODCALLTYPE Initialize(IWTSVirtualChannelManager* channelManager) override;
    HRESULT STDMETHODCALLTYPE Connected() override;
    HRESULT STDMETHODCALLTYPE Disconnected(DWORD disconnectCode) override;
    HRESULT STDMETHODCALLTYPE Terminated() override;

    // IWTSListenerCallback
    HRESULT STDMETHODCALLTYPE OnNewChannelConnection(
        IWTSVirtualChannel* channel,
        BSTR data,
        BOOL* accept,
        IWTSVirtualChannelCallback** callback) override;

private:
    ~ScriptHookPlugin();
    void ReleaseDvcObjects();

    LONG refCount_ = 1;
    hook::ProcessLauncher connectLauncher_;
    IWTSVirtualChannelManager* channelManager_ = nullptr;
    IWTSListener* listener_ = nullptr;
};