#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;
}
Bytes encodePng(const Pixels &p)
{
    if (!p.width || !p.height || uint64_t(p.width) * p.height * 4 != p.bgra.size() ||
        p.bgra.size() > UINT_MAX || p.width > UINT_MAX / 4)
        throw std::runtime_error("Invalid PNG image dimensions");
    Com<IWICImagingFactory> factory;
    check(CoCreateInstance(CLSID_WICImagingFactory, nullptr, CLSCTX_INPROC_SERVER,
                           IID_PPV_ARGS(factory.out())));
    Com<IStream> stream;
    check(CreateStreamOnHGlobal(nullptr, TRUE, stream.out()));
    Com<IWICBitmapEncoder> encoder;
    check(factory->CreateEncoder(GUID_ContainerFormatPng, nullptr, encoder.out()));
    check(encoder->Initialize(stream.p, WICBitmapEncoderNoCache));
    Com<IWICBitmapFrameEncode> frame;
    check(encoder->CreateNewFrame(frame.out(), nullptr));
    check(frame->Initialize(nullptr));
    check(frame->SetSize(p.width, p.height));
    auto format = GUID_WICPixelFormat32bppBGRA;
    check(frame->SetPixelFormat(&format));
    if (format != GUID_WICPixelFormat32bppBGRA)
        throw std::runtime_error("Unsupported PNG pixel format");
    check(frame->WritePixels(p.height, p.width * 4, (UINT)p.bgra.size(),
                             const_cast<BYTE *>(p.bgra.data())));
    check(frame->Commit());
    check(encoder->Commit());
    STATSTG stat{};
    check(stream->Stat(&stat, STATFLAG_NONAME));
    if (stat.cbSize.QuadPart > UINT_MAX)
        throw std::runtime_error("PNG image is too large");
    check(stream->Seek({}, STREAM_SEEK_SET, nullptr));
    Bytes bytes((size_t)stat.cbSize.QuadPart);
    ULONG read{};
    check(stream->Read(bytes.data(), (ULONG)bytes.size(), &read));
    if (read != bytes.size())
        throw std::runtime_error("Cannot read encoded PNG image");
    return bytes;
}
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;
}
std::filesystem::path tesseractExecutable(const Settings &s)
{
    if (!s.tesseractPath.empty())
    {
        std::filesystem::path path(s.tesseractPath);
        if (!path.is_absolute()) throw std::runtime_error("Select an absolute path to tesseract.exe");
        return path;
    }
    for (auto variable : {L"ProgramW6432", L"ProgramFiles", L"LOCALAPPDATA"})
    {
        wchar_t base[32768]{};
        auto length = GetEnvironmentVariableW(variable, base, std::size(base));
        if (!length || length >= std::size(base)) continue;
        auto root = std::filesystem::path(base);
        for (auto relative : {L"Tesseract-OCR/tesseract.exe", L"Programs/Tesseract-OCR/tesseract.exe"})
            if (std::filesystem::is_regular_file(root / relative)) return root / relative;
    }
    throw std::runtime_error("Tesseract was not found. Install the optional 64-bit Tesseract 5 package or select tesseract.exe in OCR options.");
}
std::wstring ocrProfile(const Settings &s)
{
    return s.useTesseract ? L"tesseract:" + lower(s.tesseractPath) + L":" + s.tesseractLanguages : L"windows";
}
namespace
{
struct TessLibrary
{
    HMODULE module{};
    explicit TessLibrary(const Settings &s)
    {
        auto exe = tesseractExecutable(s);
        if (!std::filesystem::is_regular_file(exe)) throw std::runtime_error("Tesseract executable does not exist");
        auto library = exe.parent_path() / L"libtesseract-5.dll";
        module = LoadLibraryExW(library.c_str(), nullptr,
                               LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR | LOAD_LIBRARY_SEARCH_DEFAULT_DIRS);
        if (!module) throw std::runtime_error("Cannot load Tesseract 5. Select a complete 64-bit installation with libtesseract-5.dll and its dependencies.");
    }
    ~TessLibrary() { if (module) FreeLibrary(module); }
    template<class T> T symbol(const char *name)
    {
        auto address = GetProcAddress(module, name);
        if (!address) throw std::runtime_error(std::string("Unsupported Tesseract library: missing ") + name);
        T result; static_assert(sizeof(result) == sizeof(address));
        memcpy(&result, &address, sizeof(result));
        return result;
    }
};
struct TessEngine
{
    TessLibrary library;
    void *handle{};
    void (*destroy)(void *){};
    explicit TessEngine(const Settings &s) : library(s)
    {
        auto folder = tesseractExecutable(s).parent_path() / L"tessdata";
        size_t start = 0;
        do
        {
            auto end = s.tesseractLanguages.find(L'+', start);
            auto language = s.tesseractLanguages.substr(start, end == std::wstring::npos ? end : end - start);
            if (!std::filesystem::is_regular_file(folder / (language + L".traineddata")))
                throw std::runtime_error("Tesseract language is not installed: " + utf8(language));
            if (end == std::wstring::npos) break;
            start = end + 1;
        } while (true);
        auto create = library.symbol<void *(*)()>("TessBaseAPICreate");
        destroy = library.symbol<void (*)(void *)>("TessBaseAPIDelete");
        auto init = library.symbol<int (*)(void *, const char *, const char *)>("TessBaseAPIInit3");
        handle = create();
        if (!handle) throw std::runtime_error("Cannot create Tesseract engine");
        if (init(handle, utf8(folder.wstring()).c_str(), utf8(s.tesseractLanguages).c_str()) != 0)
        {
            destroy(handle); handle = nullptr;
            throw std::runtime_error("Cannot initialize the selected Tesseract languages");
        }
    }
    ~TessEngine() { if (handle) destroy(handle); }
};
}
void validateTesseract(const Settings &s)
{
    if (s.useTesseract) { TessEngine engine(s); }
}
std::wstring recognizeText(const Clip &clip, const Settings &s)
{
    if (!s.useTesseract) return recognizeText(clip);
    auto apartmentResult = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
    if (FAILED(apartmentResult) && apartmentResult != RPC_E_CHANGED_MODE) check(apartmentResult);
    struct Apartment { bool owned; ~Apartment() { if (owned) CoUninitialize(); } } apartment{SUCCEEDED(apartmentResult)};
    // Load the optional C API dynamically: Windows OCR users need no Tesseract DLLs.
    TessEngine engine(s);
    auto &lib = engine.library;
    auto image = lib.symbol<void (*)(void *, const unsigned char *, int, int, int, int)>("TessBaseAPISetImage");
    auto mode = lib.symbol<void (*)(void *, int)>("TessBaseAPISetPageSegMode");
    auto recognize = lib.symbol<int (*)(void *, void *)>("TessBaseAPIRecognize");
    auto getText = lib.symbol<char *(*)(void *)>("TessBaseAPIGetUTF8Text");
    auto freeText = lib.symbol<void (*)(const char *)>("TessDeleteText");
    auto createMonitor = lib.symbol<void *(*)()>("TessMonitorCreate");
    auto freeMonitor = lib.symbol<void (*)(void *)>("TessMonitorDelete");
    auto deadline = lib.symbol<void (*)(void *, int)>("TessMonitorSetDeadlineMSecs");
    std::unique_ptr<void, void (*)(void *)> monitor(createMonitor(), freeMonitor);
    if (!monitor) throw std::runtime_error("Cannot create Tesseract progress monitor");
    auto pixels = decodeImage(clip);
    Bytes gray(size_t(pixels.width) * pixels.height);
    for (size_t i = 0; i < gray.size(); ++i)
        gray[i] = BYTE((29 * pixels.bgra[4*i] + 150 * pixels.bgra[4*i+1] + 77 * pixels.bgra[4*i+2]) >> 8);
    image(engine.handle, gray.data(), pixels.width, pixels.height, 1, pixels.width);
    mode(engine.handle, 11); // Sparse text: screenshots often contain separate labels and paragraphs.
    deadline(monitor.get(), 30000);
    if (recognize(engine.handle, monitor.get()) != 0)
        throw std::runtime_error("Tesseract recognition failed or exceeded 30 seconds");
    std::unique_ptr<char, void (*)(const char *)> text(getText(engine.handle), freeText);
    if (!text) throw std::runtime_error("Tesseract returned no result");
    return wide(text.get());
}
} // namespace ks