import asyncio
from typing import Dict

import playwright
from browsergym.core.action.base import execute_python_code

# from playwright._impl._errors import TimeoutError as PlaywrightTimeoutError

from agentS.agent_state import AgentState

from langchain_core.tools import tool
import PyPDF2
import io
import requests
import re

from agentS.tool_processor import transform_action

TOOLS = {
    "Click",
    "Type",
    "Scroll",
    "Google",
    "ANSWER",
    "GoBack",
    "ReadPage",
}


async def read_pdf_content(page, pages=None):
    # Get the PDF URL
    pdf_url = page.url

    # Download the PDF content
    response = requests.get(pdf_url)
    pdf_content = io.BytesIO(response.content)

    # Use PyPDF2 to read the PDF content
    pdf_reader = PyPDF2.PdfReader(pdf_content)

    # Extract text from all pages
    text = ""
    for page_num in range(len(pdf_reader.pages)):
        if pages is None or page_num in pages:
            text += pdf_reader.pages[page_num].extract_text()
        elif pages is not None and max(pages) < page_num:
            break
        else:
            text += pdf_reader.pages[page_num].extract_text()

    return text


async def read_pdf_content_sync(page, pages=None):
    # Get the PDF URL
    pdf_url = page.url

    # Download the PDF content
    response = requests.get(pdf_url)
    pdf_content = io.BytesIO(response.content)

    # Use PyPDF2 to read the PDF content
    pdf_reader = PyPDF2.PdfReader(pdf_content)

    # Extract text from all pages
    text = ""
    for page_num in range(len(pdf_reader.pages)):
        if pages is None or page_num in pages:
            text += pdf_reader.pages[page_num].extract_text()
        elif pages is not None and max(pages) < page_num:
            break
        else:
            text += pdf_reader.pages[page_num].extract_text()

    return text


