import glob
import hashlib
import json
import os
import re
import subprocess
import time
from pathlib import Path
from typing import Any, Callable

import pandas as pd
import requests

from browser_env import Action, ActionTypes, Trajectory
from browser_env.actions import is_equivalent
from browser_env.auto_login import is_expired_for_sites
from browser_env.env_config import (
    CLASSIFIEDS_RESET_TOKEN,
    NO_COOKIE_RESET_SITES,
    VWA_DOMAINS,
    WA_DOMAINS,
    Website,
    get_url_from_file,
)
from browser_env.utils import get_running_websites
from constants.constants import AUTH_DIR, CLEAN_ENV_SCRIPT, RESET_ENV_SCRIPT
from utils.file_utils import contains_file
from utils.logger_utils import logger

# ===============================================================================
# Data loaders
# ===============================================================================
print_template = f"{__file__}:{{msg}}"


def get_tasks_with_trajectory(trajectory_html_path: str, test_config_base_dir: str, test_file_list: list[str] = []):
    html_files = glob.glob(os.path.join(trajectory_html_path, "*.html"))
    matches = [re.search(r"(\d+)\.html$", html_file) for html_file in html_files]
    task_ids = sorted([int(match.group(1)) for match in matches if match])
    test_files_with_trajectory = [os.path.join(test_config_base_dir, f"{i}.json") for i in task_ids]
    if test_file_list:
        test_file_list = [t for t in test_files_with_trajectory if t in test_file_list]
    else:
        test_file_list = test_files_with_trajectory
    return test_file_list


# ===============================================================================
# Helpers for evaluation loop
# ===============================================================================


def early_stop(trajectory: Trajectory, max_steps: int, thresholds: dict[str, int]) -> tuple[bool, str]:
    """Check whether need to stop early"""

    # reach the max step
    num_steps = (len(trajectory) - 1) / 2
    if num_steps >= max_steps:
        # return True, f"Reach max steps {max_steps}"
        return True, f"Unable to achieve this task after {max_steps} steps"

    last_k_actions: list[Action]
    action_seq: list[Action]

    # Case: parsing failure for k times
    k = thresholds["parsing_failure"]
    last_k_actions = trajectory[1::2][-k:]  # type: ignore[assignment]

    if len(last_k_actions) >= k:
        if all([action["action_type"] == ActionTypes.NONE for action in last_k_actions]):
            return True, f"Failed to parse actions for {k} times"

    # Case: same action for k times
    k = thresholds["repeating_action"]
    last_k_actions = trajectory[1::2][-k:]  # type: ignore[assignment]
    action_seq = trajectory[1::2]  # type: ignore[assignment]

    if len(action_seq) == 0:
        return False, ""

    last_action: Action = action_seq[-1]

    if last_action["action_type"] != ActionTypes.TYPE:
        if len(last_k_actions) >= k:
            if all([is_equivalent(action, last_action) for action in last_k_actions]):
                return True, f"Same action for {k} times"
    else:
        # check the action sequence
        if sum([is_equivalent(action, last_action) for action in action_seq]) >= k:
            return True, f"Same typing action for {k} times"
    return False, ""


# ===============================================================================
# Helpers for env setup, start, reset
# ===============================================================================
def build_site_list(sites: str | list[str], exc_sites: list[Website] = []) -> list[str]:
    if isinstance(sites, str):
        sites = [sites]

    if "all_vwa" in sites:
        sites = [website.value for website in VWA_DOMAINS if website not in exc_sites]
    elif "all_wa" in sites:
        sites = [website.value for website in WA_DOMAINS if website not in exc_sites]

    else:
        strings_to_check = [website.value for website in VWA_DOMAINS] + [website.value for website in WA_DOMAINS]
        invalid_sites = [site for site in sites if site not in strings_to_check]

        if len(invalid_sites) > 0:
            raise ValueError(f"Invalid site: {invalid_sites}. Valid sites: {strings_to_check}")
        return sites

    return sites


def get_websites_up_down(instance_id: int, env: str = "vwa") -> tuple[list[str], list[str]]:
    websites_up = get_running_websites(process_id=str(instance_id))
    websites_down = [website for website in VWA_DOMAINS if website not in websites_up]
    if env == "vwa":
        return [website.value for website in VWA_DOMAINS if website in websites_up], [
            website.value for website in websites_down
        ]
    elif env == "wa":
        return [website.value for website in WA_DOMAINS if website in websites_up], [
            website.value for website in websites_down
        ]
    return [], []


