    def teardown(self) -> None:
        pass

class ImpersonationTask(AbstractServiceNowTask):
    """
    Task to impersonate a user.

    Parameters:
    -----------

    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task. If provided, the task will use the provided configuration instead of
        selecting a random one. See browsergym/workarena/data_files/task_configs/impersonation_users.json
        for an example of a configuration file.

    """

    def __init__(
        self, seed: int = None, instance=None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(seed=seed, instance=instance, start_rel_url="/now/nav/ui/home")
        self.fixed_config = fixed_config
        with open(IMPERSONATION_CONFIG_PATH, "r") as f:
            self.all_configs = json.load(f)
        self.__dict__.update(kwargs)

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)

        # Get task configuration
        self.user_full_name = (
            self.fixed_config if self.fixed_config else self.random.choice(self.all_configs)
        )
        assert self.user_full_name in self.all_configs

        # Generate goal
        goal = f"Impersonate the user {self.user_full_name}."
        info = {}

        return goal, info

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses


    -----------

    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task. If provided, the task will use the provided configuration instead of
        selecting a random one. See browsergym/workarena/data_files/task_configs/impersonation_users.json
        for an example of a configuration file.

    """

    def __init__(
        self, seed: int = None, instance=None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(seed=seed, instance=instance, start_rel_url="/now/nav/ui/home")
        self.fixed_config = fixed_config
        with open(IMPERSONATION_CONFIG_PATH, "r") as f:
            self.all_configs = json.load(f)
        self.__dict__.update(kwargs)

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)

        # Get task configuration
        self.user_full_name = (
            self.fixed_config if self.fixed_config else self.random.choice(self.all_configs)
        )
        assert self.user_full_name in self.all_configs

        # Generate goal
        goal = f"Impersonate the user {self.user_full_name}."
        info = {}

        return goal, info

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        task_info = f"- Impersonate the user {self.user_full_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)
        impersonate_user(self.user_full_name, page)

    def validate(
        self, page: playwright.sync_api.Page, chat_messages: list[str]
    ) -> Tuple[float, bool, str, dict]:
        page.wait_for_function("window.NOW && window.NOW.user")

        user_info = page.evaluate("window.NOW")["user"]


        self.fixed_config = fixed_config
        with open(IMPERSONATION_CONFIG_PATH, "r") as f:
            self.all_configs = json.load(f)
        self.__dict__.update(kwargs)

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)

        # Get task configuration
        self.user_full_name = (
            self.fixed_config if self.fixed_config else self.random.choice(self.all_configs)
        )
        assert self.user_full_name in self.all_configs

        # Generate goal
        goal = f"Impersonate the user {self.user_full_name}."
        info = {}

        return goal, info

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        task_info = f"- Impersonate the user {self.user_full_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)
        impersonate_user(self.user_full_name, page)

    def validate(
        self, page: playwright.sync_api.Page, chat_messages: list[str]
    ) -> Tuple[float, bool, str, dict]:
        page.wait_for_function("window.NOW && window.NOW.user")

        user_info = page.evaluate("window.NOW")["user"]

        # If the current user is not being impersonated, fail.
        if not user_info["isImpersonating"]:
            return 0, False, "", {"message": "Not currently impersonating a user."}

        # Fetch user's full name from database
        user_fullname = table_api_call(
            instance=self.instance,
            table="sys_user",


cheat

        self.user_full_name = (
            self.fixed_config if self.fixed_config else self.random.choice(self.all_configs)
        )
        assert self.user_full_name in self.all_configs

        # Generate goal
        goal = f"Impersonate the user {self.user_full_name}."
        info = {}

        return goal, info

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        task_info = f"- Impersonate the user {self.user_full_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)
        impersonate_user(self.user_full_name, page)

    def validate(
        self, page: playwright.sync_api.Page, chat_messages: list[str]
    ) -> Tuple[float, bool, str, dict]:
        page.wait_for_function("window.NOW && window.NOW.user")

        user_info = page.evaluate("window.NOW")["user"]

        # If the current user is not being impersonated, fail.
        if not user_info["isImpersonating"]:
            return 0, False, "", {"message": "Not currently impersonating a user."}

        # Fetch user's full name from database
        user_fullname = table_api_call(
            instance=self.instance,
            table="sys_user",
            params={
                "sysparm_query": f"sys_id={user_info['userID']}",
                "sysparm_fields": "name",
            },


validate

        # Generate goal
        goal = f"Impersonate the user {self.user_full_name}."
        info = {}

        return goal, info

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        task_info = f"- Impersonate the user {self.user_full_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)
        impersonate_user(self.user_full_name, page)

    def validate(
        self, page: playwright.sync_api.Page, chat_messages: list[str]
    ) -> Tuple[float, bool, str, dict]:
        page.wait_for_function("window.NOW && window.NOW.user")

        user_info = page.evaluate("window.NOW")["user"]

        # If the current user is not being impersonated, fail.
        if not user_info["isImpersonating"]:
            return 0, False, "", {"message": "Not currently impersonating a user."}

        # Fetch user's full name from database
        user_fullname = table_api_call(
            instance=self.instance,
            table="sys_user",
            params={
                "sysparm_query": f"sys_id={user_info['userID']}",
                "sysparm_fields": "name",
            },
        )["result"][0]["name"]

        # If the name matches, success.
        if user_fullname == self.user_full_name:
            return (
                1,
                True,
                "Nice work, thank you!",
                {"message": "Correct user impersonated."},
            )

        # Otherwise, fail.
        return 0, False, "", {"message": "Currently impersonating the wrong user."}

    def teardown(self) -> None:
        pass

__TASKS__ = [AllMenuTask, ImpersonationTask]


teardown

            instance=self.instance,
            table="sys_user",
            params={
                "sysparm_query": f"sys_id={user_info['userID']}",
                "sysparm_fields": "name",
            },
        )["result"][0]["name"]

        # If the name matches, success.
        if user_fullname == self.user_full_name:
            return (
                1,
                True,
                "Nice work, thank you!",
                {"message": "Correct user impersonated."},
            )

        # Otherwise, fail.
        return 0, False, "", {"message": "Currently impersonating the wrong user."}

    def teardown(self) -> None:
        pass

__TASKS__ = [AllMenuTask, ImpersonationTask]


"""
Tasks related to knowledge bases.

