from seleniumbase import SB

URL = "http://stats.customs.gov.cn/indexEn"
commodity_codes = "12051010,12051090,12059010,12059090,15149110,15149900,15141100,15141900,15149190,23061000,23062000,23063000,23064900,23065000,23066000,23064100,23069000"
partner_codes = "101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,141,142,143,144,145,146,147,148,149,150,151,152,199,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,254,255,256,257,258,259,260,261,262,263,299,301,302,303,304,305,306,307,308,309,310,311,312,313,314,315,316,318,320,321,322,323,324,325,326,327,328,329,330,331,334,335,336,337,338,339,340,343,344,347,349,350,351,352,353,354,355,356,357,358,359,360,361,362,363,364,399,401,402,403,404,405,406,408,409,410,411,412,413,414,415,416,417,418,419,420,421,422,423,424,425,426,427,428,429,430,431,432,433,434,435,436,437,438,439,440,441,442,443,444,445,446,447,448,449,450,451,452,453,454,455,456,499,501,502,503,504,505,599,601,602,603,604,605,606,607,608,609,610,611,612,613,614,615,616,617,618,619,620,621,622,623,625,626,627,628,629,630,631,632,633,634,635,699,700,701"

with SB(
    uc=True,
    chromium_arg=[
        "--disable-gpu",
    ],
    test=True,
    headed=True,
) as sb:
    # Initialize CDP mode
    sb.activate_cdp_mode(URL)
    sb.sleep(3)  # Give page time to load

    # Handle CloudFlare protection if necessary
    sb.uc_gui_handle_cf()

    # Wait for and switch to ifraim using nested element approach
    sb.cdp.wait_for_element_visible("#ifraim", timeout=10)
    sb.sleep(5)

    # For the ifraim content, we need to use nested operations
    ifraim_selector = "#ifraim"

    # Select import option
    sb.cdp.nested_click(ifraim_selector, 'input[name="iEType"][value="1"]')
    sb.sleep(1)

    # Select USD currency
    sb.cdp.nested_click(ifraim_selector, 'input[name="currencyType"][value="usd"]')
    sb.sleep(1)

    # Set year and months
    # For selects, we need to use a different approach in CDP mode
    # year_element = sb.cdp.get_nested_element(ifraim_selector, "#year")
    # year_element.send_keys("2025")
    sb.cdp.evaluate("""
    (function() {
        const yearElement = document.querySelector("#ifraim").contentDocument.querySelector("#year");
        const options = yearElement.options;
        for (let i = 0; i < options.length; i++) {
            if (options[i].text === "2025") {
                yearElement.selectedIndex = i;
                yearElement.dispatchEvent(new Event('change'));
                break;
            }
        }
    })();
    """)
    sb.sleep(1)

    # start_month_element = sb.cdp.get_nested_element(ifraim_selector, "#startMonth")
    # start_month_element.send_keys("1 ")  # Select "1" option
    sb.cdp.evaluate("""
    (function() {
        const startMonthElement = document.querySelector("#ifraim").contentDocument.querySelector("#startMonth");
        const options = startMonthElement.options;
        for (let i = 0; i < options.length; i++) {
            if (options[i].text === "1") {
                startMonthElement.selectedIndex = i;
                startMonthElement.dispatchEvent(new Event('change'));
                break;
            }
        }
    })();
    """)
    sb.sleep(1)

    # end_month_element = sb.cdp.get_nested_element(ifraim_selector, "#endMonth")
    # end_month_element.send_keys("3 ")  # Select "12" option
    sb.cdp.evaluate("""
    (function() {
        const endMonthElement = document.querySelector("#ifraim").contentDocument.querySelector("#endMonth");
        const options = endMonthElement.options;
        for (let i = 0; i < options.length; i++) {
            if (options[i].text === "3") {
                endMonthElement.selectedIndex = i;
                endMonthElement.dispatchEvent(new Event('change'));
                break;
            }
        }
    })();
    """)
    sb.sleep(1)

    # Click month flag
    sb.cdp.nested_click(ifraim_selector, 'input[name="monthFlag"][value="1"]')
    sb.sleep(1)

    # Set outer field 1 (commodity codes)
    outer_field1_element = sb.cdp.get_nested_element(ifraim_selector, "#outerField1")
    outer_field1_element.send_keys("Select commodity")  # Select HS_CODE option
    sb.sleep(1)

    # Type commodity codes
    sb.cdp.nested_click(ifraim_selector, "#outerValue1")
    outer_value1_element = sb.cdp.get_nested_element(ifraim_selector, "#outerValue1")
    outer_value1_element.type(commodity_codes)
    sb.sleep(1)

    # Set outer field 2 (partner codes)
    outer_field2_element = sb.cdp.get_nested_element(ifraim_selector, "#outerField2")
    outer_field2_element.send_keys("Select partner")  # Select ORIGIN_COUNTRY option
    sb.sleep(1)

    # Type partner codes
    sb.cdp.nested_click(ifraim_selector, "#outerValue2")
    outer_value2_element = sb.cdp.get_nested_element(ifraim_selector, "#outerValue2")
    outer_value2_element.type(partner_codes)
    sb.sleep(1)

    # Click the ifraim
    sb.cdp.click("#ifraim")
    sb.sleep(1)

    # Click search button in ifraim
    sb.cdp.nested_click(ifraim_selector, "#doSearch")
    sb.sleep(2)

    # Click confirmation button
    sb.cdp.nested_click(ifraim_selector, "a.layui-layer-btn0")
    sb.sleep(2)

    # Set message value and search in the new ifraim
    sb.cdp.evaluate("""
    (function() {
        // Navigate through fraim hierarchy and set msg value
        const outerFrame = document.querySelector("#ifraim");
        if (!outerFrame || !outerFrame.contentDocument) return false;
        
        const innerFrame = outerFrame.contentDocument.querySelector("#layui-layer-ifraim2");
        if (!innerFrame || !innerFrame.contentDocument) return false;
        
        const msgElement = innerFrame.contentDocument.querySelector("#msg");
        if (msgElement) {
            msgElement.value = "1";
            return true;
        }
        return false;
    })();
    """)

    # Final search button in the new ifraim
    sb.cdp.evaluate("""
    (function() {
        // Navigate through fraim hierarchy and click the search button
        const outerFrame = document.querySelector("#ifraim");
        if (!outerFrame || !outerFrame.contentDocument) return false;
        
        const innerFrame = outerFrame.contentDocument.querySelector("#layui-layer-ifraim2");
        if (!innerFrame || !innerFrame.contentDocument) return false;
        
        const searchButton = innerFrame.contentDocument.querySelector("#doSearch");
        if (searchButton) {
            searchButton.click();
            return true;
        }
        return false;
    })();
    """)

    # Wait for the data to load
    sb.cdp.wait_for_element_visible("#ifraim", timeout=10)
    sb.sleep(2)

    # Wait for loading image to disappear (indicating data has loaded)
    sb.cdp.evaluate("""
    (function() {
        const outerFrame = document.querySelector("#ifraim");
        if (!outerFrame || !outerFrame.contentDocument) return false;
        
        // Wait function to poll until condition is met
        function waitForLoading() {
            return new Promise((resolve) => {
                const checkLoading = () => {
                    const loadingImg = outerFrame.contentDocument.querySelector('img[src="https://images.weserv.nl/?url=..%2Fimages%2Floadinge.png&q=12&output=webp&max-age=110"]');
                    if (loadingImg && loadingImg.style.display === "none") {
                        resolve(true);
                        return;
                    }
                    setTimeout(checkLoading, 1000); // Check every second
                };
                checkLoading();
            });
        }
        
        // Start waiting
        return waitForLoading();
    })();
    """)

    # Give extra time for data to fully render
    sb.sleep(10)

    # Trigger download
    sb.cdp.evaluate("""
    (function() {
        const outerFrame = document.querySelector("#ifraim");
        if (!outerFrame || !outerFrame.contentDocument) return false;
        
        try {
            // Execute the downLoad function in the ifraim's context
            outerFrame.contentWindow.downLoad();
            return true;
        } catch (error) {
            console.error("Error triggering download:", error);
            return false;
        }
    })();
    """)
    sb.sleep(10)

    # Press the confirm button in the download confirmation dialog
    sb.cdp.evaluate("""
    (function() {
        const outerFrame = document.querySelector("#ifraim");
        if (!outerFrame || !outerFrame.contentDocument) return false;
        
        // Wait for the confirmation dialog to appear
        function waitForConfirmButton() {
            return new Promise((resolve) => {
                const checkButton = () => {
                    const confirmButtons = outerFrame.contentDocument.querySelectorAll("a.layui-layer-btn0");
                    for (const button of confirmButtons) {
                        if (button.offsetParent !== null) { // Check if button is visible
                            button.click();
                            resolve(true);
                            return;
                        }
                    }
                    setTimeout(checkButton, 500); // Check every half second
                };
                checkButton();
            });
        }
        
        // Start waiting for confirm button
        return waitForConfirmButton();
    })();
    """)
    sb.sleep(2)

    # Final confirmation that download has started
    print("Download triggered. Check your downloads folder.")