def clean_envs(instance_ids: list[int]) -> None:
    """Clean docker containers"""
    cmd = [CLEAN_ENV_SCRIPT, *[str(instance_id) for instance_id in instance_ids]]
    subprocess.run(cmd, check=True)


def _reset_environments(
    instance_id: int,
    domains: list[str],
    reset_env_script: str = RESET_ENV_SCRIPT,
    wait_for_reset: bool = True,
    env: str = "vwa",
    parallel: bool = False,
) -> None:
    # TODO: clean logic dependent on strings for websites. Need to move args from shell script to constants or change to python
    domains = domains.copy()
    if env == "vwa":
        if "homepage" in domains:
            domains.remove("homepage")
            domains.append("homepagevwa")
    elif env == "wa":
        if "homepage" in domains:
            domains.remove("homepage")
            domains.append("homepagewa")

    if not parallel:
        cmd_args = [reset_env_script, "-p", str(instance_id), *domains]
        try:
            logger.info(f"Resetting {domains}. Command: {cmd_args}. Parallel: {parallel}")
            if wait_for_reset:
                subprocess.run(cmd_args, check=True)
            else:
                subprocess.Popen(cmd_args)
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to reset environments: {e.stderr}")

    else:
        # Run non-homepage domains in parallel using Popen, then run homepage last
        homepage_variants = {"homepage", "homepagevwa", "homepagewa"}
        homepage_domain = next((d for d in domains if d in homepage_variants), None)
        other_domains = [d for d in domains if d != homepage_domain]

        procs: list[tuple[subprocess.Popen[bytes], str]] = []
        for domain in other_domains:
            domain_cmd = [reset_env_script, "-p", str(instance_id), domain]
            logger.info(f"Resetting {domain}. Command: {domain_cmd}")
            try:
                proc = subprocess.Popen(domain_cmd)
                procs.append((proc, domain))
            except Exception as e:
                logger.error(f"Failed to start reset for {domain}: {e}")

        failures: list[tuple[str, int]] = []
        if wait_for_reset and procs:
            for proc, domain in procs:
                ret = proc.wait()
                if ret != 0:
                    failures.append((domain, ret))

        # Homepage runs last to ensure all links in it are correct
        if homepage_domain:
            homepage_cmd = [reset_env_script, "-p", str(instance_id), homepage_domain]
            logger.info(f"Resetting {homepage_domain} last. Command: {homepage_cmd}")
            try:
                if wait_for_reset:
                    ret = subprocess.run(homepage_cmd, check=False).returncode
                    if ret != 0:
                        failures.append((homepage_domain, ret))
                else:
                    subprocess.Popen(homepage_cmd)
            except Exception as e:
                logger.error(f"Failed to start reset for {homepage_domain}: {e}")

        if failures:
            failed_str = ", ".join([f"{d} (code {rc})" for d, rc in failures])
            logger.error(f"One or more resets failed: {failed_str}")


