import time
from DrissionPage import ChromiumPage 


class CloudflareBypasser:
    def __init__(self, driver: ChromiumPage, log=True):
        self.driver = driver
        self.log = log
        
    def log_message(self, message):
        if self.log:
            print(message)
    
    def search_recursively_shadow_root_with_iframe(self, ele, depth=0, max_depth=3):
        if depth > max_depth:
            self.log_message(f"Max depth {max_depth} reached during iframe search.")
            return None

        if ele.shadow_root:
            if ele.shadow_root.child().tag == "iframe":
                return ele.shadow_root.child()
        else:
            for child in ele.children():
                result = self.search_recursively_shadow_root_with_iframe(child, depth + 1, max_depth)
                if result:
                    return result
        return None

    def search_recursively_shadow_root_with_cf_input(self,ele):
        if ele.shadow_root:
            if ele.shadow_root.ele("tag:input"):
                return ele.shadow_root.ele("tag:input")
        else:
            for child in ele.children():
                result = self.search_recursively_shadow_root_with_cf_input(child)
                if result:
                    return result
        return None
    
    def locate_cf_button(self, dfs=False):
        try:
            button = None
            eles = self.driver.eles("tag:input")
            for ele in eles:
                attrs = ele.attrs
                if "name" in attrs and "type" in attrs:
                    if "turnstile" in attrs["name"] and attrs["type"] == "hidden":
                        if ele.parent() and ele.parent().shadow_root:
                            button = ele.parent().shadow_root.child()("tag:body").shadow_root("tag:input")
                            break

            if button:
                return button
            else:
                if dfs:
                    self.log_message("Basic search failed. Searching for button recursively.")
                    ele = self.driver.ele("tag:body")
                    iframe = self.search_recursively_shadow_root_with_iframe(ele)
                    if iframe:
                        return self.search_recursively_shadow_root_with_cf_input(iframe("tag:body"))
                    else:
                        self.log_message("Iframe not found. Button search failed.")
            return None

        except Exception as e:
            self.log_message(f"Error locating verification button: {e}")
            return None
    
    def click_verification_button(self, is_dfs):
        try:
            button = self.locate_cf_button(dfs=is_dfs)
            if button:
                self.log_message("Verification button found. Attempting to click.")
                if button.states.is_displayed and button.states.is_enabled:
                    button.click()
                    time.sleep(1)  # 确保事件触发
                else:
                    self.log_message("Button is not clickable.")
            else:
                self.log_message("Verification button not found.")

        except Exception as e:
            self.log_message(f"Error clicking verification button: {e}")

    def is_bypassed(self):
        try:
            title = self.driver.title.lower()
            return "just a moment" not in title and "请稍候" not in title
        except Exception as e:
            self.log_message(f"Error checking page title: {e}")
        return False

    def bypass(self, max_retry=5):
        for i in range(max_retry):
            if self.is_bypassed():
                return True
            self.log_message(f"Verification page detected.  Trying to bypass times={i}...")
            self.click_verification_button(False)
            time.sleep(2)
        return self.is_bypassed()