// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
// PARTICULAR PURPOSE.
//
// Copyright (c) Microsoft Corporation. All rights reserved

#ifndef __FILTERBASE_H__
#define __FILTERBASE_H__

#include <strsafe.h>
#include <propkey.h>
#include <propsys.h>
#include <filter.h>
#include <filterr.h>

class CChunkValue
{
public:
    CChunkValue() : m_fIsValid(false), m_pszValue(nullptr)
    {
        PropVariantInit(&m_propVariant);
        Clear();
    }

    ~CChunkValue() { Clear(); };

    void Clear()
    {
        m_fIsValid = false;
        ZeroMemory(&m_chunk, sizeof(m_chunk));
        PropVariantClear(&m_propVariant);
        CoTaskMemFree(m_pszValue);
        m_pszValue = nullptr;
    }

    BOOL IsValid() { return m_fIsValid; }

    HRESULT GetValue(PROPVARIANT **ppPropVariant)
    {
        if (!ppPropVariant)
            return E_INVALIDARG;

        *ppPropVariant = nullptr;

        PROPVARIANT *pPropVariant = static_cast<PROPVARIANT*>(CoTaskMemAlloc(sizeof(PROPVARIANT)));
        if (!pPropVariant)
            return E_OUTOFMEMORY;

        HRESULT hr = PropVariantCopy(pPropVariant, &m_propVariant);
        if (SUCCEEDED(hr))
            *ppPropVariant = pPropVariant;
        else
            CoTaskMemFree(pPropVariant);

        return hr;
    }

    PWSTR GetString() { return m_pszValue; };

    HRESULT CopyChunk(STAT_CHUNK *pStatChunk)
    {
        if (!pStatChunk)
            return E_INVALIDARG;

        *pStatChunk = m_chunk;
        return S_OK;
    }

    CHUNKSTATE GetChunkType() { return m_chunk.flags; }

    HRESULT SetTextValue(REFPROPERTYKEY pkey, PCWSTR pszValue, CHUNKSTATE chunkType = CHUNK_VALUE,
                         LCID locale = 0, DWORD cwcLenSource = 0, DWORD cwcStartSource = 0,
                         CHUNK_BREAKTYPE chunkBreakType = CHUNK_NO_BREAK)
    {
        if (pszValue == nullptr)
            return E_INVALIDARG;

        HRESULT hr = SetChunk(pkey, chunkType, locale, cwcLenSource, cwcStartSource, chunkBreakType);
        if (FAILED(hr))
            return hr;

        size_t cch = wcslen(pszValue) + 1;
        PWSTR pszCoTaskValue = static_cast<PWSTR>(CoTaskMemAlloc(cch * sizeof(WCHAR)));
        if (!pszCoTaskValue)
            return E_OUTOFMEMORY;

        StringCchCopyW(pszCoTaskValue, cch, pszValue);
        m_fIsValid = true;
        if (chunkType == CHUNK_VALUE)
        {
            m_propVariant.vt = VT_LPWSTR;
            m_propVariant.pwszVal = pszCoTaskValue;
        }
        else
        {
            m_pszValue = pszCoTaskValue;
        }

        return hr;
    };

    HRESULT SetFileTimeValue(REFPROPERTYKEY pkey, FILETIME dtVal, CHUNKSTATE chunkType = CHUNK_VALUE,
                             LCID locale = 0, DWORD cwcLenSource = 0, DWORD cwcStartSource = 0,
                             CHUNK_BREAKTYPE chunkBreakType = CHUNK_NO_BREAK)
    {
        HRESULT hr = SetChunk(pkey, chunkType, locale, cwcLenSource, cwcStartSource, chunkBreakType);
        if (FAILED(hr))
            return hr;

        m_propVariant.vt = VT_FILETIME;
        m_propVariant.filetime = dtVal;
        m_fIsValid = true;

        return hr;
    };

protected:
    HRESULT SetChunk(REFPROPERTYKEY pkey, CHUNKSTATE chunkType=CHUNK_VALUE, LCID locale=0, DWORD cwcLenSource=0, DWORD cwcStartSource=0, CHUNK_BREAKTYPE chunkBreakType=CHUNK_NO_BREAK);

private:
    bool m_fIsValid;
    STAT_CHUNK  m_chunk;
    PROPVARIANT m_propVariant;
    PWSTR m_pszValue;

};