def reset_envs_with_retry(
    instance_id: int,
    reset_env_script: str = RESET_ENV_SCRIPT,
    domains: list[str] | str = "all_vwa",
    wait_for_reset: bool = True,
    env: str = "vwa",
    force_homepage: bool = True,
    max_attempts: int = 3,
    parallel: bool = False,
) -> tuple[bool, list[str]]:
    """
    Reset environments with retry.

    Args:
        instance_id (int): The instance ID of the Docker container.
        reset_env_script (str, optional): The script to reset the environments. Defaults to RESET_ENV_SCRIPT.
        domains (list[str] | str, optional): The domains to reset. Defaults to "all_vwa".
        wait_for_reset (bool, optional): Whether to wait for the reset to complete. Defaults to True.
        env (str, optional): String identifying the benchmark. Defaults to "vwa". Options: "vwa", "wa".
        force_homepage (bool, optional): Whether to force the reset of the homepage. Defaults to True.
        max_attempts (int, optional): The maximum number of attempts to reset the environments. Defaults to 3.

    Returns:
        tuple[int, list[str]]:
            - int: 1 if success, 0 if failed
            - list[str]: list of domains that failed to reset
    """
    sites_to_reset = build_site_list(domains)

    if not sites_to_reset:
        return True, []

    # If force_homepage is True, add homepage to domains
    homepage_added = False
    if force_homepage and "homepage" not in sites_to_reset:
        homepage_added = True
        sites_to_reset.append("homepage")

    # Secure that all domains in domains are valid for the environment
    strings_to_check = (
        [website.value for website in VWA_DOMAINS] if env == "vwa" else [website.value for website in WA_DOMAINS]
    )
    for domain in sites_to_reset:
        assert domain in strings_to_check, f"Domain {domain} not found in {strings_to_check} for env {env}"

    if "classifieds" in sites_to_reset:
        classifieds_reset_success = try_token_reset_classifieds(instance_id)
        if classifieds_reset_success:
            sites_to_reset.remove("classifieds")

            if len(sites_to_reset) == 1 and sites_to_reset[0] == "homepage" and homepage_added:
                _, websites_down = get_websites_up_down(instance_id=instance_id, env=env)
                if "homepage" not in websites_down:
                    logger.info(
                        "Classifieds is the only site to reset, token reset successful, and only homepage is left. Skipping homepage."
                    )
                    return True, []

    # Initial start/reset
    _reset_environments(
        instance_id=instance_id,
        reset_env_script=reset_env_script,
        domains=sites_to_reset,
        wait_for_reset=wait_for_reset,
        parallel=parallel,
    )

    # If some domains failed, retry up to max_attempts times.
    must_reset_homepage = False
    for i in range(max_attempts):
        sites_to_retry = []
        _, websites_down = get_websites_up_down(instance_id=instance_id, env=env)
        for req_site in sites_to_reset:
            if req_site in websites_down:
                sites_to_retry.append(req_site)

        if not sites_to_retry:
            break

        print(f"{__file__}: Start/Reset of for instance {instance_id} failed for {sites_to_retry}. Retrying...")
        print(f"[INFO] Resetting {sites_to_retry}. Attempt {i + 1} of {max_attempts}.")
        _reset_environments(
            instance_id=instance_id,
            reset_env_script=reset_env_script,
            domains=sites_to_retry,
            wait_for_reset=wait_for_reset,
            parallel=parallel,
        )
        must_reset_homepage = True

    # If some envs were retried, force reset homepage to ensure all links in it are correct
    if must_reset_homepage:
        sites_to_retry.append("homepage")
        for i in range(max_attempts):
            _reset_environments(
                instance_id=instance_id,
                reset_env_script=reset_env_script,
                domains=["homepage"],
                wait_for_reset=wait_for_reset,
                parallel=False,
            )
            _, websites_down = get_websites_up_down(instance_id=instance_id, env=env)
            if "homepage" not in websites_down:
                sites_to_retry.remove("homepage")
                break

    if not sites_to_retry:
        return True, []
    else:
        return False, sites_to_retry


def maybe_reset_environments(
    instance_id: int,
    env: str,
    config_dict: dict[str, Any],
    require_all_sites_up: bool = True,
    force_reset: bool = False,
    force_reset_site_list: list[str] = [],
    force_homepage: bool = True,
) -> tuple[bool, list[str]]:
    """Reset environments for each task before its execution.

    Args:
        instance_id (int): Docker instance ID
        env (str): Environment to run the test on.
        config_dict (dict[str, Any] | None, optional): Task configuration dictionary. Defaults to None.
        require_all_sites_up (bool, optional): Whether to require all sites to be up. Defaults to True.
        force_reset (bool, optional): Whether to force reset the websites for current task. Defaults to False.
    """
    sites_to_reset = set()
    sites_failed_reset = []
    websites_up, websites_down = get_websites_up_down(instance_id=instance_id, env=env)

    # If no websites are up and all are required, reset and return
    if len(websites_up) == 0 and require_all_sites_up:
        env_str = "all_vwa" if env == "vwa" else "all_wa"
        reset_success, sites_failed_reset = reset_envs_with_retry(
            instance_id=instance_id,
            domains=env_str,
            wait_for_reset=True,
            force_homepage=force_homepage,
            env=env,
        )
        return reset_success, sites_failed_reset

    # Else, build list of sites to reset
    if require_all_sites_up and len(websites_down) > 0:
        sites_to_reset.update(websites_down)

    # If any required sites are down, add them to the reset list
    for req_site in config_dict["sites"]:
        if req_site not in websites_up:
            logger.info(f"{__file__}: Required site {req_site} is down. Adding to reset list.")
            sites_to_reset.add(req_site)

    # If the task requires reset, add the sites to the reset list
    if config_dict["require_reset"]:
        logger.info(f"{__file__}: Task requires reset. Adding {config_dict['sites']} to reset list.")
        [sites_to_reset.add(site) for site in config_dict["sites"]]

    # If force_reset, add sites for the config file to reset list or use the provided list
    if force_reset:
        if not force_reset_site_list:
            force_reset_site_list = config_dict["sites"]
            logger.info(f"{__file__}: Force reset enabled but site list not provided. Using {config_dict['sites']}.")
        else:
            logger.info(f"{__file__}: Force reset enabled. Adding {force_reset_site_list} to reset list.")
            [sites_to_reset.add(site) for site in force_reset_site_list]

    # Reset websites
    if sites_to_reset:
        logger.info(f"{__file__}: Resetting environments for {sites_to_reset}")
        reset_success, sites_failed_reset = reset_envs_with_retry(
            instance_id=instance_id,
            domains=list(sites_to_reset),
            wait_for_reset=True,
            max_attempts=3,
            force_homepage=force_homepage,
            env=env,
        )
    else:
        return True, []


