# -*- 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 ElementIn(object):

    default_find_fn = 'find_element'

    def __init__(self, element, locator, find_fn=None):
        self.element = element
        self.locator = locator
        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.locator)
        except:
            return False


class ElementsIn(ElementIn):

    default_find_fn = 'find_elements'


class SeleniumExtensions(object):

    @classmethod
    def wait_elem(cls, driver, locator, inside_el=None, wait_time=10):
        if inside_el:
            return WebDriverWait(driver, wait_time).until(ElementIn(inside_el, locator))
        return WebDriverWait(driver, wait_time).until(
            EC.presence_of_element_located(locator)
        )

    @classmethod
    def wait_elems(cls, driver, locator, inside_el=None, wait_time=10):
        if inside_el:
            return WebDriverWait(driver, wait_time).until(ElementsIn(inside_el, locator))
        return WebDriverWait(driver, wait_time).until(
            EC.presence_of_all_elements_located(locator)
        )