#pragma once

#include <string>

#include <windows.h>
#include <d3d11.h>
#include <dxgi1_2.h>
#include <mftransform.h>
#include <wrl/client.h>

#include <flutter/event_sink.h>
#include <flutter/encodable_value.h>

#include <atomic>
#include <memory>
#include <mutex>
#include <thread>
#include <vector>

class ScreenEncoder {
 public:
  ScreenEncoder();
  ~ScreenEncoder();

  void Start(std::unique_ptr<flutter::EventSink<flutter::EncodableValue>> sink);
  void Stop();
  void ForceKeyframe();

 private:
  void CaptureLoop();
  bool InitD3D();
  bool InitEncoder(UINT width, UINT height);
  bool CaptureFrame(std::vector<uint8_t>& bgra, UINT& width, UINT& height);
  void EncodeFrame(const std::vector<uint8_t>& bgra, UINT width, UINT height,
                   bool keyframe);
  void BgraToNv12(const uint8_t* bgra, std::vector<uint8_t>& nv12, UINT width,
                  UINT height);
  void SendEvent(std::vector<uint8_t> data);

  Microsoft::WRL::ComPtr<ID3D11Device>           d3d_dev_;
  Microsoft::WRL::ComPtr<ID3D11DeviceContext>    d3d_ctx_;
  Microsoft::WRL::ComPtr<IDXGIOutputDuplication> dupl_;
  Microsoft::WRL::ComPtr<ID3D11Texture2D>        staging_;
  Microsoft::WRL::ComPtr<IMFTransform>           encoder_;

  UINT      enc_width_   = 0;
  UINT      enc_height_  = 0;
  bool      config_sent_ = false;
  LONGLONG  sample_ts_   = 0;

  std::unique_ptr<flutter::EventSink<flutter::EncodableValue>> sink_;
  std::mutex sink_mu_;

  std::atomic<bool> running_   {false};
  std::atomic<bool> force_kf_  {false};
  std::thread       thread_;
};