def try_token_reset_classifieds(instance_id: int) -> bool:
    logger.info(f"{__file__}: Trying to reset classifieds with token {CLASSIFIEDS_RESET_TOKEN}.")
    try:
        classifieds_url = get_url_from_file(Website.CLASSIFIEDS, instance_id)
        reset_success = False

        if not classifieds_url:
            return reset_success

        if classifieds_url:
            response = requests.post(
                f"{classifieds_url}/index.php?page=reset",
                data={"token": CLASSIFIEDS_RESET_TOKEN},
            )

            # Check if the request was successful
            if response.status_code == 200:
                reset_success = True
                logger.info(f"{__file__}: Classifieds token reset successful.")
            else:
                reset_success = False
                logger.info(f"{__file__}: Classifieds token reset failed. Retrying with docker.")
        return reset_success
    except Exception as e:
        logger.info(f"{__file__}: Classifieds token reset failed. Retrying with docker.")
        return False


def reset_cookies_with_retry(
    docker_instance_id: str,
    sites: list[str] | str = [],
    exc_comb: bool = False,
    expired_only: bool = False,
    auth_folder: str = AUTH_DIR,
    wait_for_cookies_reset: bool = True,
    max_attempts: int = 3,
    env: str = "vwa",
) -> tuple[bool, list[str]]:
    """
    Reset cookies with retry.

    Args:
        docker_instance_id (str): The instance ID of the Docker container.
        sites (list[str] | str, optional): The sites to reset. Defaults to [] = all sites.
        exc_comb (bool, optional): Whether to exclude combinations of sites. Defaults to False.
        expired_only (bool, optional): Whether to reset only expired cookies. Defaults to False.
        auth_folder (str, optional): The folder to store the auth state. Defaults to AUTH_DIR.
    """
    if not sites:
        sites = "all_vwa" if env == "vwa" else "all_wa"
    initial_sites = build_site_list(sites, exc_sites=NO_COOKIE_RESET_SITES)
    sites_to_reset = initial_sites.copy()

    if expired_only:
        expired_sites = is_expired_for_sites(
            docker_instance_id=docker_instance_id, site_names=initial_sites, auth_folder=auth_folder, exc_comb=exc_comb
        )
        sites_to_reset = [site for site in initial_sites if site in expired_sites]

        if not sites_to_reset:
            # print(print_template.format(msg="[INFO] Expired cookies set to true and no expired cookies to reset."))
            return True, []
        else:
            logger.info(f"Expired sites: {expired_sites}. Resetting cookies for {sites_to_reset}.")

    try:
        # First attempt to reset cookies
        code, msg = reset_cookies(
            docker_instance_id=docker_instance_id,
            sites=sites_to_reset,
            exc_comb=exc_comb,
            auth_folder=auth_folder,
            wait_for_cookies_reset=wait_for_cookies_reset,
        )
        # If errors during reset_cookies script execution, log and return
        if code == -1:
            logger.error(f"Failed to reset cookies: {msg}", exc_info=True)
            return False, initial_sites

        expired_sites = is_expired_for_sites(docker_instance_id, sites_to_reset, auth_folder, exc_comb)
        if not expired_sites:
            return True, []

        # If no error but cookies still expired, Ignores expired_only and full reset of cookies for all sites.
        sites_to_reset = initial_sites
        logger.info(f"Cookies expired for {expired_sites}. Retrying. All: {sites_to_reset}")
        for i in range(max_attempts):
            logger.info(
                f"Retrying cookies reset for env {docker_instance_id}, sites {sites_to_reset}. Attempt {i + 1} of {max_attempts}."
            )
            try:
                reset_cookies(
                    docker_instance_id=docker_instance_id,
                    sites=sites_to_reset,
                    exc_comb=exc_comb,
                    auth_folder=auth_folder,
                    wait_for_cookies_reset=wait_for_cookies_reset,
                )
            except Exception as e:
                logger.error(f"Failed to reset cookies: {e}. Retrying...", exc_info=True)
                time.sleep(1)

            expired_sites = is_expired_for_sites(docker_instance_id, sites_to_reset, auth_folder, exc_comb)
            if not expired_sites:
                logger.info(f"Cookies reset successful for {sites_to_reset}.")
                return True, []

        logger.info(f"Failed to reset cookies for {sites_to_reset} after {max_attempts} attempts.")
        return False, sites_to_reset

    except Exception as e:
        logger.error(f"Failed to reset cookies: {e}", exc_info=True)
        return False, sites_to_reset


