#include "screen_capturer_wgc.h"

#include <Windows.h>
#include <d3d11_4.h>
#include <winrt/Windows.Foundation.Metadata.h>
#include <winrt/Windows.Graphics.Capture.h>

#include <iostream>

#include "libyuv.h"
#include "rd_log.h"

static std::vector<DisplayInfo> gs_display_list;

std::string WideToUtf8(const wchar_t *wideStr) {
  int size_needed = WideCharToMultiByte(CP_UTF8, 0, wideStr, -1, nullptr, 0,
                                        nullptr, nullptr);
  std::string result(size_needed, 0);
  WideCharToMultiByte(CP_UTF8, 0, wideStr, -1, &result[0], size_needed, nullptr,
                      nullptr);
  result.pop_back();
  return result;
}

BOOL WINAPI EnumMonitorProc(HMONITOR hmonitor, [[maybe_unused]] HDC hdc,
                            [[maybe_unused]] LPRECT lprc, LPARAM data) {
  MONITORINFOEX monitor_info_;
  monitor_info_.cbSize = sizeof(MONITORINFOEX);

  if (GetMonitorInfo(hmonitor, &monitor_info_)) {
    if (monitor_info_.dwFlags & MONITORINFOF_PRIMARY) {
      gs_display_list.insert(
          gs_display_list.begin(),
          {(void *)hmonitor, WideToUtf8(monitor_info_.szDevice),
           (monitor_info_.dwFlags & MONITORINFOF_PRIMARY) ? true : false,
           monitor_info_.rcMonitor.left, monitor_info_.rcMonitor.top,
           monitor_info_.rcMonitor.right, monitor_info_.rcMonitor.bottom});
      *(HMONITOR *)data = hmonitor;
    } else {
      gs_display_list.push_back(DisplayInfo(
          (void *)hmonitor, WideToUtf8(monitor_info_.szDevice),
          (monitor_info_.dwFlags & MONITORINFOF_PRIMARY) ? true : false,
          monitor_info_.rcMonitor.left, monitor_info_.rcMonitor.top,
          monitor_info_.rcMonitor.right, monitor_info_.rcMonitor.bottom));
    }
  }

  if (monitor_info_.dwFlags == DISPLAY_DEVICE_MIRRORING_DRIVER) return true;

  return true;
}

HMONITOR GetPrimaryMonitor() {
  HMONITOR hmonitor = nullptr;

  gs_display_list.clear();
  ::EnumDisplayMonitors(NULL, NULL, EnumMonitorProc, (LPARAM)&hmonitor);

  return hmonitor;
}

ScreenCapturerWgc::ScreenCapturerWgc() : monitor_(nullptr) {}

ScreenCapturerWgc::~ScreenCapturerWgc() {
  Stop();
  CleanUp();

  if (nv12_frame_) {
    delete nv12_frame_;
    nv12_frame_ = nullptr;
  }

  if (nv12_frame_scaled_) {
    delete nv12_frame_scaled_;
    nv12_frame_scaled_ = nullptr;
  }
}

bool ScreenCapturerWgc::IsWgcSupported() {
  try {
    /* no contract for IGraphicsCaptureItemInterop, verify 10.0.18362.0 */
    return winrt::Windows::Foundation::Metadata::ApiInformation::
        IsApiContractPresent(L"Windows.Foundation.UniversalApiContract", 8);
  } catch (const winrt::hresult_error &) {
    return false;
  } catch (...) {
    return false;
  }
}

int ScreenCapturerWgc::Init(const int fps, cb_desktop_data cb) {
  int error = 0;
  if (inited_ == true) return error;

  // nv12_frame_ = new unsigned char[rect.right * rect.bottom * 3 / 2];
  // nv12_frame_scaled_ = new unsigned char[1280 * 720 * 3 / 2];

  fps_ = fps;

  on_data_ = cb;

  if (!IsWgcSupported()) {
    LOG_ERROR("WGC not supported");
    error = 2;
    return error;
  }

  monitor_ = GetPrimaryMonitor();

  display_info_list_ = gs_display_list;

  if (display_info_list_.empty()) {
    LOG_ERROR("No display found");
    return -1;
  }

  for (int i = 0; i < display_info_list_.size(); i++) {
    const auto &display = display_info_list_[i];
    LOG_INFO(
        "index: {}, display name: {}, is primary: {}, bounds: ({}, {}) - "
        "({}, {})",
        i, display.name, (display.is_primary ? "yes" : "no"), display.left,
        display.top, display.right, display.bottom);

    sessions_.push_back(
        {std::make_unique<WgcSessionImpl>(i), false, false, false});
    sessions_.back().session_->RegisterObserver(this);
    error = sessions_.back().session_->Initialize((HMONITOR)display.handle);
    if (error != 0) {
      return error;
    }
    sessions_[i].inited_ = true;
    inited_ = true;
  }

  LOG_INFO("Default on monitor {}:{}", monitor_index_,
           display_info_list_[monitor_index_].name);

  return 0;
}