class Tools:

    @staticmethod
    @tool
    async def create_action_plan(state: Dict) -> str:
        """
        Create an action plan based on the prediction
        """
        page = state.get("page")
        # read the page content
        text_content = await page.evaluate('() => document.body.innerText')

        prediction = state.get("prediction")
        actions = prediction.get("actions")
        state["prediction"]["actions"] = actions.split("|")
        return f"Created action plan: {actions}"

    @staticmethod
    @tool
    async def read_page(state) -> str:
        """
        Read the content of the page
        """
        page = state.get("page")
        url = page.url

        # Check if the URL contains a PDF
        if 'pdf' in url.split('/'):
            pdf_content = await read_pdf_content(page=page, pages=[
                0])  # To change the pages, change the pages=[0] to pages=[0, 1, 2, ...]
            return f"Read PDF content: {pdf_content}"
        else:
            text_content = await page.evaluate('() => document.body.innerText')
            return f"Read page content: {text_content}"
        # return f"Tried to read the page, but the URL does not contain a PDF, {url}. Proceed with other actions."

    @staticmethod
    @tool
    async def goback(state) -> str:
        """
        Navigate back to the previous
        """
        page = state.get("page")
        await page.go_back()
        return f'Navigated to previous page: {page.url}'

    @staticmethod
    @tool
    # async def click(state) -> str:
    # async def click(state, additional_info) -> str:
    async def click(state) -> str:
        """
        Click on an element on the page
        """
        page = state.get("page")
        extension_obj = state.get("extension_obj")
        annotation = state.get("annotations")
        args = state["prediction"]["args"]
        click_args = args[list(args.keys())[0]]
        # if click_args is None or len(click_args) != 1:
        # if click_args is None or len(click_args.split()) != 1:
        if click_args is None:
            return f"Failed to click {args} due to incorrect arguments."
        selector = str(click_args)
        selector = extract_number(selector)
        if 'Error' in selector:
            return 'Failed to click because no number found in the string.'

        # Make sure that the selector is event exists
        if selector not in annotation[1]['browser_content']:
            return f"Failed to click {selector}. Element not found."

        element_info = annotation[1]['browser_content'][selector]
        # if element_info["element_type"] not in ["input", "button", "link", "form", "select_toggle"]:
        #     return f"Failed to click {element_info}. Element type not supported."
        element_id = element_info['element_id']
        # element_to_perform_action = extension_obj.get_selector(element_id)
        element_to_perform_action = extension_obj.get_selector(page, element_id)
        try:
            # Attempt to type with a timeout of 30 seconds
            await element_to_perform_action.click(timeout=30000)
            observation = f'Clicked {element_info["element_name"]}, {element_info["element_type"]}'
        except:
            # Handling timeout exception
            observation = f'Failed to click {element_info["element_name"]} due to timeout.'

        return observation

    @staticmethod
    @tool
    async def goback(state) -> str:
        """
        Navigate back to the previous
        """
        page = state.get("page")
        await page.go_back()
        return f'Navigated to previous page: {page.url}'

    @staticmethod
    @tool
    async def to_google(state) -> str:
        """
        Navigate to Google
        """
        page = state.get("page")
        await page.goto("https://www.google.com")
        return f'Navigated to Google'

    @staticmethod
    @tool
    async def type(state) -> str:
        """
        Type in an element on the page
        """
        page = state.get("page")
        type_args = state["prediction"]["args"]
        annotation = state.get("annotations")
        extension_obj = state.get("extension_obj")
        # if type_args is None or len(type_args) != 2:
        #     return f"Failed to type {type_args} due to incorrect arguments."
        selector = str(type_args[list(type_args.keys())[0]]) if not isinstance(type_args[list(type_args.keys())[0]],
                                                                               str) else type_args[
            list(type_args.keys())[0]]
        text = str(type_args[list(type_args.keys())[1]])  # Convert to string
        selector = extract_number(selector)
        if 'Error' in selector:
            print(f"\n ********** {selector} ********** \n")
            return 'Failed to click because no number found in the string.'
        # Make sure that the selector is event exists
        if selector not in annotation[1]['browser_content']:
            return f"Failed to type {text} due to incorrect selector: {selector}."

        element_info = annotation[1]['browser_content'][selector]
        element_id = element_info['element_id']
        # element_to_perform_action = extension_obj.get_selector(element_id)
        element_to_perform_action = extension_obj.get_selector(page, element_id)

        try:
            # Attempt to type with a timeout of 30 seconds
            await element_to_perform_action.type("", timeout=30000)  # Clear the input field
            await element_to_perform_action.type(text, timeout=30000)
            observation = f'Typed "{text}" in {element_info["element_name"]}, {element_info["element_type"]}'
        except:
            # Handling timeout exception
            observation = f'Failed to type "{text}" in {element_info["element_name"]} due to timeout.'

        return observation

    @staticmethod
    @tool
    def answer(state) -> str:
        """
        Answer a question
        """
        return state["prediction"]["args"][0]

    @staticmethod
    @tool
    async def scroll(state) -> str:
        """
        Scroll the page
        """
        page = state.get("page")
        scroll_args = state["prediction"]["args"]
        # if scroll_args is None or len(scroll_args) != 1:
        #     return "Failed to scroll due to incorrect arguments."
        direction = scroll_args[0]
        if direction not in ["up", "down"]:
            return f"Failed to scroll due to incorrect direction: {direction}."
        await page.scroll(direction)
        return f'Scrolled {direction}'


def extract_number(input_string):
    # Check if the input string is an actual number
    try:
        # Try to convert the string to a float
        num = int(input_string)
        return input_string
    except ValueError:
        # If it raises a ValueError, it means the string is not a valid number
        # Extract numbers from the string using regex
        numbers = re.findall(r'[-+]?\d*\.\d+|\d+', input_string)
        if numbers:
            # Convert extracted strings to float and handle multiple numbers
            extracted_numbers = [str(num) for num in numbers]
            print(f"Extracted numbers from the {input_string}: {extracted_numbers}")
            return extracted_numbers[0] if len(extracted_numbers) == 1 else extracted_numbers
        else:
            print("No numbers found in the string.")
            return "Error: No numbers found in the string."