inline HRESULT CChunkValue::SetChunk(REFPROPERTYKEY pkey,
                                     CHUNKSTATE chunkType/*=CHUNK_VALUE*/,
                                     LCID locale /*=0*/,
                                     DWORD cwcLenSource /*=0*/,
                                     DWORD cwcStartSource /*=0*/,
                                     CHUNK_BREAKTYPE chunkBreakType /*= CHUNK_NO_BREAK */)
{
    Clear();

    m_chunk.attribute.psProperty.ulKind = PRSPEC_PROPID;
    m_chunk.attribute.psProperty.propid = pkey.pid;
    m_chunk.attribute.guidPropSet = pkey.fmtid;
    m_chunk.flags = chunkType;
    m_chunk.locale = locale;
    m_chunk.cwcLenSource = cwcLenSource;
    m_chunk.cwcStartSource = cwcStartSource;
    m_chunk.breakType = chunkBreakType;

    return S_OK;
}

class CFilterBase : public IFilter, public IInitializeWithStream, public IPersistStream, public IPersistFile
{
public:
    // OnInit() is called when the IFilter is initialized (at the end of IFilter::Init)
    virtual HRESULT OnInit() = 0;

    // When GetNextChunkValue() is called you should fill in the ChunkValue by calling SetXXXValue() with the property.
    // example:  chunkValue.SetTextValue(PKYE_ItemName, L"foo bar");
    // return FILTER_E_END_OF_CHUNKS when there are no more chunks
    virtual HRESULT GetNextChunkValue(CChunkValue &chunkValue) = 0;

protected:
    inline DWORD GetChunkId() const { return m_dwChunkId; }

public:
    CFilterBase(long *plRefCount) : m_lRef(1), m_plModuleRef(plRefCount),
        m_dwChunkId(0), m_iText(0), m_pStream(nullptr) {
        InterlockedIncrement(m_plModuleRef);
    }

    virtual ~CFilterBase() {
        if (m_pStream)
            m_pStream->Release();
        InterlockedDecrement(m_plModuleRef);
    }

    // IUnknown
    IFACEMETHODIMP QueryInterface(REFIID riid, void **ppv) {
        static const QITAB qit[] = {
            QITABENT(CFilterBase, IPersistStream),
            QITABENT(CFilterBase, IPersistFile),
            QITABENTMULTI(CFilterBase, IPersist, IPersistStream),
            QITABENT(CFilterBase, IInitializeWithStream),
            QITABENT(CFilterBase, IFilter),
            { 0 }
        };
        return QISearch(this, qit, riid, ppv);
    }
    IFACEMETHODIMP_(ULONG) AddRef() {
        return InterlockedIncrement(&m_lRef);
    }
    IFACEMETHODIMP_(ULONG) Release() {
        long cRef = InterlockedDecrement(&m_lRef);
        if (cRef == 0)
            delete this;
        return cRef;
    }

    // IFilter
    IFACEMETHODIMP Init(ULONG grfFlags, ULONG cAttributes, const FULLPROPSPEC *aAttributes, ULONG *pFlags) {
        UNUSED(grfFlags);
        if (cAttributes > 0 && !aAttributes)
            return E_INVALIDARG;

        m_dwChunkId = 0;
        m_iText = 0;
        m_currentChunk.Clear();
        if (pFlags)
            *pFlags = 0;

        return OnInit();
    }

    IFACEMETHODIMP GetChunk(STAT_CHUNK *pStat) {
        m_currentChunk.Clear();
        HRESULT hr = GetNextChunkValue(m_currentChunk);
        if (hr != S_OK)
            return hr;
        if (!m_currentChunk.IsValid())
             return E_INVALIDARG;

        m_iText = 0;
        m_currentChunk.CopyChunk(pStat);
        pStat->idChunk = ++m_dwChunkId;
        if (pStat->flags == CHUNK_TEXT)
            pStat->idChunkSource = pStat->idChunk;

        return hr;
    }

