#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 <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) {
  winrt::init_apartment();

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

        try {
          const auto png_bytes = ExtractBytes(*call.arguments());
          const auto text = RecognizePng(png_bytes);
          result->Success(flutter::EncodableValue(text));
        } catch (const winrt::hresult_error& error) {
          result->Error("ocr_failed", winrt::to_string(error.message()));
        } catch (...) {
          result->Error("ocr_failed", "Unknown OCR error");
        }
      });
}