class SychronousTools:
    @staticmethod
    async def get_elem_by_bid_async(page, bid, scroll_into_view=False):
        if not isinstance(bid, str):
            raise ValueError(f"expected a string, got {repr(bid)}")

        current_frame = page

        i = 0
        while bid[i:] and not bid[i:].isnumeric():
            i += 1
            frame_bid = bid[:i]
            frame_elem = await current_frame.get_by_test_id(frame_bid)
            if not await frame_elem.count():
                raise ValueError(f'Could not find element with bid "{bid}"')
            if scroll_into_view:
                await frame_elem.scroll_into_view_if_needed(timeout=500)
            current_frame = frame_elem.frame_locator(":scope")

        elem = await current_frame.get_by_test_id(bid)
        if not await elem.count():
            raise ValueError(f'Could not find element with bid "{bid}"')
        if scroll_into_view:
            await elem.scroll_into_view_if_needed(timeout=500)
        return elem

    @staticmethod
    def get_elem_by_bid(page: playwright.sync_api.Page, bid: str,
                        scroll_into_view: bool = False) -> playwright.sync_api.Locator:
        if not isinstance(bid, str):
            raise ValueError(f"expected a string, got {repr(bid)}")

        current_frame = page

        i = 0
        while bid[i:] and not bid[i:].isnumeric():
            i += 1
            frame_bid = bid[:i]
            frame_elem = current_frame.get_by_test_id(frame_bid)
            if not frame_elem.count():
                raise ValueError(f'Could not find element with bid "{bid}"')
            if scroll_into_view:
                frame_elem.scroll_into_view_if_needed(timeout=500)
            current_frame = frame_elem.frame_locator(":scope")

        elem = current_frame.get_by_test_id(bid)
        if not elem.count():
            raise ValueError(f'Could not find element with bid "{bid}"')
        if scroll_into_view:
            elem.scroll_into_view_if_needed(timeout=500)
        return elem

    @staticmethod
    @tool
    def update_policy(state: Dict) -> str:
        """
        Requests updated policy from Planner Agent based on current state and changes.

        Parameters:
        - reason (str): Explain why update is needed, including:
            1. Specific reason (obstacle, new info, etc.)
            2. How situation differs from expectations
            3. Relevant changes in task/environment
            4. New insights and failed attempts
            5. Suggested new approach

        Returns:
        - str: Updated policy

        Example:
        update_policy("Expected login form, found 'Create Account'. Login link absent.
                       Consider temp account or alternative access. Check 'Help/FAQ'.")

        Use when:
        * Unexpected obstacles or failures
        * New pages/sections not in current policy
        * Policy inadequate for current state
        * Major progress needing new guidance
        * Strategy-altering discoveries
        * All current options exhausted
        * Significßßant page structure changes

        Provide detailed reasons for effective policy update.
        """
        pass

    @staticmethod
    @tool
    def select_option(state) -> str:

        """
        Select option(s) in <select>, dropdown, or combobox elements. Works with HTML selects
        and complex UI components like custom dropdowns/comboboxes.

        Function simulates user interaction:
        - Clicks to open dropdown if needed
        - Selects specified option(s) by value or visible text

        Parameters:
        - bid (str): Unique identifier for the select/dropdown/combobox element
        - options (str | list[str]): Option(s) to select. Single string or list of strings

        Can select by option value or label. Supports multi-select for applicable components.

        Examples:
            # Single option in dropdown/combobox
            select_option('dropdown_id', "blue")

            # Multiple options in multi-select
            select_option('multiselect_id', ["red", "green", "blue"])

            # Select by visible text in combobox
            select_option('combobox_id', "Desired Option Text")

        Note: Adapts to various selection components. May perform extra steps (e.g., clicking
        to open dropdown) for custom implementations.
        """
        dict_keys = list(state.keys())
        bid = state[dict_keys[0]]
        options = state[dict_keys[1]]
        env = state.get("env")
        page = state.get("env").page
        action = {
            "name": "select_option",
            "args": {
                'state': state
            }
        }
        action = transform_action(action)
        info = {}
        try:
            if env.action_mapping:
                code = env.action_mapping(action)
            else:
                code = action
            execute_python_code(
                code,
                env.page,
                send_message_to_user=None,
                report_infeasible_instructions=None,
            )
        except Exception as e:
            last_action_error = f"{type(e).__name__}: {e}"
            match = re.match("TimeoutError: Timeout ([0-9]+)ms exceeded.", last_action_error)
            if match:
                info["action_exec_timeout"] = float(match.groups()[0]) / 1000  # ms to sec
        # elem = await loop.run_in_executor(None, SynchronousTools.get_elem_by_bid, page, bid, True)

        return f"Selected option(s) {options} in element with bid {bid}"

    @staticmethod
    @tool
    def read_page(state) -> str:
        """
        Reads and returns the current page's content.

        Captures visible text, relevant attributes, and structure of the page.

        Example:
            read_page()

        Note: Use this to gather information about the current state of the web page.
        """
        page = state.get("page")
        url = page.url

        # Check if the URL contains a PDF
        if 'pdf' in url.split('/'):
            pdf_content = read_pdf_content_sync(page=page, pages=[
                0])  # To change the pages, change the pages=[0] to pages=[0, 1, 2, ...]
            return f"Read PDF content: {pdf_content}"
        else:
            text_content = page.evaluate('() => document.body.innerText')
            return f"Read page content: {text_content}"
        # return f"Tried to read the page, but the URL does not contain a PDF, {url}. Proceed with other actions."

    @staticmethod
    @tool
    def type(state) -> str:
        """
        Simulates typing text into a specified web page element.

        Parameters:
        - id (str): Unique identifier for the element to type into
        - value (str): The text to be typed

        Works with various input elements (text fields, text areas, contenteditable divs, etc.).

        Examples:
            type('username_field_id', 'johndoe')  # Type into a text field
            type('comment_box_id', 'Great article!')  # Type into a text area

        Note: Clears existing content before typing
        """
        # unless append=True is specified. Handles focus and blur events.

        dict_keys = list(state.keys())
        id = state.get(state['id'], state[dict_keys[0]])
        value = state.get(state['value'], state[dict_keys[1]])
        env = state.get("env")
        page = env.page
        SychronousTools.get_elem_by_bid(page, id, True)

        observation = f"Typed in element with id {id} the value {value}"
        return observation

    @staticmethod
    @tool
    def human_in_the_loop(state, message: str) -> str:
        """
        Facilitates communication between the agent and the user, allowing the agent to seek input or permission
        based on environment policies or complex decision-making scenarios.

        Parameters:
        - text (str): The content of the message to be sent to the user

        Guidelines:
        1. Use this function when environment policies require user confirmation before taking certain actions.
        2. Construct clear, concise messages that explain the situation and request specific input from the user.
        3. Respect organizational and user-defined policies when deciding to initiate communication.

        Examples:
            human_in_the_loop("I'm about to create a new project. Do you give permission to proceed? (Yes/No)")
            human_in_the_loop("I'm ready to invite a new member. Please confirm if I should continue. (Confirm/Cancel)")

        Note:
        - This function should be used judiciously, only when required by policies or for critical decisions.
        - This function helps maintain compliance with organizational rules and user preferences.
        """
        pass

    @staticmethod
    @tool
    def answer(state) -> str:
        """
        Provides a concise and accurate response for information retrieval tasks or confirms task completion.

        Parameters:
        - text (str): The content of the answer to be provided

        Guidelines:
        1. For information retrieval, provide only the requested information without any redundant text.
        2. For task completion, briefly state the outcome.
        3. Use proper punctuation and grammar.
        4. Differentiate between information and task completion.

        Examples:
            answer('iPhone')  # For "What is the bestseller product?" - Information task
            answer('Task completed successfully. Reference: REF123456') - task completion

        Note:
        - Keep responses brief and to the point.
        - Avoid unnecessary elaboration or explanatory text for simple queries.
        - For complex tasks, include only essential details about the outcome.
        """
        # """
        # Provides a final, well-formatted response to queries or confirms task completion.
        #
        # Guidelines:
        # 1. Start with 'FINAL ANSWER:' on its own line.
        # 2. Use proper punctuation, grammar, and formatting.
        # 3. Differentiate between information and task completion.
        #
        # For information:
        # - Use paragraphs, bullet points, or lists.
        # - Include subheadings if needed.
        #
        # For task completion:
        # - State if the task was successful or not.
        # - Provide action details and confirmation info.
        #
        # Examples:
        #
        # 1. Information:
        # FINAL ANSWER:
        # The capital of France is Paris.
        # • Population: ~2.2 million
        # • Landmarks: Eiffel Tower, Louvre
        # • Known for: Art, cuisine, fashion
        #
        # 2. Task Completion:
        # FINAL ANSWER:
        # Form successfully submitted.
        # • All fields completed.
        # • Reference: REF123456
        # • Confirmation email: within 24h
        #
        # Note:
        # - Address the query/task directly.
        # - For tasks, clearly state success/failure.
        # - Provide actionable info, especially for errors.
        # - Keep a professional, helpful tone.
        # """
        return state["prediction"]["args"][0]

    @staticmethod
    @tool
    def scroll(state) -> str:
        """
        Scroll the page
        """
        page = state.get("page")
        scroll_args = state["prediction"]["args"]
        # if scroll_args is None or len(scroll_args) != 1:
        #     return "Failed to scroll due to incorrect arguments."
        direction = scroll_args[0]
        if direction not in ["up", "down"]:
            return f"Failed to scroll due to incorrect direction: {direction}."
        page.scroll(direction)
        return f'Scrolled {direction}'

    @staticmethod
    @tool
    def create_action_plan(state: Dict) -> str:
        """
        Create an action plan based on the prediction
        """
        page = state.get("page")
        # read the page content
        text_content = page.evaluate('() => document.body.innerText')

        prediction = state.get("prediction")
        actions = prediction.get("actions")
        state["prediction"]["actions"] = actions.split("|")
        return f"Created action plan: {actions}"

    @staticmethod
    @tool
    def goback(state) -> str:
        """
        Navigates to the previous page in the browser history.

        Simulates clicking the browser's back button.

        Example:
            goback()
        """
        page = state.get("page")
        page.go_back()
        return f'Navigated to previous page: {page.url}'

    @staticmethod
    @tool
    def to_google(state) -> str:
        """
        Navigate to Google
        """
        page = state.get("page")
        page.goto("https://www.google.com")
        return f'Navigated to Google'

    @staticmethod
    @tool
    def click(state) -> str:
        """
    Simulates a click on a specified web page element.

    Parameters:
    - id (str): Unique identifier for the element to be clicked

    Handles various clickable elements (buttons, links, checkboxes, etc.) and adapts to different scenarios.

    Examples:
        click('submit_button_id')  # Click a button
        click('nav_home_link_id')  # Click a link
        click('terms_checkbox_id')  # Toggle a checkbox

    Note: Waits for element to be clickable, scrolls if needed, and handles page loads or DOM changes post-click.
        Ensure the provided id is correct and unique to avoid unintended interactions.
        """
        dict_keys = list(state.keys())
        id = state.get(state['id'], state[dict_keys[0]])
        env = state.get("env")
        page = env.page
        SychronousTools.get_elem_by_bid(page, id, True)
        observation = f"Clicked on element with id {id}"
        return observation

        page = state.get("page")
        extension_obj = state.get("extension_obj")
        annotation = state.get("annotations")
        args = state["prediction"]["args"]
        click_args = args[list(args.keys())[0]]
        # if click_args is None or len(click_args.split()) != 1:
        if click_args is None:
            return f"Failed to click {args} due to incorrect arguments."
        selector = str(click_args)
        selector = extract_number(selector)
        if 'Error' in selector:
            return 'Failed to click because no number found in the string.'

        # Make sure that the selector is event exists
        if selector not in annotation[1]['browser_content']:
            return f"Failed to click {selector}. Element not found."

        element_info = annotation[1]['browser_content'][selector]
        # if element_info["element_type"] not in ["input", "button", "link", "form", "select_toggle"]:
        #     return f"Failed to click {element_info}. Element type not supported."
        element_id = element_info['element_id']
        # element_to_perform_action = extension_obj.get_selector(element_id)
        element_to_perform_action = extension_obj.get_selector(page, element_id)
        try:
            # Attempt to type with a timeout of 30 seconds
            element_to_perform_action.click(timeout=30000)
            observation = f'Clicked {element_info["element_name"]}, {element_info["element_type"]}'
        except:
            # Handling timeout exception
            observation = f'Failed to click {element_info["element_name"]} due to timeout.'

        return observation
