from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.common.keys import Keys
from selenium.common.exceptions import TimeoutException, NoSuchElementException, ElementNotVisibleException
import time


class SeleniumBase:
    def __init__(self, driver):
        """
        初始化Selenium基础操作类
        :param driver: WebDriver实例
        """
        self.driver = driver
        self.wait = WebDriverWait(self.driver, 10)  # 默认等待10秒

    def set_implicit_wait(self, seconds):
        """设置隐式等待时间"""
        self.driver.implicitly_wait(seconds)

    def set_explicit_wait(self, seconds):
        """设置显式等待时间"""
        self.wait = WebDriverWait(self.driver, seconds)

    # ==================== 浏览器操作 ====================
    def open_url(self, url):
        """打开指定URL"""
        self.driver.get(url)

    def get_current_url(self):
        """获取当前URL"""
        return self.driver.current_url

    def get_page_title(self):
        """获取页面标题"""
        return self.driver.title

    def refresh_page(self):
        """刷新页面"""
        self.driver.refresh()

    def navigate_forward(self):
        """浏览器前进"""
        self.driver.forward()

    def navigate_back(self):
        """浏览器后退"""
        self.driver.back()

    def maximize_window(self):
        """最大化窗口"""
        self.driver.maximize_window()

    def set_window_size(self, width, height):
        """设置窗口大小"""
        self.driver.set_window_size(width, height)

    def get_window_size(self):
        """获取窗口大小"""
        return self.driver.get_window_size()

    def switch_to_frame(self, frame_reference):
        """切换到指定frame"""
        self.driver.switch_to.frame(frame_reference)

    def switch_to_default_content(self):
        """切换回默认内容"""
        self.driver.switch_to.default_content()

    def switch_to_window(self, window_handle):
        """切换到指定窗口"""
        self.driver.switch_to.window(window_handle)

    def switch_to_newest_window(self):
        """切换到最新打开的窗口"""
        handles = self.driver.window_handles
        self.driver.switch_to.window(handles[-1])
        return handles[-1]

    def close_current_window(self):
        """关闭当前窗口"""
        self.driver.close()

    def get_all_window_handles(self):
        """获取所有窗口句柄"""
        return self.driver.window_handles

    def get_current_window_handle(self):
        """获取当前窗口句柄"""
        return self.driver.current_window_handle

    def execute_script(self, script, *args):
        """执行JavaScript脚本"""
        return self.driver.execute_script(script, *args)

    def scroll_to_element(self, element):
        """滚动到指定元素"""
        self.execute_script("arguments[0].scrollIntoView();", element)

    def scroll_to_bottom(self):
        """滚动到页面底部"""
        self.execute_script("window.scrollTo(0, document.body.scrollHeight);")

    def scroll_to_top(self):
        """滚动到页面顶部"""
        self.execute_script("window.scrollTo(0, 0);")

    def take_screenshot(self, filename):
        """截取屏幕截图"""
        self.driver.save_screenshot(filename)

    # ==================== 元素定位 ====================
    def find_element(self, by, value, timeout=None):
        """
        查找元素（带等待）
        :param by: 定位方式 (By.ID, By.XPATH, By.CSS_SELECTOR, etc.)
        :param value: 定位值
        :param timeout: 超时时间，默认使用类初始化时的等待时间
        :return: WebElement
        """
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        try:
            return wait.until(EC.presence_of_element_located((by, value)))
        except TimeoutException:
            raise NoSuchElementException(f"元素未找到: {by}={value}")

    def find_elements(self, by, value, timeout=None):
        """
        查找多个元素（带等待）
        :param by: 定位方式
        :param value: 定位值
        :param timeout: 超时时间
        :return: 元素列表
        """
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        try:
            return wait.until(EC.presence_of_all_elements_located((by, value)))
        except TimeoutException:
            return []

    def find_visible_element(self, by, value, timeout=None):
        """
        查找可见元素
        :param by: 定位方式
        :param value: 定位值
        :param timeout: 超时时间
        :return: WebElement
        """
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        try:
            return wait.until(EC.visibility_of_element_located((by, value)))
        except TimeoutException:
            raise ElementNotVisibleException(f"元素不可见: {by}={value}")

    def find_clickable_element(self, by, value, timeout=None):
        """
        查找可点击元素
        :param by: 定位方式
        :param value: 定位值
        :param timeout: 超时时间
        :return: WebElement
        """
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        try:
            return wait.until(EC.element_to_be_clickable((by, value)))
        except TimeoutException:
            raise ElementNotVisibleException(f"元素不可点击: {by}={value}")

    # 快捷定位方法
    def by_id(self, id_value, timeout=None):
        """通过ID定位元素"""
        return self.find_element(By.ID, id_value, timeout)

    def by_name(self, name_value, timeout=None):
        """通过Name定位元素"""
        return self.find_element(By.NAME, name_value, timeout)

    def by_xpath(self, xpath_value, timeout=None):
        """通过XPath定位元素"""
        return self.find_element(By.XPATH, xpath_value, timeout)

    def by_css(self, css_value, timeout=None):
        """通过CSS选择器定位元素"""
        return self.find_element(By.CSS_SELECTOR, css_value, timeout)

    def by_class(self, class_value, timeout=None):
        """通过Class定位元素"""
        return self.find_element(By.CLASS_NAME, class_value, timeout)

    def by_link_text(self, link_text, timeout=None):
        """通过链接文本定位元素"""
        return self.find_element(By.LINK_TEXT, link_text, timeout)

    def by_partial_link_text(self, partial_link_text, timeout=None):
        """通过部分链接文本定位元素"""
        return self.find_element(By.PARTIAL_LINK_TEXT, partial_link_text, timeout)

    def by_tag(self, tag_name, timeout=None):
        """通过标签名定位元素"""
        return self.find_element(By.TAG_NAME, tag_name, timeout)

    # ==================== 元素操作 ====================
    def click(self, by, value, timeout=None):
        """点击元素"""
        element = self.find_clickable_element(by, value, timeout)
        element.click()

    def input_text(self, by, value, text, timeout=None):
        """输入文本"""
        element = self.find_element(by, value, timeout)
        element.clear()
        element.send_keys(text)

    def clear_text(self, by, value, timeout=None):
        """清除文本"""
        element = self.find_element(by, value, timeout)
        element.clear()

    def get_text(self, by, value, timeout=None):
        """获取元素文本"""
        element = self.find_element(by, value, timeout)
        return element.text

    def get_attribute(self, by, value, attribute_name, timeout=None):
        """获取元素属性"""
        element = self.find_element(by, value, timeout)
        return element.get_attribute(attribute_name)

    def is_displayed(self, by, value, timeout=None):
        """检查元素是否显示"""
        try:
            element = self.find_visible_element(by, value, timeout)
            return element.is_displayed()
        except (NoSuchElementException, ElementNotVisibleException):
            return False

    def is_enabled(self, by, value, timeout=None):
        """检查元素是否启用"""
        try:
            element = self.find_element(by, value, timeout)
            return element.is_enabled()
        except NoSuchElementException:
            return False

    def is_selected(self, by, value, timeout=None):
        """检查元素是否被选中"""
        try:
            element = self.find_element(by, value, timeout)
            return element.is_selected()
        except NoSuchElementException:
            return False

    def select_dropdown_by_value(self, by, value, option_value, timeout=None):
        """通过值选择下拉选项"""
        from selenium.webdriver.support.ui import Select
        element = self.find_element(by, value, timeout)
        select = Select(element)
        select.select_by_value(option_value)

    def select_dropdown_by_text(self, by, value, option_text, timeout=None):
        """通过文本选择下拉选项"""
        from selenium.webdriver.support.ui import Select
        element = self.find_element(by, value, timeout)
        select = Select(element)
        select.select_by_visible_text(option_text)

    def select_dropdown_by_index(self, by, value, index, timeout=None):
        """通过索引选择下拉选项"""
        from selenium.webdriver.support.ui import Select
        element = self.find_element(by, value, timeout)
        select = Select(element)
        select.select_by_index(index)

    def get_dropdown_options(self, by, value, timeout=None):
        """获取下拉选项"""
        from selenium.webdriver.support.ui import Select
        element = self.find_element(by, value, timeout)
        select = Select(element)
        return [option.text for option in select.options]

    def hover(self, by, value, timeout=None):
        """鼠标悬停"""
        element = self.find_element(by, value, timeout)
        ActionChains(self.driver).move_to_element(element).perform()

    def double_click(self, by, value, timeout=None):
        """双击元素"""
        element = self.find_element(by, value, timeout)
        ActionChains(self.driver).double_click(element).perform()

    def right_click(self, by, value, timeout=None):
        """右键点击元素"""
        element = self.find_element(by, value, timeout)
        ActionChains(self.driver).context_click(element).perform()

    def drag_and_drop(self, source_by, source_value, target_by, target_value, timeout=None):
        """拖放元素"""
        source_element = self.find_element(source_by, source_value, timeout)
        target_element = self.find_element(target_by, target_value, timeout)
        ActionChains(self.driver).drag_and_drop(source_element, target_element).perform()

    def press_key(self, by, value, key, timeout=None):
        """按下键盘按键"""
        element = self.find_element(by, value, timeout)
        element.send_keys(key)

    def press_enter(self, by, value, timeout=None):
        """按下回车键"""
        self.press_key(by, value, Keys.ENTER, timeout)

    def press_tab(self, by, value, timeout=None):
        """按下Tab键"""
        self.press_key(by, value, Keys.TAB, timeout)

    def press_escape(self, by, value, timeout=None):
        """按下Esc键"""
        self.press_key(by, value, Keys.ESCAPE, timeout)

    # ==================== 等待方法 ====================
    def wait_for_element_present(self, by, value, timeout=None):
        """等待元素出现"""
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        return wait.until(EC.presence_of_element_located((by, value)))

    def wait_for_element_visible(self, by, value, timeout=None):
        """等待元素可见"""
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        return wait.until(EC.visibility_of_element_located((by, value)))

    def wait_for_element_invisible(self, by, value, timeout=None):
        """等待元素不可见"""
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        return wait.until(EC.invisibility_of_element_located((by, value)))

    def wait_for_element_clickable(self, by, value, timeout=None):
        """等待元素可点击"""
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        return wait.until(EC.element_to_be_clickable((by, value)))

    def wait_for_text_present(self, by, value, text, timeout=None):
        """等待元素包含特定文本"""
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        return wait.until(EC.text_to_be_present_in_element((by, value), text))

    def wait_for_title_contains(self, text, timeout=None):
        """等待标题包含特定文本"""
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        return wait.until(EC.title_contains(text))

    def wait_for_title_is(self, text, timeout=None):
        """等待标题等于特定文本"""
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        return wait.until(EC.title_is(text))

    def wait_for_alert_present(self, timeout=None):
        """等待警告框出现"""
        wait = self.wait if timeout is None else WebDriverWait(self.driver, timeout)
        return wait.until(EC.alert_is_present())

    # ==================== 弹窗处理 ====================
    def accept_alert(self, timeout=None):
        """接受警告框"""
        alert = self.wait_for_alert_present(timeout)
        alert.accept()

    def dismiss_alert(self, timeout=None):
        """取消警告框"""
        alert = self.wait_for_alert_present(timeout)
        alert.dismiss()

    def get_alert_text(self, timeout=None):
        """获取警告框文本"""
        alert = self.wait_for_alert_present(timeout)
        return alert.text

    def input_alert_text(self, text, timeout=None):
        """向警告框输入文本"""
        alert = self.wait_for_alert_present(timeout)
        alert.send_keys(text)
        alert.accept()

    # ==================== Cookie操作 ====================
    def get_all_cookies(self):
        """获取所有Cookie"""
        return self.driver.get_cookies()

    def get_cookie_by_name(self, name):
        """根据名称获取Cookie"""
        return self.driver.get_cookie(name)

    def add_cookie(self, cookie_dict):
        """添加Cookie"""
        self.driver.add_cookie(cookie_dict)

    def delete_cookie(self, name):
        """删除指定Cookie"""
        self.driver.delete_cookie(name)

    def delete_all_cookies(self):
        """删除所有Cookie"""
        self.driver.delete_all_cookies()

    # ==================== JS执行 ====================
    def js_click(self, by, value, timeout=None):
        """通过JS点击元素"""
        element = self.find_element(by, value, timeout)
        self.execute_script("arguments[0].click();", element)

    def js_input_text(self, by, value, text, timeout=None):
        """通过JS输入文本"""
        element = self.find_element(by, value, timeout)
        self.execute_script("arguments[0].value = arguments[1];", element, text)

    def js_set_attribute(self, by, value, attr_name, attr_value, timeout=None):
        """通过JS设置元素属性"""
        element = self.find_element(by, value, timeout)
        self.execute_script(f"arguments[0].setAttribute('{attr_name}', '{attr_value}');", element)

    def js_remove_attribute(self, by, value, attr_name, timeout=None):
        """通过JS移除元素属性"""
        element = self.find_element(by, value, timeout)
        self.execute_script(f"arguments[0].removeAttribute('{attr_name}');", element)

    def js_scroll_to_element(self, by, value, timeout=None):
        """通过JS滚动到元素"""
        element = self.find_element(by, value, timeout)
        self.execute_script("arguments[0].scrollIntoView(true);", element)

    # ==================== 其他实用方法 ====================
    def switch_to_active_element(self):
        """切换到当前活动元素"""
        return self.driver.switch_to.active_element

    def get_page_source(self):
        """获取页面源码"""
        return self.driver.page_source

    def wait_for_page_load(self, timeout=30):
        """等待页面完全加载"""
        old_page = self.driver.find_element(By.TAG_NAME, 'html')
        yield
        WebDriverWait(self.driver, timeout).until(
            EC.staleness_of(old_page)
        )

    def is_alert_present(self, timeout=5):
        """检查警告框是否存在"""
        try:
            WebDriverWait(self.driver, timeout).until(EC.alert_is_present())
            return True
        except TimeoutException:
            return False

    def wait_for_ajax_complete(self, timeout=30):
        """等待AJAX请求完成"""
        js_condition = "return jQuery.active == 0"
        WebDriverWait(self.driver, timeout).until(
            lambda driver: self.execute_script(js_condition)
        )

    def highlight_element(self, by, value, timeout=None, duration=3):
        """高亮显示元素（用于调试）"""
        element = self.find_element(by, value, timeout)
        original_style = element.get_attribute("style")
        self.execute_script("arguments[0].setAttribute('style', arguments[1]);",
                            element, "border: 2px solid red; border-style: dashed;")
        time.sleep(duration)
        self.execute_script("arguments[0].setAttribute('style', arguments[1]);",
                            element, original_style)