def reset_cookies(
    docker_instance_id: str,
    sites: list[str] | str = [],
    exc_comb: bool = False,
    auth_folder: str = AUTH_DIR,
    wait_for_cookies_reset: bool = True,
) -> tuple[bool, str]:
    """
    Reset cookies by calling the auto_login module.
    """
    if isinstance(sites, str):
        sites = [sites]
    elif isinstance(sites, list):
        sites_to_reset = sites.copy()

    logger.info(
        f"Resetting cookies for {sites_to_reset}, auth folder {auth_folder}, docker instance id {docker_instance_id}"
    )
    cmd = ["python", "-m", "browser_env.auto_login", "--p", str(docker_instance_id), "--auth_folder", auth_folder]
    cmd.extend(["--site_list", *sites_to_reset]) if sites_to_reset else None
    cmd.append("--exc_comb") if exc_comb else None
    try:
        if wait_for_cookies_reset:
            subprocess.run(cmd, check=True)
            return True, ""
        else:
            subprocess.Popen(cmd)
            return True, ""
    except subprocess.CalledProcessError as e:
        return False, str(e.stderr)
    except Exception as e:
        return False, str(e)


# ===============================================================================
# Other helpers
# ===============================================================================


def log_or_print(logger: Any, log_type: str, message: str) -> None:
    if logger:
        if log_type == "info":
            logger.info(message)
        elif log_type == "warning":
            logger.warning(message)
        elif log_type == "error":
            logger.error(message)
        elif log_type == "debug":
            logger.debug(message)
        elif log_type == "critical":
            logger.critical(message)
    else:
        print(message)


# ===============================================================================
# Internet helpers
# ===============================================================================


def wait_for_spinners(
    page: Any, selector: str, max_timeout_ms: float = 2 * 1000, wait_to_appear_ms: int = 0, state="detached"
) -> None:
    if page.query_selector(selector):
        page.wait_for_selector(selector, state=state, timeout=max_timeout_ms)

    if not wait_to_appear_ms:
        return

    time.sleep(wait_to_appear_ms / 1000)
    if page.query_selector(selector):
        page.wait_for_selector(selector, state=state, timeout=max_timeout_ms)
    else:
        return


def get_dom_hash(page: Any) -> str:
    """Compute a hash of the page DOM to check for stability."""
    dom = page.content()
    return hashlib.md5(dom.encode("utf-8")).hexdigest()


def wait_for_dom_hash_stability(page: Any, interval_ms: float = 500, checks: int = 3) -> None:
    """Waits until the DOM is stable for `checks` consecutive intervals."""
    previous_hash = get_dom_hash(page)
    stable_count = 0

    for _ in range(checks):
        time.sleep(interval_ms / 1000)
        current_hash = get_dom_hash(page)
        if current_hash == previous_hash:
            stable_count += 1
        else:
            stable_count = 0
            previous_hash = current_hash

    if stable_count == checks:
        return
    else:
        raise Exception(f"DOM hash is not stable after {checks} checks")