"""

import json
import logging
import re

from playwright.sync_api import Page
from typing import Tuple
from urllib import parse

from .base import AbstractServiceNowTask
from .comp_building_block import CompositionalBuildingBlockTask
from .utils.utils import check_url_suffix_match

from ..api.utils import table_api_call
from ..config import KB_FILEPATH, KB_CONFIG_PATH, KB_NAME, SNOW_BROWSER_TIMEOUT
from ..install import check_knowledge_base
from ..instance import SNowInstance

class KnowledgeBaseSearchTask(AbstractServiceNowTask):
    """
    Generic task to create a search for information in the knowledge base.

    Parameters:
    -----------

    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task. If provided, the task will use the provided configuration instead of
        selecting a random one. See browsergym/workarena/data_files/task_configs/knowledge_base_configs.json
        for an example of a configuration file.

    """

    def __init__(
        self,
        instance=None,
        fixed_config: dict = None,
        is_correct: bool = True,
        is_only_navigating: bool = False,
        search_by_title: bool = False,
        seed: int = None,
        **kwargs,
    ) -> None:
        """
        Parameters:
        -----------
        instance: SNowInstance
            The ServiceNow instance to run the task on.
        fixed_config:
            A fixed configuration for the task, if required.
        is_correct: bool
            Used for the compositional task.
            If false, the answer is highlighted in 'red' instead of 'yellow' when using cheat.
        is_only_navigating: bool
            Used for the compositional task.
            If we only are navigating and not searching, change the goal for the agent.
        search_by_title: bool
            Used for the compositional task.
            If true, clicks on the article title using the article name, else opens the first article.
        """
        super().__init__(
            seed=seed,
            instance=instance,
            start_rel_url="/now/nav/ui/classic/params/target/kb?id=kb_home",
        )

        # Load the knowledge base and check its integrity
        with open(KB_FILEPATH, "r") as f:
            self.kb_entries = json.load(f)
        if hasattr(self, "_base_initial_instance"):
            _, requires_install, requires_delete = check_knowledge_base(
                self._base_initial_instance,  # if user does not have permission to view the kb then this breaks
                kb_name=KB_NAME,
                kb_data=self.kb_entries,  # Need admin permissions to check
            )
        else:
            _, requires_install, requires_delete = check_knowledge_base(
                SNowInstance(),  # instance would be the non-admin instance here and this might break in case user does not have required permissions
                kb_name=KB_NAME,
                kb_data=self.kb_entries,  # Need admin permissions to check
            )
        with open(KB_CONFIG_PATH, "r") as f:
            self.all_configs = json.load(f)
        if any([requires_install, requires_delete]):
            raise RuntimeError(
                f"The knowledge base in instance {self.instance.snow_url} is missing or corrupted. "
                "See README for setup instructions."
            )
        self.fixed_config = fixed_config
        self.config = None

        # Attributes for compositional task
        self.is_correct = is_correct
        self.is_only_navigating = is_only_navigating
        self.search_by_title = search_by_title

        self.__dict__.update(kwargs)

    def _wait_for_ready(self, page: Page) -> None:
        """
        Checks that the main iframe is fully loaded

        """
        # TODO: We don't use the flag-based method used in other tasks
        #       because gsft_main doesn't have the event we register
        #       on this page. Not sure why.
        logging.debug(f"Waiting for page to be fully loaded")
        page.wait_for_load_state("networkidle")
        page.wait_for_selector('iframe[name="gsft_main"]')
        logging.debug(f"Detected page ready")

        # Get main iframe
        # XXX: We use a loop because sometimes the iframe evaluates to None
        #      even though we wait for it to be ready. This seems like a
        #      playwright bug.
        timeout = SNOW_BROWSER_TIMEOUT
        while timeout > 0:
            iframe = page.frame(name="gsft_main")
            if iframe:
                break
            page.wait_for_timeout(100)
            timeout -= 100
        else:
            raise TimeoutError(
                f"Timed out waiting for iframe to be ready in {self.instance.snow_url}"
            )

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)

        # Get task configuration
        config = self.fixed_config if self.fixed_config else self.random.choice(self.all_configs)
        self.item = config["item"]
        self.answer = config["value"]
        self.alternative_answers = config["alternative_answers"]
        self.question = config["question"]
        if self.search_by_title:
            self.kb_article_title = config["kb_article_title"]

        # Generate goal
        if self.is_only_navigating:
            goal = f'Navigate to a relevant article in the knowledge base by searching for: "{self.item}" and open the article: "{self.kb_article_title}"'
        else:
            goal = f'Answer the following question using the knowledge base: "{self.question}"'
        info = {}

        return goal, info

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        class_name = self.__class__.__name__
        class_name = class_name.replace("Task", "")
        # Split the words
        words = re.findall(r"[A-Z][^A-Z]*", class_name)
        class_name_formatted = " ".join(words)

        task_info = f"- {class_name_formatted}: {self.item} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)
        self._wait_for_ready(page)

        iframe = page.frame(name="gsft_main")
        search = iframe.locator('input[aria-label="Search"][role="textbox"]')
        search.fill(f'"{self.item}"')

        with page.expect_navigation():
            self.page.keyboard.press("Enter")

        # Click on the article
        with page.expect_navigation():
            if self.search_by_title:
                iframe.locator(f'a.kb-title:has-text("{self.kb_article_title}")').click()
            else:
                iframe.locator("a.kb-title").first.click()

        # Color the query and answer (this is just for visualization, it changes nothing to the validation)
        paragraphs = iframe.locator("p")
        for i in range(paragraphs.count()):
            paragraph = paragraphs.nth(i)
            inner_html = paragraph.inner_html()
            if self.item in inner_html:
                # Edit the inner html to change the background color of the answer
                inner_html = inner_html.replace(
                    self.item,
                    f'{self.item}',
                )
                if self.is_correct:
                    inner_html = inner_html.replace(
                        str(self.answer),
                        f'{self.answer}',
                    )
                else:
                    inner_html = inner_html.replace(
                        str(self.answer),
                        f'{self.answer}',
                    )
                paragraph.evaluate(f"element => element.innerHTML = `{inner_html}`")
                break

        # Add the "extracted" answer to the chat messages
        # TODO: this is a hack, the message will not be displayed in the html
        chat_messages.append({"role": "assistant", "message": str(self.answer)})

    def validate(self, page: Page, chat_messages: list[str]) -> tuple[float, bool, str, dict]:
        if chat_messages and chat_messages[-1]["role"] == "assistant":
            answer = chat_messages[-1]["message"]
        else:
            return (
                0,
                False,
                "",
                {"message": "The assistant did not provide an answer."},
            )

        accepted_answers = [a.lower() for a in [self.answer] + self.alternative_answers]
        answer = answer.lower()
        if any(a in answer for a in accepted_answers):
            return (
                1,
                True,
                "That is correct, thank you!",
                {"message": "Correct answer."},
            )
        else:
            return (
                0,
                False,
                "",
                {"message": "Incorrect answer provided by the assistant."},
            )

class AddCommentToKnowledgeArticleTask(AbstractServiceNowTask, CompositionalBuildingBlockTask):
    """
    Task to add a comment to a knowledge base article. Only used as a part of the compositional task for edit knowledge base
    Parameters:
    -----------
    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task.
    """

    def __init__(
        self, seed: int = None, instance=None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(
            seed=seed,
            instance=instance,
            start_rel_url="/now/nav/ui/classic/params/target/kb?id=kb_home",
            user_roles=[],
        )
        self.fixed_config = fixed_config
        if self.fixed_config is None:
            raise Exception("Please provide a config for the add comment task.")
        self.__dict__.update(kwargs)

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)
        config = self.fixed_config

        if "kb_article_title" not in config.keys():
            raise Exception("Need title in config file...")
        self.article_name = config["kb_article_title"]
        adhoc_kb_response = table_api_call(
            instance=self.instance,  # admin permissions to contribute to the KB
            table="kb_knowledge",
            method="GET",
            params={
                "sysparm_query": f"short_description={self.article_name}",
            },
        )["result"]
        if len(adhoc_kb_response) != 1:
            raise Exception("Required article not found, please fix config...")

        self.kb_article_sys_id = adhoc_kb_response[0]["sys_id"]
        self.comment = config["comment"]

        goal = f'Add the following comment to the knowledge base: "{self.comment}"'
        info = {}

        return goal, info

    def _wait_for_ready(self, page: Page) -> None:
        """
        Checks that the main iframe is fully loaded

        """
        # TODO: We don't use the flag-based method used in other tasks
        #       because gsft_main doesn't have the event we register
        #       on this page. Not sure why.
        logging.debug(f"Waiting for page to be fully loaded")
        page.wait_for_load_state("networkidle")
        page.wait_for_selector('iframe[name="gsft_main"]')
        logging.debug(f"Detected page ready")

        # Get main iframe
        # XXX: We use a loop because sometimes the iframe evaluates to None
        #      even though we wait for it to be ready. This seems like a
        #      playwright bug.
        timeout = SNOW_BROWSER_TIMEOUT
        while timeout > 0:
            iframe = page.frame(name="gsft_main")
            if iframe:
                break
            page.wait_for_timeout(100)
            timeout -= 100
        else:
            raise TimeoutError(
                f"Timed out waiting for iframe to be ready in {self.instance.snow_url}"
            )

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        class_name = self.__class__.__name__
        class_name = class_name.replace("Task", "")
        # Split the words
        words = re.findall(r"[A-Z][^A-Z]*", class_name)
        class_name_formatted = " ".join(words)

        task_info = f"- {class_name_formatted}: Add the comment '{self.comment}' for the article with title {self.article_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)

        # Check if we need to do something else, gsft_main is not loading, it seems to load when navigating from the search, so might need for compositional tasks
        self._wait_for_ready(page)
        frame = page.frame("gsft_main")
        frame.locator("button.comment-text").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').locator("html").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').get_by_label("Comments").fill(
            self.comment
        )
        frame.get_by_role("button", name="Submit").click()

    def validate(self, page: Page, chat_messages: list[str]) -> Tuple[float, bool, str, dict]:
        return super().validate(page, chat_messages)

__TASKS__ = [
    KnowledgeBaseSearchTask,
]


KnowledgeBaseSearchTask

"""

import json
import logging
import re

from playwright.sync_api import Page
from typing import Tuple
from urllib import parse

from .base import AbstractServiceNowTask
from .comp_building_block import CompositionalBuildingBlockTask
from .utils.utils import check_url_suffix_match

from ..api.utils import table_api_call
from ..config import KB_FILEPATH, KB_CONFIG_PATH, KB_NAME, SNOW_BROWSER_TIMEOUT
from ..install import check_knowledge_base
from ..instance import SNowInstance

