#include "ocr_channel.h"

#include <flutter/encodable_value.h>
#include <flutter/flutter_engine.h>
#include <flutter/method_channel.h>
#include <flutter/standard_method_codec.h>

#include <memory>
#include <string>
#include <thread>
#include <vector>

#include <winrt/Windows.Foundation.h>
#include <winrt/Windows.Graphics.Imaging.h>
#include <winrt/Windows.Media.Ocr.h>
#include <winrt/Windows.Storage.Streams.h>
#include <winrt/base.h>

namespace {

std::vector<uint8_t> ExtractBytes(const flutter::EncodableValue& value) {
  if (const auto* bytes = std::get_if<std::vector<uint8_t>>(&value)) {
    return *bytes;
  }
  if (const auto* list = std::get_if<flutter::EncodableList>(&value)) {
    std::vector<uint8_t> out;
    out.reserve(list->size());
    for (const auto& item : *list) {
      if (const auto* i = std::get_if<int32_t>(&item)) {
        out.push_back(static_cast<uint8_t>(*i));
      }
    }
    return out;
  }
  return {};
}

std::string RecognizePng(const std::vector<uint8_t>& bytes) {
  if (bytes.empty()) {
    return {};
  }

  using namespace winrt;
  using namespace Windows::Storage::Streams;
  using namespace Windows::Graphics::Imaging;
  using namespace Windows::Media::Ocr;

  InMemoryRandomAccessStream stream;
  {
    DataWriter writer{stream};
    writer.WriteBytes(bytes);
    writer.StoreAsync().get();
    writer.DetachStream();
  }
  stream.Seek(0);

  BitmapDecoder decoder = BitmapDecoder::CreateAsync(stream).get();
  SoftwareBitmap bitmap = decoder.GetSoftwareBitmapAsync().get();

  OcrEngine engine = OcrEngine::TryCreateFromUserProfileLanguages();
  if (!engine) {
    return {};
  }

  OcrResult result = engine.RecognizeAsync(bitmap).get();
  return winrt::to_string(result.Text());
}

std::unique_ptr<flutter::MethodChannel<flutter::EncodableValue>> g_ocr_channel;

}  // namespace

void RegisterOcrChannel(flutter::FlutterEngine* engine) {
  // NOTE: do NOT call winrt::init_apartment() here. This runs on the Flutter
  // platform thread, which main.cpp has already initialized as an STA via
  // CoInitializeEx(COINIT_APARTMENTTHREADED). Calling init_apartment() (which
  // defaults to multi-threaded) would throw RPC_E_CHANGED_MODE and crash the
  // app before the window appears. The WinRT OCR work runs on its own
  // dedicated MTA worker thread per call instead (see below).
  g_ocr_channel =
      std::make_unique<flutter::MethodChannel<flutter::EncodableValue>>(
          engine->messenger(), "badnote/ocr",
          &flutter::StandardMethodCodec::GetInstance());

  g_ocr_channel->SetMethodCallHandler(
      [](const flutter::MethodCall<flutter::EncodableValue>& call,
         std::unique_ptr<flutter::MethodResult<flutter::EncodableValue>>
             result) {
        if (call.method_name() != "recognize") {
          result->NotImplemented();
          return;
        }

        if (!call.arguments()) {
          result->Error("invalid_args", "Expected PNG byte data");
          return;
        }

        const auto png_bytes = ExtractBytes(*call.arguments());

        // Run the WinRT OCR on a dedicated multi-threaded-apartment worker so
        // its blocking .get() calls behave correctly and the platform thread's
        // STA apartment is left untouched. We join before returning so the
        // MethodResult is completed on the platform thread.
        std::string text;
        std::string error;
        std::thread worker([&]() {
          try {
            winrt::init_apartment(winrt::apartment_type::multi_threaded);
            text = RecognizePng(png_bytes);
            winrt::uninit_apartment();
          } catch (const winrt::hresult_error& e) {
            error = winrt::to_string(e.message());
          } catch (...) {
            error = "Unknown OCR error";
          }
        });
        worker.join();

        if (!error.empty()) {
          result->Error("ocr_failed", error);
        } else {
          result->Success(flutter::EncodableValue(text));
        }
      });
}