def wait_with_timeouts(wait_func: Callable, timeouts: list[int], logger=None, description="operation") -> bool:
    """
    Attempts to perform a waiting operation using a series of progressive timeouts.

    Args:
        wait_func: A callable that receives a timeout (in ms) and performs the waiting operation.
        timeouts: A list of timeout values (in ms) to try sequentially.
        logger: Optional logger for reporting status messages.
        description: A textual description of the operation for logging purposes.

    Returns:
        bool: True if one attempt was successful, False if all attempts failed.
    """
    for t in timeouts:
        try:
            wait_func(t)
            if logger:
                logger.debug(f"{description} succeeded with timeout {t}ms")
            return True
        except Exception as e:
            if logger:
                logger.debug(f"{description} failed with timeout {t}ms: {e}")
    log_or_print(logger, "warning", f"All attempts for {description} failed. Tried timeouts: {timeouts}")
    return False


def wait_for_page_to_stabilize(
    page: Any,
    max_timeout_per_check_ms: float = 2 * 1000,
    logger: Any | None = None,
    min_num_trues: int = 3,
    return_early: bool = False,
    return_after: int | None = None,
    hard_sleep: float = 0.0,
    max_overall_timeout_seconds: float = 2,
    min_wait_time_seconds: float = 0.0,
) -> bool:
    """
    Wait for a page to fully stabilize using multiple waiting mechanisms.
    Stops early if min_num_trues checks pass.

    Args:
        page: Playwright page object (sync API)
        max_timeout_ms: Maximum timeout per wait function
        logger: Optional logger for reporting
        min_num_trues: Minimum number of successful checks to consider the page stable
        return_early: If True, return after min_num_trues checks pass.
        return_after: If not None, return after `return_after` successful checks; overrides return_early.
    Returns:
        bool: True if at least min_num_trues checks passed; False otherwise.
    """
    if hard_sleep > 0.0:
        if logger:
            logger.info(f"Hard sleep for {hard_sleep} seconds.")
        time.sleep(hard_sleep)
        return True
    successful_checks = 0
    start_time = time.time()

    checks = [
        # Check 1: DOMContentLoaded
        # --------------------------
        # Wait for the DOMContentLoaded event which fires after the initial HTML
        # is loaded and parsed. This indicates that the core DOM structure is available.
        (
            lambda t: page.wait_for_load_state("domcontentloaded", timeout=min(t, max_timeout_per_check_ms)),
            [500, 500, 500],  # Wait for a max of 1500ms
            "DOMContentLoaded",
        ),
        # Check 2: Document Ready State
        # -----------------------------
        # Wait until document.readyState is 'complete', ensuring that all sub-resources,
        # such as images and stylesheets, have been fully loaded.
        (
            lambda t: page.wait_for_function(
                "document.readyState === 'complete'", timeout=min(t, max_timeout_per_check_ms)
            ),
            [500, 500, 500],
            "document ready state",
        ),
        # Check 3: Network Idle
        # ----------------------
        # Wait for the network to become idle. This check ensures that asynchronous
        # background tasks (e.g., API calls or lazy-loaded resources) have mostly finished.
        (
            lambda t: page.wait_for_load_state("networkidle", timeout=min(t, max_timeout_per_check_ms)),
            [500, 500, 500, 500, 500],  # Wait for a total of 2500ms
            "network idle",
        ),
        # Check 4: DOM Stability
        # -----------------------
        # Utilizes a MutationObserver to monitor changes in the DOM.
        # If no DOM mutations occur for 300ms, it is assumed that the page has stabilized.
        (
            lambda t: page.wait_for_function(
                """() => {
                    return new Promise(resolve => {
                        let timeout;
                        const observer = new MutationObserver(() => {
                            clearTimeout(timeout);
                            timeout = setTimeout(resolve, 300);
                        });
                        observer.observe(document.body, {
                            childList: true,
                            subtree: true,
                            attributes: true,
                            characterData: true
                        });
                        timeout = setTimeout(resolve, 300);
                    });
                }""",
                timeout=min(t, max_timeout_per_check_ms),
            ),
            [500, 500, 500],  # Wait for a total of 1500ms
            "DOM stability",
        ),
        (
            lambda t: wait_for_dom_hash_stability(page, interval_ms=t, checks=3),
            [50, 100],
            "DOM stability",
        ),
        # Check 5: Animation Frames
        # --------------------------
        # Uses two consecutive requestAnimationFrame calls to ensure that all rendering updates,
        # such as animations and transitions, have been flushed.
        (
            lambda t: page.wait_for_function(
                """() => new Promise(resolve => {
                    requestAnimationFrame(() => requestAnimationFrame(resolve));
                })""",
                timeout=min(t, max_timeout_per_check_ms),
            ),
            [200, 200, 100],  # Wait for a max of 500ms
            "animation frames",
        ),
        (
            lambda t: wait_for_spinners(
                page, selector="#checkout-loader", max_timeout_ms=min(t, max_timeout_per_check_ms), wait_to_appear_ms=5
            ),
            [500, 500, 500],  # Wait for a max of 2000ms
            "spinner disappearance",
        ),
    ]

    # Regularize params
    min_num_trues = min(min_num_trues, len(checks))

    if not return_after:
        # If return_after is not set, use return_early or default to len(checks)
        if return_early:
            return_after = min_num_trues
        else:
            return_after = len(checks)
    else:
        # If return_after is set, use it and ensure it's not greater than len(checks)
        return_after = min(return_after, len(checks))

    # emulate_slow_network(page)
    is_stable = False
    try:
        for wait_func, timeouts, description in checks:
            if time.time() - start_time > max_overall_timeout_seconds:
                log_or_print(logger, "warning", f"Page stabilization timed out after {max_overall_timeout_seconds}s.")
                return False

            try:
                if wait_with_timeouts(wait_func, timeouts, logger=logger, description=description):
                    successful_checks += 1

                # Stop early if the minimum number of successful checks has been reached.
                if successful_checks >= return_after:
                    log_or_print(logger, "debug", f"Page stabilized early after {successful_checks} successful checks.")
                    is_stable = True
                    break
            except Exception as e:
                log_or_print(logger, "warning", f"Error in timeout check {description}: {e}")
                continue

    except Exception as e:
        log_or_print(logger, "warning", f"Error in wait_for_page_to_stabilize: {e}")

    if is_stable:
        log_or_print(logger, "debug", f"Page stabilized with {successful_checks} successful checks.")
    else:
        log_or_print(logger, "warning", f"Page did not stabilize.")

    if min_wait_time_seconds:
        elapsed = time.time() - start_time
        if elapsed < min_wait_time_seconds:
            remaining_time = min_wait_time_seconds - elapsed
            log_or_print(
                logger,
                "debug",
                f"Stable checks finished before min_wait_time_seconds {min_wait_time_seconds}s. Sleeping for {remaining_time} seconds.",
            )
            time.sleep(remaining_time)

    return is_stable