class KnowledgeBaseSearchTask(AbstractServiceNowTask):
    """
    Generic task to create a search for information in the knowledge base.

    Parameters:
    -----------

    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task. If provided, the task will use the provided configuration instead of
        selecting a random one. See browsergym/workarena/data_files/task_configs/knowledge_base_configs.json
        for an example of a configuration file.

    """

    def __init__(
        self,
        instance=None,
        fixed_config: dict = None,
        is_correct: bool = True,
        is_only_navigating: bool = False,
        search_by_title: bool = False,
        seed: int = None,
        **kwargs,
    ) -> None:
        """
        Parameters:
        -----------
        instance: SNowInstance
            The ServiceNow instance to run the task on.
        fixed_config:
            A fixed configuration for the task, if required.
        is_correct: bool
            Used for the compositional task.
            If false, the answer is highlighted in 'red' instead of 'yellow' when using cheat.
        is_only_navigating: bool
            Used for the compositional task.
            If we only are navigating and not searching, change the goal for the agent.
        search_by_title: bool
            Used for the compositional task.
            If true, clicks on the article title using the article name, else opens the first article.
        """
        super().__init__(
            seed=seed,
            instance=instance,
            start_rel_url="/now/nav/ui/classic/params/target/kb?id=kb_home",
        )

        # Load the knowledge base and check its integrity
        with open(KB_FILEPATH, "r") as f:
            self.kb_entries = json.load(f)
        if hasattr(self, "_base_initial_instance"):
            _, requires_install, requires_delete = check_knowledge_base(
                self._base_initial_instance,  # if user does not have permission to view the kb then this breaks
                kb_name=KB_NAME,
                kb_data=self.kb_entries,  # Need admin permissions to check
            )
        else:
            _, requires_install, requires_delete = check_knowledge_base(
                SNowInstance(),  # instance would be the non-admin instance here and this might break in case user does not have required permissions
                kb_name=KB_NAME,
                kb_data=self.kb_entries,  # Need admin permissions to check
            )
        with open(KB_CONFIG_PATH, "r") as f:
            self.all_configs = json.load(f)
        if any([requires_install, requires_delete]):
            raise RuntimeError(
                f"The knowledge base in instance {self.instance.snow_url} is missing or corrupted. "
                "See README for setup instructions."
            )
        self.fixed_config = fixed_config
        self.config = None

        # Attributes for compositional task
        self.is_correct = is_correct
        self.is_only_navigating = is_only_navigating
        self.search_by_title = search_by_title

        self.__dict__.update(kwargs)

    def _wait_for_ready(self, page: Page) -> None:
        """
        Checks that the main iframe is fully loaded

        """
        # TODO: We don't use the flag-based method used in other tasks
        #       because gsft_main doesn't have the event we register
        #       on this page. Not sure why.
        logging.debug(f"Waiting for page to be fully loaded")
        page.wait_for_load_state("networkidle")
        page.wait_for_selector('iframe[name="gsft_main"]')
        logging.debug(f"Detected page ready")

        # Get main iframe
        # XXX: We use a loop because sometimes the iframe evaluates to None
        #      even though we wait for it to be ready. This seems like a
        #      playwright bug.
        timeout = SNOW_BROWSER_TIMEOUT
        while timeout > 0:
            iframe = page.frame(name="gsft_main")
            if iframe:
                break
            page.wait_for_timeout(100)
            timeout -= 100
        else:
            raise TimeoutError(
                f"Timed out waiting for iframe to be ready in {self.instance.snow_url}"
            )

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)

        # Get task configuration
        config = self.fixed_config if self.fixed_config else self.random.choice(self.all_configs)
        self.item = config["item"]
        self.answer = config["value"]
        self.alternative_answers = config["alternative_answers"]
        self.question = config["question"]
        if self.search_by_title:
            self.kb_article_title = config["kb_article_title"]

        # Generate goal
        if self.is_only_navigating:
            goal = f'Navigate to a relevant article in the knowledge base by searching for: "{self.item}" and open the article: "{self.kb_article_title}"'
        else:
            goal = f'Answer the following question using the knowledge base: "{self.question}"'
        info = {}

        return goal, info

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        class_name = self.__class__.__name__
        class_name = class_name.replace("Task", "")
        # Split the words
        words = re.findall(r"[A-Z][^A-Z]*", class_name)
        class_name_formatted = " ".join(words)

        task_info = f"- {class_name_formatted}: {self.item} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)
        self._wait_for_ready(page)

        iframe = page.frame(name="gsft_main")
        search = iframe.locator('input[aria-label="Search"][role="textbox"]')
        search.fill(f'"{self.item}"')

        with page.expect_navigation():
            self.page.keyboard.press("Enter")

        # Click on the article
        with page.expect_navigation():
            if self.search_by_title:
                iframe.locator(f'a.kb-title:has-text("{self.kb_article_title}")').click()
            else:
                iframe.locator("a.kb-title").first.click()

        # Color the query and answer (this is just for visualization, it changes nothing to the validation)
        paragraphs = iframe.locator("p")
        for i in range(paragraphs.count()):
            paragraph = paragraphs.nth(i)
            inner_html = paragraph.inner_html()
            if self.item in inner_html:
                # Edit the inner html to change the background color of the answer
                inner_html = inner_html.replace(
                    self.item,
                    f'{self.item}',
                )
                if self.is_correct:
                    inner_html = inner_html.replace(
                        str(self.answer),
                        f'{self.answer}',
                    )
                else:
                    inner_html = inner_html.replace(
                        str(self.answer),
                        f'{self.answer}',
                    )
                paragraph.evaluate(f"element => element.innerHTML = `{inner_html}`")
                break

        # Add the "extracted" answer to the chat messages
        # TODO: this is a hack, the message will not be displayed in the html
        chat_messages.append({"role": "assistant", "message": str(self.answer)})

    def validate(self, page: Page, chat_messages: list[str]) -> tuple[float, bool, str, dict]:
        if chat_messages and chat_messages[-1]["role"] == "assistant":
            answer = chat_messages[-1]["message"]
        else:
            return (
                0,
                False,
                "",
                {"message": "The assistant did not provide an answer."},
            )

        accepted_answers = [a.lower() for a in [self.answer] + self.alternative_answers]
        answer = answer.lower()
        if any(a in answer for a in accepted_answers):
            return (
                1,
                True,
                "That is correct, thank you!",
                {"message": "Correct answer."},
            )
        else:
            return (
                0,
                False,
                "",
                {"message": "Incorrect answer provided by the assistant."},
            )

class AddCommentToKnowledgeArticleTask(AbstractServiceNowTask, CompositionalBuildingBlockTask):
    """
    Task to add a comment to a knowledge base article. Only used as a part of the compositional task for edit knowledge base
    Parameters:
    -----------
    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task.
    """

    def __init__(
        self, seed: int = None, instance=None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(
            seed=seed,
            instance=instance,
            start_rel_url="/now/nav/ui/classic/params/target/kb?id=kb_home",


AddCommentToKnowledgeArticleTask

            )

        accepted_answers = [a.lower() for a in [self.answer] + self.alternative_answers]
        answer = answer.lower()
        if any(a in answer for a in accepted_answers):
            return (
                1,
                True,
                "That is correct, thank you!",
                {"message": "Correct answer."},
            )
        else:
            return (
                0,
                False,
                "",
                {"message": "Incorrect answer provided by the assistant."},
            )

class AddCommentToKnowledgeArticleTask(AbstractServiceNowTask, CompositionalBuildingBlockTask):
    """
    Task to add a comment to a knowledge base article. Only used as a part of the compositional task for edit knowledge base
    Parameters:
    -----------
    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task.
    """

    def __init__(
        self, seed: int = None, instance=None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(
            seed=seed,
            instance=instance,
            start_rel_url="/now/nav/ui/classic/params/target/kb?id=kb_home",
            user_roles=[],
        )
        self.fixed_config = fixed_config
        if self.fixed_config is None:
            raise Exception("Please provide a config for the add comment task.")
        self.__dict__.update(kwargs)

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)
        config = self.fixed_config

        if "kb_article_title" not in config.keys():
            raise Exception("Need title in config file...")
        self.article_name = config["kb_article_title"]
        adhoc_kb_response = table_api_call(
            instance=self.instance,  # admin permissions to contribute to the KB
            table="kb_knowledge",
            method="GET",
            params={
                "sysparm_query": f"short_description={self.article_name}",
            },
        )["result"]
        if len(adhoc_kb_response) != 1:
            raise Exception("Required article not found, please fix config...")

        self.kb_article_sys_id = adhoc_kb_response[0]["sys_id"]
        self.comment = config["comment"]

        goal = f'Add the following comment to the knowledge base: "{self.comment}"'
        info = {}

        return goal, info

    def _wait_for_ready(self, page: Page) -> None:
        """
        Checks that the main iframe is fully loaded

        """
        # TODO: We don't use the flag-based method used in other tasks
        #       because gsft_main doesn't have the event we register
        #       on this page. Not sure why.
        logging.debug(f"Waiting for page to be fully loaded")
        page.wait_for_load_state("networkidle")
        page.wait_for_selector('iframe[name="gsft_main"]')
        logging.debug(f"Detected page ready")

        # Get main iframe
        # XXX: We use a loop because sometimes the iframe evaluates to None
        #      even though we wait for it to be ready. This seems like a
        #      playwright bug.
        timeout = SNOW_BROWSER_TIMEOUT
        while timeout > 0:
            iframe = page.frame(name="gsft_main")
            if iframe:
                break
            page.wait_for_timeout(100)
            timeout -= 100
        else:
            raise TimeoutError(
                f"Timed out waiting for iframe to be ready in {self.instance.snow_url}"
            )

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        class_name = self.__class__.__name__
        class_name = class_name.replace("Task", "")
        # Split the words
        words = re.findall(r"[A-Z][^A-Z]*", class_name)
        class_name_formatted = " ".join(words)

        task_info = f"- {class_name_formatted}: Add the comment '{self.comment}' for the article with title {self.article_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)

        # Check if we need to do something else, gsft_main is not loading, it seems to load when navigating from the search, so might need for compositional tasks
        self._wait_for_ready(page)
        frame = page.frame("gsft_main")
        frame.locator("button.comment-text").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').locator("html").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').get_by_label("Comments").fill(
            self.comment
        )
        frame.get_by_role("button", name="Submit").click()

    def validate(self, page: Page, chat_messages: list[str]) -> Tuple[float, bool, str, dict]:
        return super().validate(page, chat_messages)

__TASKS__ = [
    KnowledgeBaseSearchTask,
]


        else:
            return (
                0,
                False,
                "",
                {"message": "Incorrect answer provided by the assistant."},
            )

