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

bool loadXISF(const LibXISF::ByteArray &data, std::shared_ptr<RawImage> &rawImage);
bool loadFITS(const LibXISF::ByteArray &data, std::shared_ptr<RawImage> &rawImage);

void RawImageToHTBITMAP(std::shared_ptr<RawImage> &rawImage, HBITMAP *hbmp, UINT thumbSize)
{
    rawImage->calcStats();

    DWORD thre = 20;
    DWORD dataSize = 4;
    HRESULT hr = HRESULT_FROM_WIN32(RegGetValueW(HKEY_CURRENT_USER, L"SOFTWARE\\nou\\Tenmon\\settings", L"thumbnailstretchthreshold", RRF_RT_DWORD, NULL, &thre, &dataSize));

    float thref = 0.1f;
    if(hr == S_OK)
        thref = thre / 100.0f;

    if(rawImage->imageStats().m_median[0] < rawImage->norm() * thref)
    {
        //OutputDebugStringA("Stretch image");
        MTFParam params = rawImage->calcMTFParams();
        rawImage->applySTF(params);
    }

    UINT w = rawImage->width();
    UINT h = rawImage->height();

    UINT cw = thumbSize;
    UINT ch = thumbSize;
    if (w > h)
        ch = h * thumbSize / w;
    else
        cw = w * thumbSize / h;


    rawImage->resize(cw, ch);
    rawImage->convertToType(RawImage::UINT8);

    BITMAPINFO bmi = {};
    bmi.bmiHeader.biSize = sizeof(bmi.bmiHeader);
    bmi.bmiHeader.biWidth = cw;
    bmi.bmiHeader.biHeight = -static_cast<LONG>(ch);
    bmi.bmiHeader.biPlanes = 1;
    bmi.bmiHeader.biBitCount = 32;
    bmi.bmiHeader.biCompression = BI_RGB;
    UINT lw = cw * 4;

    BYTE *pBits;
    HBITMAP bmp = CreateDIBSection(NULL, &bmi, DIB_RGB_COLORS, reinterpret_cast<void **>(&pBits), NULL, 0);

    const unsigned char *p = (const unsigned char*)rawImage->data();
    const unsigned short *ps = (const unsigned short*)rawImage->data();
    if(rawImage->channels() == 1)
    {
        for(UINT y = 0; y < ch; y++)
        {
            for(UINT x = 0; x < cw; x++)
            {
                pBits[(y * lw) + x * 4 + 0] = p[y * cw + x];
                pBits[(y * lw) + x * 4 + 1] = p[y * cw + x];
                pBits[(y * lw) + x * 4 + 2] = p[y * cw + x];
                pBits[(y * lw) + x * 4 + 3] = 255;
            }
        }
    }
    else
    {
        for(UINT y = 0; y < ch; y++)
        {
            for(UINT x = 0; x < cw; x++)
            {
                pBits[(y * lw) + x * 4 + 0] = p[y * cw * 4 + x * 4 + 2];
                pBits[(y * lw) + x * 4 + 1] = p[y * cw * 4 + x * 4 + 1];
                pBits[(y * lw) + x * 4 + 2] = p[y * cw * 4 + x * 4 + 0];
                pBits[(y * lw) + x * 4 + 3] = 255;
            }
        }
    }

    *hbmp = bmp;
}

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);

    std::shared_ptr<RawImage> rawImage;
    if(data[0] == 'X' && data[1] == 'I' && data[2] == 'S' && data[3] == 'F')
    {
        if(!loadXISF(data, rawImage))
            return E_FAIL;
    }
    else
    {
        if(!loadFITS(data, rawImage))
            return E_FAIL;
    }

    RawImageToHTBITMAP(rawImage, phbmp, cx);
    return S_OK;
}