# useful for debugging purposes
def emulate_slow_network(page):
    # Create a new CDP session for this page.
    client = page.context.new_cdp_session(page)
    # Enable the Network domain.
    client.send("Network.enable")
    # Emulate network conditions.
    client.send(
        "Network.emulateNetworkConditions",
        {
            "offline": False,
            "latency": 500,  # 500 ms of additional latency.
            "downloadThroughput": 50 * 1024 / 8,  # ~50 kb/s download.
            "uploadThroughput": 50 * 1024 / 8,  # ~50 kb/s upload.
        },
    )
    # Now, when you navigate, the page will load under these throttled conditions.
    # Now, when you navigate, the page will load under these throttled conditions.


# ===============================================================================
# File helpers
# ===============================================================================
def get_trajectory_json_file(file_path: str | Path) -> str | None:
    # if file is .png, get task_id from trajectory.json in parent dir
    file_path_obj = Path(file_path)

    traj_file = list(file_path_obj.parent.glob(r"*trajectory*.json")) + list(
        file_path_obj.parent.parent.glob(r"*trajectory*.json")
    )
    if traj_file:
        return str(traj_file[0])
    else:
        return None


def get_task_id_from_trajectory_json_file(file_path: str | Path) -> str:
    file_path_obj = Path(file_path)
    with open(file_path_obj, "r") as f:
        trajectory = json.load(f)
        return trajectory.get("task_id", "")


