#include "pch.h"
#include "tasklist_positions.h"

void Tasklist::update() {
  // Get HWND of the tasklist
  auto tasklist_hwnd = FindWindowA("Shell_TrayWnd", nullptr);
  if (!tasklist_hwnd) return;
  tasklist_hwnd = FindWindowExA(tasklist_hwnd, 0, "ReBarWindow32", nullptr);
  if (!tasklist_hwnd) return;
  tasklist_hwnd = FindWindowExA(tasklist_hwnd, 0, "MSTaskSwWClass", nullptr);
  if (!tasklist_hwnd) return;
  tasklist_hwnd = FindWindowExA(tasklist_hwnd, 0, "MSTaskListWClass", nullptr);
  if (!tasklist_hwnd) return;
  if (!automation) {
    winrt::check_hresult(CoCreateInstance(CLSID_CUIAutomation,
                                          nullptr,
                                          CLSCTX_INPROC_SERVER,
                                          IID_IUIAutomation,
                                          automation.put_void()));
    winrt::check_hresult(automation->CreateTrueCondition(true_condition.put()));
  }
  element = nullptr;
  winrt::check_hresult(automation->ElementFromHandle(tasklist_hwnd, element.put()));
}

bool Tasklist::update_buttons(std::vector<TasklistButton>& buttons) {  
  if (!automation || !element) {
    return false;
  }
  winrt::com_ptr<IUIAutomationElementArray> elements;
  if (element->FindAll(TreeScope_Children, true_condition.get(), elements.put()) < 0)
    return false;
  if (!elements)
    return false;
  int count;
  if (elements->get_Length(&count) < 0)
    return false;
  winrt::com_ptr<IUIAutomationElement> child;
  std::vector<TasklistButton> found_butttons;
  found_butttons.reserve(count);
  for (int i = 0; i < count; ++i) {
    child = nullptr;
    if (elements->GetElement(i, child.put()) < 0)
      return false;
    TasklistButton button;
    if (VARIANT var_rect; child->GetCurrentPropertyValue(UIA_BoundingRectanglePropertyId, &var_rect) >= 0) {
      if (var_rect.vt == (VT_R8 | VT_ARRAY)) {
        LONG pos;
        double value;
        pos = 0; SafeArrayGetElement(var_rect.parray, &pos, &value);
        button.x = (long)value;
        pos = 1; SafeArrayGetElement(var_rect.parray, &pos, &value);
        button.y = (long)value;
        pos = 2; SafeArrayGetElement(var_rect.parray, &pos, &value);
        button.width = (long)value;
        pos = 3; SafeArrayGetElement(var_rect.parray, &pos, &value);
        button.height = (long)value;
      }
      VariantClear(&var_rect);
    } else {
      return false;
    }
    if (BSTR automation_id; child->get_CurrentAutomationId(&automation_id) >= 0) {
      button.name = automation_id;
      SysFreeString(automation_id);
    }
    found_butttons.push_back(button);
  }
  // assign keynums
  buttons.clear();
  for (auto& button : found_butttons) {
    if (buttons.empty()) {
      button.keynum = 1;
      buttons.push_back(std::move(button));
    } else {
      if (button.x < buttons.back().x || button.y < buttons.back().y) // skip 2nd row
        break;
      if (button.name == buttons.back().name)
        continue; // skip buttons from the same app
      button.keynum = buttons.back().keynum + 1;
      buttons.push_back(std::move(button));
      if (buttons.back().keynum == 10)
        break; // no more than 10 buttons
    }
  }
  return true;
}

std::vector<TasklistButton> Tasklist::get_buttons() {
  std::vector<TasklistButton> buttons;
  update_buttons(buttons);
  return buttons;
}