int ScreenCapturerWgc::Destroy() { return 0; }

int ScreenCapturerWgc::Start() {
  if (running_ == true) {
    LOG_ERROR("Screen capturer already running");
    return 0;
  }

  if (inited_ == false) {
    LOG_ERROR("Screen capturer not inited");
    return 4;
  }

  for (int i = 0; i < sessions_.size(); i++) {
    if (sessions_[i].inited_ == false) {
      LOG_ERROR("Session {} not inited", i);
      continue;
    }

    if (sessions_[i].running_) {
      LOG_ERROR("Session {} is already running", i);
    } else {
      sessions_[i].session_->Start();

      if (i != 0) {
        sessions_[i].session_->Pause();
        sessions_[i].paused_ = true;
      }
      sessions_[i].running_ = true;
    }
    running_ = true;
  }

  return 0;
}

int ScreenCapturerWgc::Pause(int monitor_index) {
  if (monitor_index >= sessions_.size() || monitor_index < 0) {
    LOG_ERROR("Invalid session index: {}", monitor_index);
    return -1;
  }

  if (!sessions_[monitor_index].paused_) {
    sessions_[monitor_index].session_->Pause();
    sessions_[monitor_index].paused_ = true;
    LOG_INFO("Pausing session {}", monitor_index);
  }
  return 0;
}

int ScreenCapturerWgc::Resume(int monitor_index) {
  if (monitor_index >= sessions_.size() || monitor_index < 0) {
    LOG_ERROR("Invalid session index: {}", monitor_index);
    return -1;
  }

  if (sessions_[monitor_index].paused_) {
    sessions_[monitor_index].session_->Resume();
    sessions_[monitor_index].paused_ = false;
    LOG_INFO("Resuming session {}", monitor_index);
  }
  return 0;
}

int ScreenCapturerWgc::Stop() {
  for (int i = 0; i < sessions_.size(); i++) {
    if (sessions_[i].running_) {
      sessions_[i].session_->Stop();
      sessions_[i].running_ = false;
    }
  }
  running_ = false;

  return 0;
}

int ScreenCapturerWgc::SwitchTo(int monitor_index) {
  if (monitor_index_ == monitor_index) {
    LOG_INFO("Already on monitor {}:{}", monitor_index_ + 1,
             display_info_list_[monitor_index_].name);
    return 0;
  }

  if (monitor_index >= display_info_list_.size()) {
    LOG_ERROR("Invalid monitor index: {}", monitor_index);
    return -1;
  }

  if (!sessions_[monitor_index].inited_) {
    LOG_ERROR("Monitor {} not inited", monitor_index);
    return -1;
  }

  Pause(monitor_index_);

  monitor_index_ = monitor_index;
  LOG_INFO("Switching to monitor {}:{}", monitor_index_,
           display_info_list_[monitor_index_].name);

  Resume(monitor_index);

  return 0;
}

void ScreenCapturerWgc::OnFrame(const WgcSession::wgc_session_frame &frame,
                                int id) {
  if (on_data_) {
    if (!nv12_frame_) {
      nv12_frame_ = new unsigned char[frame.width * frame.height * 3 / 2];
    }

    libyuv::ARGBToNV12((const uint8_t *)frame.data, frame.width * 4,
                       (uint8_t *)nv12_frame_, frame.width,
                       (uint8_t *)(nv12_frame_ + frame.width * frame.height),
                       frame.width, frame.width, frame.height);

    on_data_(nv12_frame_, frame.width * frame.height * 3 / 2, frame.width,
             frame.height, display_info_list_[id].name.c_str());
  }
}

void ScreenCapturerWgc::CleanUp() {
  if (inited_) {
    for (auto &session : sessions_) {
      if (session.session_) {
        session.session_->Stop();
      }
    }
    sessions_.clear();
  }
}