def get_task_id_from_file_path(file_path: str | Path) -> str:
    """
    Get the task id from a file path.
    The task id is the last integer that appears before the file extension (which can be empty, e.g. if `file_path` is a folder).
    Args:
        file_path (str): The path to the file or folder.

    Returns:
        str: The task id if found, otherwise an empty string.

    Examples:
        - local/experiments_1/domain_2/task234.json -> 234
        - tasks1/render_23_failed.html -> 23
        - tasks1/render_23 -> 23
        - tasks1/render_failed.html -> ""
    """
    file_path_obj = Path(file_path)
    if json_file := get_trajectory_json_file(file_path_obj):
        return get_task_id_from_trajectory_json_file(json_file)

    # if file is trajectory.json, open it and get the task_id
    if re.search(r".*trajectory.*.json", str(file_path_obj)):
        with open(file_path_obj, "r") as f:
            trajectory = json.load(f)
            return trajectory.get("task_id", "")

    # if file is .png, get task_id from trajectory.json in parent dir
    if file_path_obj.suffix in [".png", ".jpg", ".jpeg"]:
        traj_file = list(file_path_obj.parent.glob(r"*trajectory*.json")) + list(
            file_path_obj.parent.parent.glob(r"*trajectory*.json")
        )
        if traj_file:
            with open(traj_file[0], "r") as f:
                trajectory = json.load(f)
                return trajectory.get("task_id", "")
        else:
            return ""

    ext = file_path_obj.suffix
    # Gets the LAST integer that appear before the file extension (which can be empty, for e.g. if file_path is a folder)
    # e.g: tasks1/render_23_failed.html -> 23
    match = re.search(rf".*?(\d+)[^0-9]*{ext}$", str(file_path_obj))
    if match:
        return match.group(1)
    else:
        return ""


def get_config_base_dir_from_txt(txt_path) -> str:
    with open(txt_path, "r") as f:
        for line in f:
            if not line.strip().isdigit():
                return line.strip()
    return ""


def get_ids_from_tst_config_list(txt_path: str = "", tst_config_list: list | None = None) -> set[int]:
    # If .txt provided, read tst_config_list from it
    tst_config_list = tst_config_list if tst_config_list else []
    if txt_path:
        with open(txt_path, "r") as file:
            for line in file:
                line = line.strip()
                tst_config_list.append(line)

    task_ids = set()
    for tst_config in tst_config_list:
        # Case 1: JSON file path pattern (e.g., "config_files/vwa/test_shopping/102.json")
        if tst_config.endswith(".json"):
            match = re.search(r"(\d+)\.json$", tst_config)
            if match:
                task_ids.add(int(match.group(1)))
            continue

        # Case 2: Plain integer (e.g., "102")
        try:
            task_ids.add(int(tst_config))
        except ValueError:
            continue

    return task_ids


def get_task_ids_from_dir(dir_path: str) -> set[int]:
    dir_path = re.sub(r"-raw", "", dir_path)
    task_ids = set()
    for file in os.listdir(dir_path):
        task_id = get_task_id_from_file_path(file)
        if task_id:
            task_ids.add(int(task_id))
    return task_ids


def get_task_ids_from_csv(csv_path: str | Path) -> set[int]:
    try:
        if not os.path.exists(csv_path):
            return set()
        df = pd.read_csv(csv_path)
        return set(df["task_id"].astype(int))
    except Exception as e:
        print(f"Error reading CSV file {csv_path}: {e}")
        return set()


def get_attribute_from_args_file(args_file_path: str, attribute: str):
    with open(args_file_path, "r") as f:
        args = json.load(f)
        return args.get(attribute, "")


def get_domain_from_args_file(args_file_path: str):
    domain = get_attribute_from_args_file(args_file_path=args_file_path, attribute="test_config_base_dir")
    if not domain:
        raise ValueError(f"Domain not found in {args_file_path}")
    return domain.split("/")[-1].split("_")[-1]


def is_experiments_dir(dir_path: str, identifier_files: list[str] = ["args.json", "log_files.txt"], match_any=True):
    """
    Returns `True` if `dir_path` and its nested subdirs contain any or all of the `identifier_files`.
    Args:
        dir_path (str): The path to start the traversal search for the `identifier_files`.
        identifier_files (list[str]): The list of files to check for.
        match_any (bool): If `True`, returns `True` if `dir_path` contains ANY of the `identifier_files`.
        If `False`, returns `True` if `dir_path` contains ALL of the `identifier_files`.
    """
    if not os.path.isdir(dir_path):
        return False

    if isinstance(identifier_files, str):
        identifier_files = [identifier_files]

    if match_any:
        return any(contains_file(dir_path, file) for file in identifier_files)
    else:
        return all(contains_file(dir_path, file) for file in identifier_files)


def get_domain_from_test_config_dir(test_config_dir: str) -> str:
    return Path(test_config_dir).name.replace("test_", "")