class AddCommentToKnowledgeArticleTask(AbstractServiceNowTask, CompositionalBuildingBlockTask):
    """
    Task to add a comment to a knowledge base article. Only used as a part of the compositional task for edit knowledge base
    Parameters:
    -----------
    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task.
    """

    def __init__(
        self, seed: int = None, instance=None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(
            seed=seed,
            instance=instance,
            start_rel_url="/now/nav/ui/classic/params/target/kb?id=kb_home",
            user_roles=[],
        )
        self.fixed_config = fixed_config
        if self.fixed_config is None:
            raise Exception("Please provide a config for the add comment task.")
        self.__dict__.update(kwargs)

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)
        config = self.fixed_config

        if "kb_article_title" not in config.keys():
            raise Exception("Need title in config file...")
        self.article_name = config["kb_article_title"]
        adhoc_kb_response = table_api_call(
            instance=self.instance,  # admin permissions to contribute to the KB
            table="kb_knowledge",
            method="GET",
            params={
                "sysparm_query": f"short_description={self.article_name}",
            },
        )["result"]
        if len(adhoc_kb_response) != 1:
            raise Exception("Required article not found, please fix config...")

        self.kb_article_sys_id = adhoc_kb_response[0]["sys_id"]


        self.article_name = config["kb_article_title"]
        adhoc_kb_response = table_api_call(
            instance=self.instance,  # admin permissions to contribute to the KB
            table="kb_knowledge",
            method="GET",
            params={
                "sysparm_query": f"short_description={self.article_name}",
            },
        )["result"]
        if len(adhoc_kb_response) != 1:
            raise Exception("Required article not found, please fix config...")

        self.kb_article_sys_id = adhoc_kb_response[0]["sys_id"]
        self.comment = config["comment"]

        goal = f'Add the following comment to the knowledge base: "{self.comment}"'
        info = {}

        return goal, info

    def _wait_for_ready(self, page: Page) -> None:
        """
        Checks that the main iframe is fully loaded

        """
        # TODO: We don't use the flag-based method used in other tasks
        #       because gsft_main doesn't have the event we register
        #       on this page. Not sure why.
        logging.debug(f"Waiting for page to be fully loaded")
        page.wait_for_load_state("networkidle")
        page.wait_for_selector('iframe[name="gsft_main"]')
        logging.debug(f"Detected page ready")

        # Get main iframe
        # XXX: We use a loop because sometimes the iframe evaluates to None
        #      even though we wait for it to be ready. This seems like a
        #      playwright bug.
        timeout = SNOW_BROWSER_TIMEOUT
        while timeout > 0:
            iframe = page.frame(name="gsft_main")
            if iframe:
                break
            page.wait_for_timeout(100)
            timeout -= 100
        else:
            raise TimeoutError(
                f"Timed out waiting for iframe to be ready in {self.instance.snow_url}"
            )

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        class_name = self.__class__.__name__
        class_name = class_name.replace("Task", "")
        # Split the words
        words = re.findall(r"[A-Z][^A-Z]*", class_name)
        class_name_formatted = " ".join(words)

        task_info = f"- {class_name_formatted}: Add the comment '{self.comment}' for the article with title {self.article_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)

        # Check if we need to do something else, gsft_main is not loading, it seems to load when navigating from the search, so might need for compositional tasks


    instance: SNowInstance
        The instance on which to create the record.
    fixed_config: dict
        Configuration to use for the task.
    """

    def __init__(
        self, seed: int = None, instance=None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(
            seed=seed,
            instance=instance,
            start_rel_url="/now/nav/ui/classic/params/target/kb?id=kb_home",
            user_roles=[],
        )
        self.fixed_config = fixed_config
        if self.fixed_config is None:
            raise Exception("Please provide a config for the add comment task.")
        self.__dict__.update(kwargs)

    def setup_goal(self, page: Page) -> tuple[str, dict]:
        super().setup_goal(page=page)
        config = self.fixed_config

        if "kb_article_title" not in config.keys():
            raise Exception("Need title in config file...")
        self.article_name = config["kb_article_title"]
        adhoc_kb_response = table_api_call(
            instance=self.instance,  # admin permissions to contribute to the KB
            table="kb_knowledge",
            method="GET",
            params={
                "sysparm_query": f"short_description={self.article_name}",
            },
        )["result"]
        if len(adhoc_kb_response) != 1:
            raise Exception("Required article not found, please fix config...")

        self.kb_article_sys_id = adhoc_kb_response[0]["sys_id"]
        self.comment = config["comment"]

        goal = f'Add the following comment to the knowledge base: "{self.comment}"'
        info = {}

        return goal, info

    def _wait_for_ready(self, page: Page) -> None:
        """
        Checks that the main iframe is fully loaded

        """
        # TODO: We don't use the flag-based method used in other tasks
        #       because gsft_main doesn't have the event we register
        #       on this page. Not sure why.
        logging.debug(f"Waiting for page to be fully loaded")
        page.wait_for_load_state("networkidle")
        page.wait_for_selector('iframe[name="gsft_main"]')
        logging.debug(f"Detected page ready")

        # Get main iframe
        # XXX: We use a loop because sometimes the iframe evaluates to None
        #      even though we wait for it to be ready. This seems like a
        #      playwright bug.
        timeout = SNOW_BROWSER_TIMEOUT
        while timeout > 0:


        page.wait_for_load_state("networkidle")
        page.wait_for_selector('iframe[name="gsft_main"]')
        logging.debug(f"Detected page ready")

        # Get main iframe
        # XXX: We use a loop because sometimes the iframe evaluates to None
        #      even though we wait for it to be ready. This seems like a
        #      playwright bug.
        timeout = SNOW_BROWSER_TIMEOUT
        while timeout > 0:
            iframe = page.frame(name="gsft_main")
            if iframe:
                break
            page.wait_for_timeout(100)
            timeout -= 100
        else:
            raise TimeoutError(
                f"Timed out waiting for iframe to be ready in {self.instance.snow_url}"
            )

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        class_name = self.__class__.__name__
        class_name = class_name.replace("Task", "")
        # Split the words
        words = re.findall(r"[A-Z][^A-Z]*", class_name)
        class_name_formatted = " ".join(words)

        task_info = f"- {class_name_formatted}: Add the comment '{self.comment}' for the article with title {self.article_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)

        # Check if we need to do something else, gsft_main is not loading, it seems to load when navigating from the search, so might need for compositional tasks
        self._wait_for_ready(page)
        frame = page.frame("gsft_main")
        frame.locator("button.comment-text").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').locator("html").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').get_by_label("Comments").fill(
            self.comment
        )
        frame.get_by_role("button", name="Submit").click()

    def validate(self, page: Page, chat_messages: list[str]) -> Tuple[float, bool, str, dict]:
        return super().validate(page, chat_messages)

__TASKS__ = [
    KnowledgeBaseSearchTask,


cheat

        else:
            raise TimeoutError(
                f"Timed out waiting for iframe to be ready in {self.instance.snow_url}"
            )

    def get_pretty_printed_description(self) -> str:
        """
        Get the task info for this task when used in a private task; Used in L3 compositional tasks.
        called by subclasses
        """
        class_name = self.__class__.__name__
        class_name = class_name.replace("Task", "")
        # Split the words
        words = re.findall(r"[A-Z][^A-Z]*", class_name)
        class_name_formatted = " ".join(words)

        task_info = f"- {class_name_formatted}: Add the comment '{self.comment}' for the article with title {self.article_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)

        # Check if we need to do something else, gsft_main is not loading, it seems to load when navigating from the search, so might need for compositional tasks
        self._wait_for_ready(page)
        frame = page.frame("gsft_main")
        frame.locator("button.comment-text").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').locator("html").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').get_by_label("Comments").fill(
            self.comment
        )
        frame.get_by_role("button", name="Submit").click()

    def validate(self, page: Page, chat_messages: list[str]) -> Tuple[float, bool, str, dict]:
        return super().validate(page, chat_messages)

__TASKS__ = [
    KnowledgeBaseSearchTask,
]


validate

        words = re.findall(r"[A-Z][^A-Z]*", class_name)
        class_name_formatted = " ".join(words)

        task_info = f"- {class_name_formatted}: Add the comment '{self.comment}' for the article with title {self.article_name} \n"

        return task_info

    def cheat(self, page: Page, chat_messages: list[str]) -> None:
        super().cheat(page=page, chat_messages=chat_messages)

        # Check if we need to do something else, gsft_main is not loading, it seems to load when navigating from the search, so might need for compositional tasks
        self._wait_for_ready(page)
        frame = page.frame("gsft_main")
        frame.locator("button.comment-text").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').locator("html").click()
        frame.frame_locator('iframe[title="Rich Text Area"]').get_by_label("Comments").fill(
            self.comment
        )
        frame.get_by_role("button", name="Submit").click()

    def validate(self, page: Page, chat_messages: list[str]) -> Tuple[float, bool, str, dict]:
        return super().validate(page, chat_messages)

__TASKS__ = [
    KnowledgeBaseSearchTask,
]


class CompositionalBuildingBlockTask:
    """Base class for compositional building block tasks. Used to exclude these tasks from the list of tasks that are tested like atomic tasks"""

    pass


import json
import logging
import numpy as np
import playwright.sync_api
import re

from abc import ABC, abstractmethod
from tenacity import retry, stop_after_attempt, wait_fixed
from typing import List, Tuple
from urllib import parse

from .base import AbstractServiceNowTask
from .comp_building_block import CompositionalBuildingBlockTask
from .utils.utils import check_url_suffix_match

from ..api.utils import table_api_call, table_column_info
from ..config import (
    DASHBOARD_RETRIEVAL_MINMAX_CONFIG_PATH,
    DASHBOARD_RETRIEVAL_VALUE_CONFIG_PATH,
    REPORT_RETRIEVAL_MINMAX_CONFIG_PATH,
    REPORT_RETRIEVAL_VALUE_CONFIG_PATH,
    REPORT_PATCH_FLAG,
)
from ..instance import SNowInstance
from .utils.string import share_tri_gram
from .utils.utils import check_url_suffix_match

# XXX: Some notes on plot types
#      - We currently don't support maps because they are clickable and would require a more evolved cheat function
SUPPORTED_PLOT_TYPES = ["area", "bar", "column", "line", "pie", "spline"]

# Get report filter config
config = SNowInstance().report_filter_config
if config is None:
    REPORT_DATE_FILTER = REPORT_TIME_FILTER = None
else:
    REPORT_DATE_FILTER = config["report_date_filter"]
    REPORT_TIME_FILTER = config["report_time_filter"]
del config

class DashboardRetrievalTask(AbstractServiceNowTask, ABC):
    """
    A task to retrieve information from a ServiceNow dashboard

    """

    def __init__(
        self, seed: int = None, instance: SNowInstance = None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(seed=seed, instance=instance, start_rel_url="")
        self.iframe_id = "gsft_main"
        self.fixed_config = fixed_config
        self.__dict__.update(kwargs)

    @abstractmethod
    def all_configs(self) -> List[dict]:
        pass

    @abstractmethod
    def all_configs(self) -> List[dict]:
        pass

    def _get_charts(self, page: playwright.sync_api.Page) -> None:
        """
        Extract all charts on the page

        Parameters:
        -----------
        page: playwright.sync_api.Page
            The playright page on which the charts are to be extracted

        Returns:
        --------
        charts: list
            A list of charts where the chart is represented as a tuple of the chart title and the id of the
            element that contains the chart.

        """
        iframe = page.frame(name=self.iframe_id)

        charts = page.evaluate(
            f"{self.iframe_id}.Highcharts.charts.map((x) => {{if(x){{return [x.renderTo.ariaLabel, x.renderTo.id];}}}})"
        )
        charts = [
            (title.replace("Highcharts interactive chart.", "").replace(".", "").strip(), id)
            for title, id in charts
            if title
            and iframe.locator(f"#{id}").count()
            > 0  # Check if the element is actually on page (sometime rendering breaks)
        ]

        return charts

    def _read_chart(self, page: playwright.sync_api.Page, element_id: str) -> str:
        """
        Read the chart at the specified index

        Parameters:
        -----------
        page: playwright.sync_api.Page
            The playright page on which the charts are to be extracted
        element_id: str
            The ID of the element that contains the chart

        Returns:
        --------
        chart_type: str
            The type of the chart
        chart_data: dict
            The data of the chart

        """
        self._wait_for_ready(page)

        # Validate plot type
        types = page.evaluate(
            f"{self.iframe_id}.Highcharts.charts.find(chart => chart && chart.renderTo.id === '{element_id}').types"
        )
        if len(set(types)) > 1:
            raise NotImplementedError("Multiple chart types in the same chart not supported")
        type = types[0]
        if type not in SUPPORTED_PLOT_TYPES:
            raise NotImplementedError(f"Chart type {type} not supported")

        # Get data
        data = page.evaluate(
            f"""
            {self.iframe_id}.Highcharts.charts.find(chart => chart && chart.renderTo.id === "{element_id}")
            .series.map(series => ({{
                name: series.name,
                data: series.data.map(
                              point => ({{
                                 label_cat: point.category,
                                 label_name: point.name,
                                 label_origx: point.origXValue,
                                 count: point.y,
                                 percent: point.percent
                              }}))
                }}));
            """
        )

        # Post-process each series
        for i in range(len(data)):
            # For each data point in the series
            for j in range(len(data[i]["data"])):
                data_point = data[i]["data"][j]

                # Remove None percent values when count is 0
                if data_point["count"] == 0:
                    data_point["percent"] = 0

                # Strip trailing spaces from labels
                data_point["label_cat"] = (
                    data_point["label_cat"].strip() if data_point["label_cat"] else ""
                )
                data_point["label_name"] = (
                    data_point["label_name"].strip() if data_point["label_name"] else ""
                )
                data_point["label_origx"] = (
                    data_point["label_origx"].strip() if data_point["label_origx"] else ""
                )

                # Determine which label to use (this is a heuristic)
                # Usually, when the origx value is present, it is more detailed than the other labels.
                # However, in some rare cases, it corresponds to some strange value that doesn't get rendered.
                # As a heuristic for the last point, let's just make sure that origx has at least a one trigram
                # overlap with any of the other labels.
                if data_point["label_origx"] != "" and any(
                    share_tri_gram(data_point["label_origx"], data_point[x])
                    for x in ["label_cat", "label_name"]
                ):
                    data_point["label"] = data_point["label_origx"]
                else:
                    if type in ["bar", "column", "spline"]:
                        data_point["label"] = data_point["label_cat"]
                    else:
                        data_point["label"] = data_point["label_name"]
                del data_point["label_cat"]
                del data_point["label_name"]
                del data_point["label_origx"]

            assert len(set([dp["label"] for dp in data[i]["data"]])) == len(
                data[i]["data"]
            ), "Detected duplicate labels in the same series"

        return type, data

    # retry because sometimes the page is not fully loaded
    @retry(stop=stop_after_attempt(3), wait=wait_fixed(1))
    def _get_chart_by_title(
        self, page: playwright.sync_api.Page, title: str = None
    ) -> Tuple[str, dict]:
        """
        Get the chart data by title

        Parameters:
        -----------
        page: playwright.sync_api.Page
            The playright page on which the charts are to be extracted
        title: str
            The title of the chart to be read. If None, returns the first chart.

        Returns:
        --------
        chart_type: str
            The type of the chart
        chart_data: dict
            The data of the chart
        element_id: str
            The ID of the element that contains the chart

        """
        # Get chart titles and element IDs
        charts = self._get_charts(page)

        if not title:
            title = charts[0][0]

        # Find chart index by title
        chart_idx = [title.lower() for title, _ in charts].index(title.lower())

        # Load chart data
        return *self._read_chart(page, element_id=charts[chart_idx][1]), charts[chart_idx][1]

    def _wait_for_ready(self, page: playwright.sync_api.Page) -> None:
        """
        Wait for the page to be ready for task execution

        Parameters:
        -----------
        page: playwright.sync_api.Page
            The page to wait on

        """
        logging.debug(f"Waiting for {self.iframe_id} to be fully loaded")
        page.wait_for_function(
            f"typeof window.{self.iframe_id} !== 'undefined' && window.{self.iframe_id}.WORKARENA_LOAD_COMPLETE",
        )
        logging.debug(f"Detected {self.iframe_id} ready")

        logging.debug("Waiting for Highcharts API to be available")
        page.wait_for_function(f"window.{self.iframe_id}.Highcharts")
        logging.debug("Detected Highcharts API ready")

        logging.debug("Waiting for all plots to be loaded available")
        page.wait_for_function(f"window.{self.iframe_id}.WORKARENA_HIGHCHARTS_ALL_LOADED")
        logging.debug("All plots loaded")

    def get_init_scripts(self) -> List[str]:
        return super().get_init_scripts() + [
            "registerGsftMainLoaded();",
            f"""
            async function renderAllCharts() {{
                waLog('Forcing load of all charts', 'loadAllCharts');

                await waitForCondition(() => window.WORKARENA_LOAD_COMPLETE, 100);

                const canvas = window.SNC.canvas;
                if (canvas) {{
                    waLog('This is a dashboard page.', 'loadAllCharts');
                    // Trigger the rendering of each widget
                    canvas.layoutJson.panes.forEach((p) => canvas.canvasUtils.renderSlowWidget(canvas.canvasUtils.getWidgetContainer(p.uuid)));
                    // Wait for all widgets to be rendered
                    await waitForCondition(() => window.SNC.canvas.layoutJson.panes.map((p) => p.isRendered).every(value => value == true), 100);
                }}
                else {{
                    waLog('This is a report page.', 'loadAllCharts');
                    // Wait for axes to be visible (we need to use this approach since there is no canvas to help us)
                    await waitForCondition(() => document.body.innerText.toLowerCase().includes("no data to display") || document.querySelectorAll(".highcharts-point").length > 0, 100);
                }}

                // Wait for Highcharts to say that the charts are rendered
                waitForCondition(() => Highcharts.charts.all((c) => c.hasLoaded), 100)
                .then(() => {{
                            window.WORKARENA_HIGHCHARTS_ALL_LOADED = true;
                            waLog('All charts loaded', 'loadAllCharts');
                        }});
            }}
            // Run on both dashboard and reports pages
            runInGsftMainOnlyAndProtectByURL(renderAllCharts, 'pa_dashboard.do');
            runInGsftMainOnlyAndProtectByURL(renderAllCharts, 'sys_report_template.do');
            """,
            f"""
            function purifyReportUIButtons() {{
                // Delete a lot of UI features that were causing issues due to the report refreshing without
                // reloading the page. This makes the task easier, but it doesn't matter because we really
                // want to evaluate retrieval and this doesn't prevent that.
                document.querySelectorAll('[ng-click*="main.runReport"], #sidebar, #nlq-over-cb, #open-tree-navigation-button, .data-filtering-wrap').forEach(element => {{
                    if (element && element.parentNode) {{
                        element.parentNode.removeChild(element);
                    }}
                }});
                document.addEventListener('click', function(event) {{
                    event.stopPropagation();
                    event.preventDefault();
                }}, true);
                waLog('Purified report UI.', 'purifyReportUIButtons');
            }}
            // Run it only on the reports page
            runInGsftMainOnlyAndProtectByURL(purifyReportUIButtons, 'sys_report_template.do');
            """,
        ]

    def setup_goal(self, page: playwright.sync_api.Page) -> Tuple[str | dict]:
        super().setup_goal(page=page)

        # Check that the report filters are properly setup
        if REPORT_DATE_FILTER is None or REPORT_TIME_FILTER is None:
            raise RuntimeError(
                "The report date and time filters are not set. Please run the install script to set them."
            )

        # Configure task
        # ... sample a configuration
        self.config = (
            self.fixed_config if self.fixed_config else self.random.choice(self.all_configs())
        )
        # ... set start URL based on config
        # ...... some of the reports have need a date filter to be applied so we do this by patching a placeholder in the URL
        self.start_url = self.instance.snow_url + self.config["url"].replace(
            "REPORT_DATE_FILTER", REPORT_DATE_FILTER
        ).replace("REPORT_TIME_FILTER", REPORT_TIME_FILTER)

        # Produce goal string based on question type
        chart_locator = (
            f"the \"{self.config['chart_series']}\" series of "
            if self.config["chart_series"]
            else ""
        ) + (
            f"the \"{self.config['chart_title']}\" chart"
            if self.config["chart_title"]
            else "the chart"
        )
        if self.config["question"].startswith("value"):
            q_info = [x.strip() for x in self.config["question"].split(";")]
            goal = f'What is the value of "{q_info[2]}" in {chart_locator} (in {q_info[1]})?'
        elif self.config["question"] == "max":
            goal = f"What is the maximum value in {chart_locator}? Give me both the label and the count. If there are many, pick one."
        elif self.config["question"] == "min":
            goal = f"What is the minimum value in {chart_locator}? Give me both the label and the count. If there are many, pick one."
        elif self.config["question"] == "mean":
            goal = f"What is the average value in {chart_locator}? Round off to the next highest integer."
        elif self.config["question"] == "median":
            goal = f"What is the median value in {chart_locator}?"
        elif self.config["question"] == "mode":
            goal = f"What is the mode value in {chart_locator}?"
        else:
            raise NotImplementedError(f"Question type {self.config['question']} not supported")

        return goal, {}

    def cheat(self, page: playwright.sync_api.Page, chat_messages: list[str]) -> None:
        super().cheat(page, chat_messages)
        # Check if the page is the report list view. If so, open the report
        page_is_report_list_view = check_url_suffix_match(
            page, "/now/nav/ui/classic/params/target/sys_report_list.do", self
        )
        chart_title = self.config["chart_title"]
        if page_is_report_list_view:
            # Open the report
            frame = page.wait_for_selector('iframe[name="gsft_main"]').content_frame()
            # Search for the report by title
            frame.get_by_label("Search a specific field of the Reports list").select_option("Title")
            search_input = frame.locator('input[aria-label="Search"]')
            search_input.click()
            search_input.fill(chart_title)
            search_input.press("Enter")
            page.wait_for_function(
                "typeof window.gsft_main !== 'undefined' && window.gsft_main.WORKARENA_LOAD_COMPLETE"
            )
            # Click on the chart preview to open it
            frame.wait_for_selector(f'a[aria-label="Preview record: {chart_title}"]').click()
            page.wait_for_timeout(1000)
            page.keyboard.press("Enter")
            # Now in the form view, wait for the page to load and click to view the report
            page.wait_for_function(
                "typeof window.gsft_main !== 'undefined' && window.gsft_main.WORKARENA_LOAD_COMPLETE"
            )
            frame = page.wait_for_selector('iframe[name="gsft_main"]').content_frame()
            frame.get_by_text("View Report").first.click()

        self._wait_for_ready(page)

        # Get the chart data
        chart_type, chart_data, chart_element_id = self._get_chart_by_title(
            page, self.config["chart_title"]
        )

        # Extract the series
        if len(chart_data) == 1:
            chart_data = chart_data[0]["data"]
        else:
            chart_data = [
                series["data"]
                for series in chart_data
                if series["name"] == self.config["chart_series"]
            ][0]

        # Scroll to the chart
        iframe = page.frame(name=self.iframe_id)
        iframe.evaluate_handle(
            f"findElementInShadowDOM('#{chart_element_id}')"
        ).scroll_into_view_if_needed()

        # Extract the value and add it to the chat
        if self.config["question"].startswith("value"):
            format = self.config["question"].split(";")[1].strip()
            label = self.config["question"].split(";")[2].strip()
            value = [
                point["count" if format == "count" else "percent"]
                for point in chart_data
                if point["label"] == label
            ][0]
            chat_messages.append({"message": str(value), "role": "assistant"})
        elif self.config["question"] == "max":
            max_point = max(chart_data, key=lambda x: x["count"])
            chat_messages.append(
                {"message": f"{max_point['label']}, {max_point['count']}", "role": "assistant"}
            )
        elif self.config["question"] == "min":
            min_point = min(chart_data, key=lambda x: x["count"])
            chat_messages.append(
                {"message": f"{min_point['label']}, {min_point['count']}", "role": "assistant"}
            )
        elif self.config["question"] == "mean":
            counts = [data["count"] for data in chart_data]
            target_count = np.mean(counts)
            chat_messages.append({"message": f"Mean / Average {target_count}", "role": "assistant"})
        elif self.config["question"] == "median":
            counts = [data["count"] for data in chart_data]
            target_count = np.median(counts)
            chat_messages.append({"message": f"Median {target_count}", "role": "assistant"})
        elif self.config["question"] == "mode":
            counts = [data["count"] for data in chart_data]
            # We select the maximum value if there are two or more modes
            frequencies = {}
            for count in counts:
                if count not in frequencies:
                    frequencies[count] = 1
                else:
                    frequencies[count] += 1
            sorted_frequencies = {
                count: frequency
                for count, frequency in sorted(
                    frequencies.items(), key=lambda item: item[1], reverse=True
                )
            }
            max_frequency = list(sorted_frequencies.values())[0]
            max_frequencies = [
                count
                for count, frequency in sorted_frequencies.items()
                if frequency == max_frequency
            ]
            target_count = max(max_frequencies)
            chat_messages.append({"message": f"Mode {target_count}", "role": "assistant"})
        else:
            raise NotImplementedError(f"Question type \"{self.config['question']}\" not supported")

    def validate(
        self, page: playwright.sync_api.Page, chat_messages: list[str]
    ) -> Tuple[float, bool, str, dict]:
        super().validate(page, chat_messages)

        # Check if the page is in the right URL
        logging.debug("Checking if the page is in the right URL to validate the task")
        right_url = check_url_suffix_match(page, expected_url=self.start_url, task=self)
        if not right_url:
            return (
# ... truncated ...


DashboardRetrievalTask

    REPORT_PATCH_FLAG,
)
from ..instance import SNowInstance
from .utils.string import share_tri_gram
from .utils.utils import check_url_suffix_match

# XXX: Some notes on plot types
#      - We currently don't support maps because they are clickable and would require a more evolved cheat function
SUPPORTED_PLOT_TYPES = ["area", "bar", "column", "line", "pie", "spline"]

# Get report filter config
config = SNowInstance().report_filter_config
if config is None:
    REPORT_DATE_FILTER = REPORT_TIME_FILTER = None
else:
    REPORT_DATE_FILTER = config["report_date_filter"]
    REPORT_TIME_FILTER = config["report_time_filter"]
del config

class DashboardRetrievalTask(AbstractServiceNowTask, ABC):
    """
    A task to retrieve information from a ServiceNow dashboard

    """

    def __init__(
        self, seed: int = None, instance: SNowInstance = None, fixed_config: dict = None, **kwargs
    ) -> None:
        super().__init__(seed=seed, instance=instance, start_rel_url="")
        self.iframe_id = "gsft_main"
        self.fixed_config = fixed_config
        self.__dict__.update(kwargs)

    @abstractmethod
    def all_configs(self) -> List[dict]:
        pass

    @abstractmethod
    def all_configs(self) -> List[dict]:
        pass

    def _get_charts(self, page: playwright.sync_api.Page) -> None:
        """
        Extract all charts on the page

        Parameters:
        -----------
        page: playwright.sync_api.Page
            The playright page on which the charts are to be extracted

        Returns:
        --------
        charts: list
            A list of charts where the chart is represented as a tuple of the chart title and the id of the
            element that contains the chart.

        """
        iframe = page.frame(name=self.iframe_id)

        charts = page.evaluate(
            f"{self.iframe_id}.Highcharts.charts.map((x) => {{if(x){{return [x.renderTo.ariaLabel, x.renderTo.id];}}}})"
        )
        charts = [
            (title.replace("Highcharts interactive chart.", "").replace(".", "").strip(), id)
            for title, id in charts
            if title
            and iframe.locator(f"#{id}").count()
            > 0  # Check if the element is actually on page (sometime rendering breaks)
        ]

        return charts

    def _read_chart(self, page: playwright.sync_api.Page, element_id: str) -> str:
        """
        Read the chart at the specified index

        Parameters:
        -----------
        page: playwright.sync_api.Page
            The playright page on which the charts are to be extracted
        element_id: str
            The ID of the element that contains the chart

        Returns:
        --------
        chart_type: str
            The type of the chart
        chart_data: dict
            The data of the chart

        """
        self._wait_for_ready(page)

        # Validate plot type
        types = page.evaluate(
            f"{self.iframe_id}.Highcharts.charts.find(chart => chart && chart.renderTo.id === '{element_id}').types"
        )
        if len(set(types)) > 1:
            raise NotImplementedError("Multiple chart types in the same chart not supported")
        type = types[0]
        if type not in SUPPORTED_PLOT_TYPES:
            raise NotImplementedError(f"Chart type {type} not supported")

        # Get data
        data = page.evaluate(
            f"""
            {self.iframe_id}.Highcharts.charts.find(chart => chart && chart.renderTo.id === "{element_id}")
            .series.map(series => ({{
                name: series.name,
                data: series.data.map(
                              point => ({{
                                 label_cat: point.category,
                                 label_name: point.name,
                                 label_origx: point.origXValue,
                                 count: point.y,
                                 percent: point.percent
                              }}))
                }}));
            """
        )

        # Post-process each series
        for i in range(len(data)):
            # For each data point in the series
            for j in range(len(data[i]["data"])):
                data_point = data[i]["data"][j]

                # Remove None percent values when count is 0
                if data_point["count"] == 0:
                    data_point["percent"] = 0

                # Strip trailing spaces from labels
                data_point["label_cat"] = (
                    data_point["label_cat"].strip() if data_point["label_cat"] else ""
                )
                data_point["label_name"] = (
                    data_point["label_name"].strip() if data_point["label_name"] else ""
                )
                data_point["label_origx"] = (
                    data_point["label_origx"].strip() if data_point["label_origx"] else ""
                )

                # Determine which label to use (this is a heuristic)
                # Usually, when the origx value is present, it is more detailed than the other labels.
                # However, in some rare cases, it corresponds to some strange value that doesn't get rendered.
                # As a heuristic for the last point, let's just make sure that origx has at least a one trigram
                # overlap with any of the other labels.
                if data_point["label_origx"] != "" and any(
                    share_tri_gram(data_point["label_origx"], data_point[x])
                    for x in ["label_cat", "label_name"]
                ):
                    data_point["label"] = data_point["label_origx"]
                else:
                    if type in ["bar", "column", "spline"]:
                        data_point["label"] = data_point["label_cat"]
                    else:
                        data_point["label"] = data_point["label_name"]
                del data_point["label_cat"]
                del data_point["label_name"]
                del data_point["label_origx"]

            assert len(set([dp["label"] for dp in data[i]["data"]])) == len(
                data[i]["data"]
            ), "Detected duplicate labels in the same series"

        return type, data

    # retry because sometimes the page is not fully loaded
    @retry(stop=stop_after_attempt(3), wait=wait_fixed(1))
    def _get_chart_by_title(
        self, page: playwright.sync_api.Page, title: str = None
    ) -> Tuple[str, dict]:
        """
        Get the chart data by title

        Parameters:
        -----------
        page: playwright.sync_api.Page
            The playright page on which the charts are to be extracted
        title: str
            The title of the chart to be read. If None, returns the first chart.

        Returns:
        --------
        chart_type: str
            The type of the chart
        chart_data: dict
            The data of the chart
        element_id: str
            The ID of the element that contains the chart

        """
        # Get chart titles and element IDs
        charts = self._get_charts(page)

        if not title:
            title = charts[0][0]

        # Find chart index by title
        chart_idx = [title.lower() for title, _ in charts].index(title.lower())

        # Load chart data
        return *self._read_chart(page, element_id=charts[chart_idx][1]), charts[chart_idx][1]

    def _wait_for_ready(self, page: playwright.sync_api.Page) -> None:
        """
        Wait for the page to be ready for task execution

        Parameters:
        -----------
        page: playwright.sync_api.Page
            The page to wait on

        """
        logging.debug(f"Waiting for {self.iframe_id} to be fully loaded")
        page.wait_for_function(
            f"typeof window.{self.iframe_id} !== 'undefined' && window.{self.iframe_id}.WORKARENA_LOAD_COMPLETE",
        )
        logging.debug(f"Detected {self.iframe_id} ready")

        logging.debug("Waiting for Highcharts API to be available")
        page.wait_for_function(f"window.{self.iframe_id}.Highcharts")
        logging.debug("Detected Highcharts API ready")

        logging.debug("Waiting for all plots to be loaded available")
        page.wait_for_function(f"window.{self.iframe_id}.WORKARENA_HIGHCHARTS_ALL_LOADED")
        logging.debug("All plots loaded")

    def get_init_scripts(self) -> List[str]:
        return super().get_init_scripts() + [
            "registerGsftMainLoaded();",
            f"""
            async function renderAllCharts() {{
                waLog('Forcing load of all charts', 'loadAllCharts');

                await waitForCondition(() => window.WORKARENA_LOAD_COMPLETE, 100);

                const canvas = window.SNC.canvas;
                if (canvas) {{
                    waLog('This is a dashboard page.', 'loadAllCharts');
                    // Trigger the rendering of each widget
                    canvas.layoutJson.panes.forEach((p) => canvas.canvasUtils.renderSlowWidget(canvas.canvasUtils.getWidgetContainer(p.uuid)));
                    // Wait for all widgets to be rendered
                    await waitForCondition(() => window.SNC.canvas.layoutJson.panes.map((p) => p.isRendered).every(value => value == true), 100);
                }}
                else {{
                    waLog('This is a report page.', 'loadAllCharts');
                    // Wait for axes to be visible (we need to use this approach since there is no canvas to help us)
                    await waitForCondition(() => document.body.innerText.toLowerCase().includes("no data to display") || document.querySelectorAll(".highcharts-point").length > 0, 100);
                }}

                // Wait for Highcharts to say that the charts are rendered
                waitForCondition(() => Highcharts.charts.all((c) => c.hasLoaded), 100)
                .then(() => {{
                            window.WORKARENA_HIGHCHARTS_ALL_LOADED = true;
                            waLog('All charts loaded', 'loadAllCharts');
                        }});
            }}
            // Run on both dashboard and reports pages
            runInGsftMainOnlyAndProtectByURL(renderAllCharts, 'pa_dashboard.do');
            runInGsftMainOnlyAndProtectByURL(renderAllCharts, 'sys_report_template.do');
            """,
            f"""
            function purifyReportUIButtons() {{
                // Delete a lot of UI features that were causing issues due to the report refreshing without
                // reloading the page. This makes the task easier, but it doesn't matter because we really
                // want to evaluate retrieval and this doesn't prevent that.
                document.querySelectorAll('[ng-click*="main.runReport"], #sidebar, #nlq-over-cb, #open-tree-navigation-button, .data-filtering-wrap').forEach(element => {{
                    if (element && element.parentNode) {{
                        element.parentNode.removeChild(element);
                    }}
                }});
                document.addEventListener('click', function(event) {{
                    event.stopPropagation();
                    event.preventDefault();
                }}, true);
                waLog('Purified report UI.', 'purifyReportUIButtons');
            }}
            // Run it only on the reports page
            runInGsftMainOnlyAndProtectByURL(purifyReportUIButtons, 'sys_report_template.do');
            """,
        ]

    def setup_goal(self, page: playwright.sync_api.Page) -> Tuple[str | dict]:
        super().setup_goal(page=page)

        # Check that the report filters are properly setup
        if REPORT_DATE_FILTER is None or REPORT_TIME_FILTER is None:
            raise RuntimeError(
                "The report date and time filters are not set. Please run the install script to set them."
            )

        # Configure task
        # ... sample a configuration
        self.config = (
            self.fixed_config if self.fixed_config else self.random.choice(self.all_configs())
        )
        # ... set start URL based on config
        # ...... some of the reports have need a date filter to be applied so we do this by patching a placeholder in the URL
        self.start_url = self.instance.snow_url + self.config["url"].replace(
            "REPORT_DATE_FILTER", REPORT_DATE_FILTER
        ).replace("REPORT_TIME_FILTER", REPORT_TIME_FILTER)

        # Produce goal string based on question type
        chart_locator = (
            f"the \"{self.config['chart_series']}\" series of "
            if self.config["chart_series"]
            else ""
        ) + (
            f"the \"{self.config['chart_title']}\" chart"
            if self.config["chart_title"]
            else "the chart"
        )
        if self.config["question"].startswith("value"):
            q_info = [x.strip() for x in self.config["question"].split(";")]
            goal = f'What is the value of "{q_info[2]}" in {chart_locator} (in {q_info[1]})?'
        elif self.config["question"] == "max":
            goal = f"What is the maximum value in {chart_locator}? Give me both the label and the count. If there are many, pick one."
        elif self.config["question"] == "min":
            goal = f"What is the minimum value in {chart_locator}? Give me both the label and the count. If there are many, pick one."
        elif self.config["question"] == "mean":
            goal = f"What is the average value in {chart_locator}? Round off to the next highest integer."
        elif self.config["question"] == "median":
            goal = f"What is the median value in {chart_locator}?"
        elif self.config["question"] == "mode":
            goal = f"What is the mode value in {chart_locator}?"
        else:
            raise NotImplementedError(f"Question type {self.config['question']} not supported")

        return goal, {}

    def cheat(self, page: playwright.sync_api.Page, chat_messages: list[str]) -> None:
        super().cheat(page, chat_messages)
        # Check if the page is the report list view. If so, open the report
        page_is_report_list_view = check_url_suffix_match(
            page, "/now/nav/ui/classic/params/target/sys_report_list.do", self
        )
        chart_title = self.config["chart_title"]
        if page_is_report_list_view:
            # Open the report
            frame = page.wait_for_selector('iframe[name="gsft_main"]').content_frame()
            # Search for the report by title
            frame.get_by_label("Search a specific field of the Reports list").select_option("Title")
            search_input = frame.locator('input[aria-label="Search"]')
            search_input.click()
            search_input.fill(chart_title)
            search_input.press("Enter")
            page.wait_for_function(
                "typeof window.gsft_main !== 'undefined' && window.gsft_main.WORKARENA_LOAD_COMPLETE"
            )
            # Click on the chart preview to open it
            frame.wait_for_selector(f'a[aria-label="Preview record: {chart_title}"]').click()
            page.wait_for_timeout(1000)
            page.keyboard.press("Enter")
            # Now in the form view, wait for the page to load and click to view the report
            page.wait_for_function(
                "typeof window.gsft_main !== 'undefined' && window.gsft_main.WORKARENA_LOAD_COMPLETE"
            )
            frame = page.wait_for_selector('iframe[name="gsft_main"]').content_frame()
            frame.get_by_text("View Report").first.click()

        self._wait_for_ready(page)

        # Get the chart data
        chart_type, chart_data, chart_element_id = self._get_chart_by_title(
            page, self.config["chart_title"]
        )

        # Extract the series
        if len(chart_data) == 1:
            chart_data = chart_data[0]["data"]
        else:
            chart_data = [
                series["data"]
                for series in chart_data
                if series["name"] == self.config["chart_series"]
            ][0]

        # Scroll to the chart
        iframe = page.frame(name=self.iframe_id)
        iframe.evaluate_handle(
            f"findElementInShadowDOM('#{chart_element_id}')"
        ).scroll_into_view_if_needed()

        # Extract the value and add it to the chat
        if self.config["question"].startswith("value"):
            format = self.config["question"].split(";")[1].strip()
            label = self.config["question"].split(";")[2].strip()
            value = [
                point["count" if format == "count" else "percent"]
                for point in chart_data
                if point["label"] == label
            ][0]
            chat_messages.append({"message": str(value), "role": "assistant"})
        elif self.config["question"] == "max":
            max_point = max(chart_data, key=lambda x: x["count"])
            chat_messages.append(
                {"message": f"{max_point['label']}, {max_point['count']}", "role": "assistant"}
            )
        elif self.config["question"] == "min":
            min_point = min(chart_data, key=lambda x: x["count"])
            chat_messages.append(
                {"message": f"{min_point['label']}, {min_point['count']}", "role": "assistant"}
            )
        elif self.config["question"] == "mean":
            counts = [data["count"] for data in chart_data]
            target_count = np.mean(counts)
            chat_messages.append({"message": f"Mean / Average {target_count}", "role": "assistant"})
        elif self.config["question"] == "median":
            counts = [data["count"] for data in chart_data]
            target_count = np.median(counts)
            chat_messages.append({"message": f"Median {target_count}", "role": "assistant"})
        elif self.config["question"] == "mode":
            counts = [data["count"] for data in chart_data]
            # We select the maximum value if there are two or more modes
            frequencies = {}
            for count in counts:
                if count not in frequencies:
                    frequencies[count] = 1
                else:
                    frequencies[count] += 1
            sorted_frequencies = {
                count: frequency
                for count, frequency in sorted(
                    frequencies.items(), key=lambda item: item[1], reverse=True
                )
            }
            max_frequency = list(sorted_frequencies.values())[0]
            max_frequencies = [
                count
                for count, frequency in sorted_frequencies.items()
                if frequency == max_frequency
            ]
            target_count = max(max_frequencies)
            chat_messages.append({"message": f"Mode {target_count}", "role": "assistant"})
        else:
            raise NotImplementedError(f"Question type \"{self.config['question']}\" not supported")

    def validate(
        self, page: playwright.sync_api.Page, chat_messages: list[str]
    ) -> Tuple[float, bool, str, dict]:
        super().validate(page, chat_messages)

        # Check if the page is in the right URL
        logging.debug("Checking if the page is in the right URL to validate the task")
        right_url = check_url_suffix_match(page, expected_url=self.start_url, task=self)
        if not right_url:
            return (
                0,
                False,
                "",
                {
                    "message": f"The page is not in the right URL to validate task {self.__class__.__name__}."
                },
            )

        self._wait_for_ready(page)

        # Get the chart data
        logging.debug("Extracting chart data")
        _, chart_data, _ = self._get_chart_by_title(page, self.config["chart_title"])

        # Extract the series
        logging.debug("Extracting the series")
        if len(chart_data) == 1:
            chart_data = chart_data[0]["data"]
        else:
            chart_data = [
                series["data"
# ... truncated ...


MultiChartValueRetrievalTask

            format = self.random.choice(["count", "percent"])

            # Select a random label from the chart data
            label = self.random.choice(labels)

            return {
                "url": url,
                "chart_title": chart_title,
                "chart_series": chart_series,
                "question": f"{question}; {format}; {label}",
            }
        else:
            return {
                "url": url,
                "chart_title": chart_title,
                "chart_series": chart_series,
                "question": question,
            }

class MultiChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class MultiChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class SingleChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartMeanMedianModeRetrievalTask(
    DashboardRetrievalTask, CompositionalBuildingBlockTask
):


MultiChartMinMaxRetrievalTask

            return {
                "url": url,
                "chart_title": chart_title,
                "chart_series": chart_series,
                "question": f"{question}; {format}; {label}",
            }
        else:
            return {
                "url": url,
                "chart_title": chart_title,
                "chart_series": chart_series,
                "question": question,
            }

class MultiChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class MultiChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class SingleChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartMeanMedianModeRetrievalTask(
    DashboardRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class WorkLoadBalancingMinMaxRetrievalTask(


SingleChartValueRetrievalTask

            }
        else:
            return {
                "url": url,
                "chart_title": chart_title,
                "chart_series": chart_series,
                "question": question,
            }

class MultiChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class MultiChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class SingleChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartMeanMedianModeRetrievalTask(
    DashboardRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class WorkLoadBalancingMinMaxRetrievalTask(
    SingleChartMinMaxRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))


SingleChartMinMaxRetrievalTask

                "chart_series": chart_series,
                "question": question,
            }

class MultiChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class MultiChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class SingleChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartMeanMedianModeRetrievalTask(
    DashboardRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class WorkLoadBalancingMinMaxRetrievalTask(
    SingleChartMinMaxRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

    def setup_goal(self, page: playwright.sync_api.Page) -> Tuple[str | dict]:
        super().setup_goal(page=page)

        # Configure task
        # ... sample a configuration


SingleChartMeanMedianModeRetrievalTask

class MultiChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class MultiChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(DASHBOARD_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class SingleChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartMeanMedianModeRetrievalTask(
    DashboardRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class WorkLoadBalancingMinMaxRetrievalTask(
    SingleChartMinMaxRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

    def setup_goal(self, page: playwright.sync_api.Page) -> Tuple[str | dict]:
        super().setup_goal(page=page)

        # Configure task
        # ... sample a configuration
        self.config = (
            self.fixed_config if self.fixed_config else self.random.choice(self.all_configs())
        )
        # ... set start URL based on config
        self.start_url = self.instance.snow_url + self.config["url"]

        goal = f"Create a filter to find reports whose title contains hashtag {self.problem_hashtag} and open the report."


WorkLoadBalancingMinMaxRetrievalTask

        return json.load(open(DASHBOARD_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartValueRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_VALUE_CONFIG_PATH, "r"))

class SingleChartMinMaxRetrievalTask(DashboardRetrievalTask):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class SingleChartMeanMedianModeRetrievalTask(
    DashboardRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

class WorkLoadBalancingMinMaxRetrievalTask(
    SingleChartMinMaxRetrievalTask, CompositionalBuildingBlockTask
):
    def all_configs(self):
        return json.load(open(REPORT_RETRIEVAL_MINMAX_CONFIG_PATH, "r"))

    def setup_goal(self, page: playwright.sync_api.Page) -> Tuple[str | dict]:
        super().setup_goal(page=page)

        # Configure task
        # ... sample a configuration
        self.config = (
            self.fixed_config if self.fixed_config else self.random.choice(self.all_configs())
        )
        # ... set start URL based on config
        self.start_url = self.instance.snow_url + self.config["url"]

        goal = f"Create a filter to find reports whose title contains hashtag {self.problem_hashtag} and open the report."
        goal += " From the report, identify the user with the most assigned problems and the user with the least assigned problems."

        return goal, {}

__TASKS__ = [
    var
    for var in locals().values()
    if isinstance(var, type)
    and issubclass(var, DashboardRetrievalTask)
    and not issubclass(var, CompositionalBuildingBlockTask)
    and var is not DashboardRetrievalTask
]