#include <shlwapi.h>
#include <thumbcache.h> // For IThumbnailProvider.
#include <new>
#include "libxisf.h"

bool loadXISF(const LibXISF::ByteArray &data, HBITMAP *hbmp, UINT thumbSize);
bool loadFITS(const LibXISF::ByteArray &data, HBITMAP *hbmp, UINT thumbSize);

class TenmonThumbProvider : public IInitializeWithStream,
                             public IThumbnailProvider
{
public:
    TenmonThumbProvider() : _cRef(1), _pStream(NULL)
    {
    }

    virtual ~TenmonThumbProvider()
    {
        if (_pStream)
        {
            _pStream->Release();
        }
    }

    // IUnknown
    IFACEMETHODIMP QueryInterface(REFIID riid, void **ppv)
    {
        static const QITAB qit[] =
        {
            QITABENT(TenmonThumbProvider, IInitializeWithStream),
            QITABENT(TenmonThumbProvider, IThumbnailProvider),
            { 0 },
        };
        return QISearch(this, qit, riid, ppv);
    }

    IFACEMETHODIMP_(ULONG) AddRef()
    {
        return InterlockedIncrement(&_cRef);
    }

    IFACEMETHODIMP_(ULONG) Release()
    {
        ULONG cRef = InterlockedDecrement(&_cRef);
        if (!cRef)
        {
            delete this;
        }
        return cRef;
    }

    // IInitializeWithStream
    IFACEMETHODIMP Initialize(IStream *pStream, DWORD grfMode);

    // IThumbnailProvider
    IFACEMETHODIMP GetThumbnail(UINT cx, HBITMAP *phbmp, WTS_ALPHATYPE *pdwAlpha);

private:

    long _cRef;
    IStream *_pStream;     // provided during initialization.
};

HRESULT TenmonThumbnailer_CreateInstance(REFIID riid, void **ppv)
{
    TenmonThumbProvider *pNew = new (std::nothrow) TenmonThumbProvider();
    HRESULT hr = pNew ? S_OK : E_OUTOFMEMORY;
    if (SUCCEEDED(hr))
    {
        hr = pNew->QueryInterface(riid, ppv);
        pNew->Release();
    }
    return hr;
}

// IInitializeWithStream
IFACEMETHODIMP TenmonThumbProvider::Initialize(IStream *pStream, DWORD)
{
    HRESULT hr = E_UNEXPECTED;  // can only be inited once
    if (_pStream == NULL)
    {
        // take a reference to the stream if we have not been inited yet
        hr = pStream->QueryInterface(&_pStream);
    }
    return hr;
}

// IThumbnailProvider
IFACEMETHODIMP TenmonThumbProvider::GetThumbnail(UINT cx, HBITMAP *phbmp, WTS_ALPHATYPE *pdwAlpha)
{
    LibXISF::ByteArray data;
    ULONG readSize = 0;
    ULONG read;
    data.resize(1024*1024);

    while(_pStream->Read(data.data() + readSize, data.size() - readSize, &read) == S_OK)
    {
        readSize += read;
        data.resize(data.size() + 1024*1024);
    }
    readSize += read;

    *pdwAlpha = WTSAT_RGB;

    data.resize(readSize);
    if(data[0] == 'X' && data[1] == 'I' && data[2] == 'S' && data[3] == 'F')
    {
        if(loadXISF(data, phbmp, cx))
            return S_OK;
        else
            return E_FAIL;
    }
    else
    {
        if(loadFITS(data, phbmp, cx))
            return S_OK;
        else
            return E_FAIL;
    }
    return E_FAIL;
}