# Copyright 2021 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging
import time

import pythoncom
import pywintypes
import win32api
import win32com
import win32com.client
import win32con
import win32gui
import win32process
import winerror


class _MessageQueueAttacher(object):
  """Wrapper class for message queue attachment."""

  def __enter__(self):
    """Attaches the current thread to the foreground window's message queue.

    This is an old and well known exploit used to bypass Windows Focus rules:
    http://www.google.com/search?q=attachthreadinput+setforegroundwindow
    """
    self._active_thread_id = 0
    active_hwnd = win32gui.GetForegroundWindow()
    if not active_hwnd:
      logging.warning('No active window is found.')
      return
    current_thread_id = win32api.GetCurrentThreadId()
    active_thread_id, _ = win32process.GetWindowThreadProcessId(active_hwnd)
    win32process.AttachThreadInput(current_thread_id, active_thread_id, 1)
    logging.info('Attached current thread input %s to active thread: %s',
        current_thread_id, active_thread_id)
    self._active_thread_id = active_thread_id

  def __exit__(self, unused_type, unused_value, unused_traceback):
    """Detaches the current thread from the active thread's message queue."""
    if not self._active_thread_id:
      return
    current_thread_id = win32api.GetCurrentThreadId()
    win32process.AttachThreadInput(current_thread_id, self._active_thread_id, 0)
    logging.info('Detached current thread input %s from thread: %s',
        current_thread_id, self._active_thread_id)


def SetForegroundWindow(hwnd):
  """Brings the given window to foreground.

  Args:
    hwnd: Handle of the window to bring to the foreground.
  """
  with _MessageQueueAttacher():
    return bool(win32gui.SetForegroundWindow(hwnd))


def FindWindowsWithText(parent, text_to_search):
  """Finds windows with given text.

  Args:
    parent: Handle to the parent window whose child windows are to be searched.
    text_to_search: Substring to search within Windows text, case-insensitive.

  Returns:
    A list of HWND that match the search condition.
  """

  class WindowFoundHandler(object):
    """Callback class for window enumeration."""

    def __init__(self, text_to_search):
      self.result = []
      self._text_to_search = text_to_search

    def Process(self, handle):
      """Callback function when enumerating a window.

      Args:
        handle: HWND to the enumerated window.
      """
      text = win32gui.GetWindowText(handle).lower()
      text_to_search = self._text_to_search.lower()

      if text_to_search in text:
        self.result.append(handle)

  def WinFoundCallback(hwnd, window_found_handler):
    window_found_handler.Process(hwnd)

  handler = WindowFoundHandler(text_to_search)
  try:
    win32gui.EnumChildWindows(parent, WinFoundCallback, handler)
  except pywintypes.error as e:
    logging.info('Error while searching [%s], error: [%s]', text_to_search, e)

  return handler.result


def FindWindowsWithTitle(title_to_search):
  """Finds windows with given title.

  Args:
    title_to_search: Window title substring to search, case-insensitive.

  Returns:
    A list of HWND that match the search condition.
  """
  desktop_handle = None
  return FindWindowsWithText(desktop_handle, title_to_search)


def FindWindow(title, class_name, parent=0, child_after=0):
  """Finds a window of a given title and class.

  Args:
    title: Title of the window to search.
    class_name: Class name of the window to search.
    parent: Handle to the parent window whose child windows are to be searched.
    child_after: HWND to a child window. Search begins with the next child
      window in the Z order.

  Returns:
    Handle of the found window, or 0 if not found.
  """
  hwnd = 0
  try:
    hwnd = win32gui.FindWindowEx(
        int(parent), int(child_after), class_name, title)
  except win32gui.error as err:
    if err[0] == winrror.ERROR_INVALID_WINDOW_HANDLE:  # Could be closed.
      pass
    elif err[0] != winrror.ERROR_FILE_NOT_FOUND:
      raise err
  if hwnd:
    win32gui.FlashWindow(hwnd, True)
  return hwnd


def FindWindowWithTitleAndText(title, text):
  """Checks if the any window has given title, and child window has the text.

  Args:
    title: Expected window title.
    text: Expected window text substring(in any child window), case-insensitive.

  Returns:
    A list how HWND that meets the search condition.
  """
  hwnds_title_matched = FindWindowsWithTitle(title)
  if not hwnds_title_matched:
    logging.info('No window has title: [%s].', title)

  hwnds = []
  for hwnd in hwnds_title_matched:
    if FindWindowsWithText(hwnd, text):
      hwnds.append(hwnd)
  return hwnds


def WaitForWindow(title, class_name, timeout=30):
  """Waits for window with given title and class to appear.

  Args:
    title: Windows title to search.
    class_name: Class name of the window to search.
    timeout: How long should wait before give up.

  Returns:
    A tuple of (HWND, title) of the found window, or (None, None) otherwise.
  """
  logging.info('ui.WaitForWindow("%s", "%s") for %s seconds', title,
               class_name, timeout)
  hwnd = None
  start = time.perf_counter()
  stop = start + int(timeout)

  while time.perf_counter() < stop:
    hwnd = FindWindow(title, class_name)
    if hwnd:
      elapsed = time.perf_counter() - start
      logging.info('Window ["%s"] found in %f seconds', title, elapsed)
      return (hwnd, title)
    logging.info('Window with title [%s] has not appeared yet.', title)
    time.sleep(0.5)

  logging.warning('WARNING: (%s,"%s") not found within %f seconds', title,
                  class_name, timeout)
  return (None, None)


def ClickButton(button_hwnd):
  """Clicks a button window by sending a BM_CLICK message.

  Per http://msdn2.microsoft.com/en-us/library/bb775985.aspx
  "If the button is in a dialog box and the dialog box is not active, the
  BM_CLICK message might fail. To ensure success in this situation, call
  the SetActiveWindow function to activate the dialog box before sending
  the BM_CLICK message to the button."

  Args:
    button_hwnd: HWND to the button to be clicked.
  """
  previous_active_window = win32gui.SetActiveWindow(button_hwnd)
  win32gui.PostMessage(button_hwnd, win32con.BM_CLICK, 0, 0)
  if previous_active_window:
    win32gui.SetActiveWindow(previous_active_window)


def ClickChildButtonWithText(parent_hwnd, button_text):
  """Clicks a child button window with given text.

  Args:
    parent_hwnd: HWND of the button parent.
    button_text: Button windows title.

  Returns:
    Whether button is clicked.
  """
  button_hwnd = FindWindow(button_text, 'Button', parent_hwnd)
  if button_hwnd:
    logging.debug('Found child button with: %s', button_text)
    ClickButton(button_hwnd)
    return True

  logging.debug('No button with [%s] found.', button_text)
  return False


def SendKeyToWindow(hwnd, key_to_press):
  """Sends a key press to the window.

  This is a blocking call until all keys are pressed.

  Args:
    hwnd: handle of  the window to press key.
    key_to_press: Actual key to send to window. eg.'{Enter}'. See `window_shell`
        documentation for key definitions.
  """
  try:
    window_shell = win32com.client.Dispatch('WScript.Shell')
    SetForegroundWindow(hwnd)
    window_shell.AppActivate(str(win32gui.GetForegroundWindow()))
    window_shell.SendKeys(key_to_press)
    logging.info('Sent %s to window %x.', key_to_press, hwnd)
  except pywintypes.error as err:
    logging.error('Failed to press key: %s', err)
    raise
  except pythoncom.com_error as err:
    logging.error('COM exception occurred: %s, is CoInitialize() called?', err)
    raise