# -*- coding: utf-8 -*-

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC


class elem_in(object):
  default_find_fn = 'find_element_by_xpath'
  def __init__(self, element, xpath, find_fn=None):
    self.element = element
    self.xpath = xpath
    self.find_fn = self.default_find_fn if find_fn is None else find_fn

  def __call__(self, driver):
    try:
        return getattr(self.element, self.find_fn)(self.xpath)
    except:
        return False

class elems_in(elem_in):
  default_find_fn = 'find_elements_by_xpath'


class SeleniumExtensions(object):

    @classmethod
    def wait_elem(cls, driver, xpath, inside_el=None, wait_time=10):
        if inside_el:
            return WebDriverWait(driver, wait_time).until(elem_in(inside_el, xpath))
        if xpath[:2] != '//':
            raise Exception('XPath in wait_elem must start with //')
        return WebDriverWait(driver, wait_time).until(
            EC.presence_of_element_located(
                (By.XPATH, xpath)
            )
        )

    @classmethod
    def wait_elems(cls, driver, xpath, inside_el=None, wait_time=10):
        if inside_el:
            return WebDriverWait(driver, wait_time).until(elems_in(inside_el, xpath))
        if xpath[:2] != '//':
            raise Exception('XPath in wait_elems must start with //')
        return WebDriverWait(driver, wait_time).until(
            EC.presence_of_all_elements_located(
                (By.XPATH, xpath)
            )
        )