#include "ocr.h"
#include <memory>
#include <roapi.h>
#include <wincodec.h>
#include <winstring.h>

namespace ks
{
namespace
{
template <class T> struct Com
{
    T *p{};
    ~Com()
    {
        if (p)
            p->Release();
    }
    T *operator->() const
    {
        return p;
    }
    T **out()
    {
        return &p;
    }
    Com() = default;
    Com(const Com &) = delete;
    Com &operator=(const Com &) = delete;
};
void check(HRESULT h)
{
    if (FAILED(h))
    {
        char text[80];
        snprintf(text, sizeof(text), "Windows image/OCR service failed (0x%08lx)", (unsigned long)h);
        throw std::runtime_error(text);
    }
}
struct String
{
    HSTRING value{};
    explicit String(const wchar_t *s)
    {
        check(WindowsCreateString(s, (UINT32)wcslen(s), &value));
    }
    ~String()
    {
        WindowsDeleteString(value);
    }
};
GUID guid(const wchar_t *text)
{
    GUID result{};
    check(IIDFromString(text, &result));
    return result;
}
template <class T> void factory(const wchar_t *name, const wchar_t *iid, Com<T> &out)
{
    String s(name);
    check(RoGetActivationFactory(s.value, guid(iid), (void **)out.out()));
}
} // namespace
// Minimal, strongly typed Windows Runtime ABI declarations. No CLR or C++/WinRT runtime is needed.
// Interface identities/method order follow Microsoft's Windows SDK metadata (Windows.Media.Ocr,
// Windows.Graphics.Imaging, Windows.Storage.Streams, Windows.Foundation).
struct Buffer : IInspectable
{
    virtual HRESULT STDMETHODCALLTYPE Capacity(UINT32 *) = 0;
    virtual HRESULT STDMETHODCALLTYPE Length(UINT32 *) = 0;
    virtual HRESULT STDMETHODCALLTYPE SetLength(UINT32) = 0;
};
struct BufferFactory : IInspectable
{
    virtual HRESULT STDMETHODCALLTYPE Create(UINT32, Buffer **) = 0;
};
struct BufferBytes : IUnknown
{
    virtual HRESULT STDMETHODCALLTYPE Data(BYTE **) = 0;
};
struct BitmapStatics : IInspectable
{
    virtual HRESULT STDMETHODCALLTYPE Copy(IInspectable *, IInspectable **) = 0;
    virtual HRESULT STDMETHODCALLTYPE Convert(IInspectable *, INT32, IInspectable **) = 0;
    virtual HRESULT STDMETHODCALLTYPE ConvertAlpha(IInspectable *, INT32, INT32, IInspectable **) = 0;
    virtual HRESULT STDMETHODCALLTYPE FromBuffer(Buffer *, INT32, INT32, INT32, IInspectable **) = 0;
    virtual HRESULT STDMETHODCALLTYPE FromBufferAlpha(Buffer *, INT32, INT32, INT32, INT32,
                                                      IInspectable **) = 0;
};
struct OcrResult : IInspectable
{
    virtual HRESULT STDMETHODCALLTYPE Lines(IInspectable **) = 0;
    virtual HRESULT STDMETHODCALLTYPE Angle(IInspectable **) = 0;
    virtual HRESULT STDMETHODCALLTYPE Text(HSTRING *) = 0;
};
struct OcrOperation : IInspectable
{
    virtual HRESULT STDMETHODCALLTYPE PutCompleted(IUnknown *) = 0;
    virtual HRESULT STDMETHODCALLTYPE GetCompleted(IUnknown **) = 0;
    virtual HRESULT STDMETHODCALLTYPE Results(OcrResult **) = 0;
};
struct AsyncInfo : IInspectable
{
    virtual HRESULT STDMETHODCALLTYPE Id(UINT32 *) = 0;
    virtual HRESULT STDMETHODCALLTYPE Status(INT32 *) = 0;
    virtual HRESULT STDMETHODCALLTYPE ErrorCode(HRESULT *) = 0;
    virtual HRESULT STDMETHODCALLTYPE Cancel() = 0;
    virtual HRESULT STDMETHODCALLTYPE Close() = 0;
};
struct OcrEngine : IInspectable
{
    virtual HRESULT STDMETHODCALLTYPE Recognize(IInspectable *, OcrOperation **) = 0;
    virtual HRESULT STDMETHODCALLTYPE Language(IInspectable **) = 0;
};
struct OcrStatics : IInspectable
{
    virtual HRESULT STDMETHODCALLTYPE MaxDimension(UINT32 *) = 0;
    virtual HRESULT STDMETHODCALLTYPE Languages(IInspectable **) = 0;
    virtual HRESULT STDMETHODCALLTYPE Supported(IInspectable *, boolean *) = 0;
    virtual HRESULT STDMETHODCALLTYPE FromLanguage(IInspectable *, OcrEngine **) = 0;
    virtual HRESULT STDMETHODCALLTYPE FromProfile(OcrEngine **) = 0;
};
Pixels decodeImage(const Clip &clip, UINT maxWidth, UINT maxHeight)
{
    Bytes encoded;
    for (const auto &f : clip.formats)
        if (f.name == L"PNG")
        {
            encoded = f.data;
            break;
        }
    if (encoded.empty())
        for (const auto &f : clip.formats)
            if ((f.id == CF_DIB || f.id == CF_DIBV5) && f.data.size() >= 40)
            {
                BITMAPINFOHEADER h;
                memcpy(&h, f.data.data(), 40);
                if (h.biSize < 40 || h.biSize > f.data.size())
                    continue;
                uint64_t colors = h.biClrUsed ? h.biClrUsed : (h.biBitCount <= 8 ? 1ull << h.biBitCount : 0);
                uint64_t offset = sizeof(BITMAPFILEHEADER) + h.biSize + 4 * colors +
                                  (h.biSize == 40 && h.biCompression == BI_BITFIELDS ? 12 : 0);
                if (offset > sizeof(BITMAPFILEHEADER) + f.data.size())
                    throw std::runtime_error("Invalid bitmap header");
                BITMAPFILEHEADER file{};
                file.bfType = 0x4d42;
                file.bfSize = (DWORD)(sizeof(file) + f.data.size());
                file.bfOffBits = (DWORD)offset;
                encoded.resize(sizeof(file));
                memcpy(encoded.data(), &file, sizeof(file));
                encoded.insert(encoded.end(), f.data.begin(), f.data.end());
                break;
            }
    if (encoded.empty())
        throw std::runtime_error("This clipboard item has no supported image");
    Com<IWICImagingFactory> wic;
    check(CoCreateInstance(CLSID_WICImagingFactory, nullptr, CLSCTX_INPROC_SERVER, IID_IWICImagingFactory,
                           (void **)wic.out()));
    Com<IWICStream> stream;
    check(wic->CreateStream(stream.out()));
    check(stream->InitializeFromMemory(encoded.data(), (DWORD)encoded.size()));
    Com<IWICBitmapDecoder> decoder;
    check(wic->CreateDecoderFromStream(stream.p, nullptr, WICDecodeMetadataCacheOnLoad, decoder.out()));
    Com<IWICBitmapFrameDecode> frame;
    check(decoder->GetFrame(0, frame.out()));
    Pixels pixels;
    check(frame->GetSize(&pixels.width, &pixels.height));
    if (!pixels.width || !pixels.height || uint64_t(pixels.width) * pixels.height > 32 * 1024 * 1024)
        throw std::runtime_error("Image is too large to decode safely");
    Com<IWICFormatConverter> converter;
    Com<IWICBitmapScaler> scaler;
    IWICBitmapSource *source = frame.p;
    if (maxWidth && maxHeight && (pixels.width > maxWidth || pixels.height > maxHeight))
    {
        const double ratio = std::min(double(maxWidth) / pixels.width, double(maxHeight) / pixels.height);
        pixels.width = std::max(1u, UINT(pixels.width * ratio));
        pixels.height = std::max(1u, UINT(pixels.height * ratio));
        check(wic->CreateBitmapScaler(scaler.out()));
        check(scaler->Initialize(frame.p, pixels.width, pixels.height, WICBitmapInterpolationModeFant));
        source = scaler.p;
    }
    check(wic->CreateFormatConverter(converter.out()));
    check(converter->Initialize(source, GUID_WICPixelFormat32bppBGRA, WICBitmapDitherTypeNone, nullptr, 0,
                                WICBitmapPaletteTypeCustom));
    pixels.bgra.resize(size_t(pixels.width) * pixels.height * 4);
    check(converter->CopyPixels(nullptr, pixels.width * 4, (UINT)pixels.bgra.size(), pixels.bgra.data()));
    return pixels;
}
Format imageDib(const Pixels &p)
{
    BITMAPINFOHEADER h{};
    h.biSize = sizeof(h);
    h.biWidth = (LONG)p.width;
    h.biHeight = -(LONG)p.height;
    h.biPlanes = 1;
    h.biBitCount = 32;
    Format format;
    format.id = CF_DIB;
    format.data.resize(sizeof(h));
    memcpy(format.data.data(), &h, sizeof(h));
    format.data.insert(format.data.end(), p.bgra.begin(), p.bgra.end());
    return format;
}
std::wstring recognizeText(const Clip &clip)
{
    check(RoInitialize(RO_INIT_MULTITHREADED));
    struct Apartment
    {
        ~Apartment()
        {
            RoUninitialize();
        }
    } apartment;
    auto pixels = decodeImage(clip);
    Com<OcrStatics> ocr;
    factory(L"Windows.Media.Ocr.OcrEngine", L"{5bffa85a-3384-3540-9940-699120d428a8}", ocr);
    UINT32 max{};
    check(ocr->MaxDimension(&max));
    if (pixels.width > max || pixels.height > max)
        throw std::runtime_error("Image exceeds Windows OCR's maximum dimensions");
    Com<OcrEngine> engine;
    check(ocr->FromProfile(engine.out()));
    if (!engine.p)
        throw std::runtime_error("No Windows OCR language is installed for your profile");
    Com<BufferFactory> buffers;
    factory(L"Windows.Storage.Streams.Buffer", L"{71af914d-c10f-484b-bc50-14bc623b3a27}", buffers);
    Com<Buffer> buffer;
    check(buffers->Create((UINT32)pixels.bgra.size(), buffer.out()));
    check(buffer->SetLength((UINT32)pixels.bgra.size()));
    Com<BufferBytes> bytes;
    check(buffer->QueryInterface(guid(L"{905a0fef-bc53-11df-8c49-001e4fc686da}"), (void **)bytes.out()));
    BYTE *memory{};
    check(bytes->Data(&memory));
    memcpy(memory, pixels.bgra.data(), pixels.bgra.size());
    Com<BitmapStatics> bitmaps;
    factory(L"Windows.Graphics.Imaging.SoftwareBitmap", L"{df0385db-672f-4a9d-806e-c2442f343e86}", bitmaps);
    Com<IInspectable> bitmap;
    check(bitmaps->FromBufferAlpha(buffer.p, 87, (INT32)pixels.width, (INT32)pixels.height, 2, bitmap.out()));
    Com<OcrOperation> operation;
    check(engine->Recognize(bitmap.p, operation.out()));
    Com<AsyncInfo> info;
    check(operation->QueryInterface(guid(L"{00000036-0000-0000-c000-000000000046}"), (void **)info.out()));
    auto deadline = GetTickCount64() + 30000;
    INT32 status{};
    do
    {
        check(info->Status(&status));
        if (status == 0)
        {
            if (GetTickCount64() > deadline)
            {
                info->Cancel();
                throw std::runtime_error("Windows OCR timed out");
            }
            Sleep(20);
        }
    } while (status == 0);
    if (status != 1)
    {
        HRESULT failure{};
        check(info->ErrorCode(&failure));
        check(FAILED(failure) ? failure : E_ABORT);
    }
    Com<OcrResult> result;
    check(operation->Results(result.out()));
    HSTRING text{};
    check(result->Text(&text));
    UINT32 length{};
    const wchar_t *raw = WindowsGetStringRawBuffer(text, &length);
    std::wstring out(raw, length);
    WindowsDeleteString(text);
    return out;
}
} // namespace ks