    IFACEMETHODIMP GetText(ULONG *pcwcBuffer, WCHAR *awcBuffer) {
        if (!pcwcBuffer || !*pcwcBuffer)
            return E_INVALIDARG;
        if (!m_currentChunk.IsValid())
            return FILTER_E_NO_MORE_TEXT;
        if (m_currentChunk.GetChunkType() != CHUNK_TEXT)
            return FILTER_E_NO_TEXT;

        ULONG cchTotal = static_cast<ULONG>(wcslen(m_currentChunk.GetString()));
        ULONG cchLeft = cchTotal - m_iText;
        ULONG cchToCopy = std::min(*pcwcBuffer - 1, cchLeft);

        if (!cchToCopy)
            return FILTER_E_NO_MORE_TEXT;

        PCWSTR psz = m_currentChunk.GetString() + m_iText;
        StringCchCopyNW(awcBuffer, *pcwcBuffer, psz, cchToCopy);
        awcBuffer[cchToCopy] = '\0';

        *pcwcBuffer = cchToCopy;
        m_iText += cchToCopy;
        cchLeft -= cchToCopy;

        if (!cchLeft)
            return FILTER_S_LAST_TEXT;

        return S_OK;
    }

    IFACEMETHODIMP GetValue(PROPVARIANT **ppPropValue) {
        if (!m_currentChunk.IsValid())
            return FILTER_E_NO_MORE_VALUES;
        if (m_currentChunk.GetChunkType() != CHUNK_VALUE)
            return FILTER_E_NO_VALUES;
        if (ppPropValue == nullptr)
            return E_INVALIDARG;

        HRESULT hr = m_currentChunk.GetValue(ppPropValue);
        m_currentChunk.Clear();

        return hr;
    }
    IFACEMETHODIMP BindRegion(FILTERREGION, REFIID, void **) { return E_NOTIMPL; }

    // IInitializeWithStream
    IFACEMETHODIMP Initialize(IStream *pStm, DWORD grfMode) {
        UNUSED(grfMode);
        if (m_pStream)
            m_pStream->Release();
        m_pStream = pStm;
        if (!m_pStream)
            return E_INVALIDARG;
        m_pStream->AddRef();
        return S_OK;
    };

    // IPersistStream
    IFACEMETHODIMP IsDirty() { return E_NOTIMPL; }
    IFACEMETHODIMP Load(IStream *pStm) { return Initialize(pStm, 0); }
    IFACEMETHODIMP Save(IStream *pStm, BOOL fClearDirty) { UNUSED(pStm); UNUSED(fClearDirty); return E_NOTIMPL; }
    IFACEMETHODIMP GetSizeMax(ULARGE_INTEGER *pcbSize) { UNUSED(pcbSize); return E_NOTIMPL; }

    // IPersistFile (for compatibility with older Windows Desktop Search versions and ifilttst.exe)
    IFACEMETHODIMP Load(LPCOLESTR pszFileName, DWORD dwMode) {
        UNUSED(dwMode);
        HANDLE hFile = CreateFileW(pszFileName, GENERIC_READ, FILE_SHARE_READ, nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
        if (hFile == INVALID_HANDLE_VALUE)
            return E_INVALIDARG;
        DWORD size = GetFileSize(hFile, nullptr), read;
        HGLOBAL data = GlobalAlloc(GMEM_MOVEABLE, size);
        if (!data) {
            CloseHandle(hFile);
            return E_OUTOFMEMORY;
        }
        BOOL ok = ReadFile(hFile, GlobalLock(data), size, &read, nullptr);
        GlobalUnlock(data);
        CloseHandle(hFile);

        IStream *pStm;
        if (!ok || FAILED(CreateStreamOnHGlobal(data, TRUE, &pStm))) {
            GlobalFree(data);
            return E_FAIL;
        }
        HRESULT res = Initialize(pStm, 0);
        pStm->Release();
        return res;
    }
    IFACEMETHODIMP Save(LPCOLESTR pszFileName, BOOL bRemember) { UNUSED(pszFileName); UNUSED(bRemember); return E_NOTIMPL; }
    IFACEMETHODIMP SaveCompleted(LPCOLESTR pszFileName) { UNUSED(pszFileName); return E_NOTIMPL; }
    IFACEMETHODIMP GetCurFile(LPOLESTR *ppszFileName) { UNUSED(ppszFileName); return E_NOTIMPL; }

protected:
    IStream*                    m_pStream;

private:
    long                        m_lRef, * m_plModuleRef;

    DWORD                       m_dwChunkId;
    DWORD                       m_iText;

    CChunkValue                 m_currentChunk;
};

#endif
