import base64
import copy
import io
import json
import re
import html

from pathlib import Path
from typing import Any

from bs4 import BeautifulSoup

from PIL import Image


from agent.prompts import *
from browser_env import (
    Action,
    ActionTypes,
    ObservationMetadata,
    StateInfo,
    action2str,
)

HTML_TEMPLATE = """
<!DOCTYPE html>
<head>
    <style>
        pre {{
            white-space: pre-wrap;
            word-wrap: break-word;
        }}
    </style>
</head>
<html>
    <body>
     {body}
    </body>
</html>
"""


def get_render_action(
    action: Action,
    observation_metadata: dict[str, ObservationMetadata],
    action_set_tag: str,
) -> str:
    """Parse the predicted actions for rendering purpose. More comprehensive information"""
    match action_set_tag:
        case "id_accessibility_tree":
            text_meta_data = observation_metadata["text"]
            if action["element_id"] in text_meta_data["obs_nodes_info"]:
                node_content = text_meta_data["obs_nodes_info"][action["element_id"]]["text"]
            else:
                node_content = "No match found"

            action_str = f"<div class='raw_parsed_prediction' style='background-color:grey'><pre>{action['raw_prediction']}</pre></div>"
            action_str += f"<div class='action_object' style='background-color:grey'><pre>{repr(action)}</pre></div>"
            action_str += f"<div class='parsed_action' style='background-color:yellow'><pre>{action2str(action, action_set_tag, node_content)}</pre></div>"

        case "som":
            meta_data = observation_metadata["image"]
            if action["element_id"] in meta_data["obs_nodes_semantic_info"]:
                node_content = meta_data["obs_nodes_semantic_info"][action["element_id"]]
            else:
                node_content = "No match found"

            action_str = f"<div class='raw_parsed_prediction' style='background-color:grey'><pre>{action['raw_prediction']}</pre></div>"
            action_str += f"<div class='action_object' style='background-color:grey'><pre>{repr(action)}</pre></div>"
            action_str += f"<div class='parsed_action' style='background-color:yellow'><pre>{action2str(action, action_set_tag, node_content)}</pre></div>"

        case "playwright":
            action_str = action["pw_code"]
        case _:
            raise ValueError(f"Unknown action type {action['action_type'], action_set_tag}")

    return action_str


def get_action_description(
    action: Action,
    observation_metadata: dict[str, ObservationMetadata],
    action_set_tag: str,
    prompt_constructor: PromptConstructor | None = None,
) -> str:
    """Generate the text version of the predicted actions to store in action history for prompt use.
    May contain hint information to recover from the failures"""

    match action_set_tag:
        case "id_accessibility_tree":
            text_meta_data = observation_metadata["text"]
            # FIXME: if the observation type is `som_image` observation_metadata["text"] will be empty, giving wrong hint
            # removed  `id_accessibility_tree` + `som_image` possibility
            if action["action_type"] in [
                ActionTypes.CLICK,
                ActionTypes.HOVER,
                ActionTypes.TYPE,
            ]:
                action_name = str(action["action_type"]).split(".")[1].lower()
                if action["element_id"] in text_meta_data["obs_nodes_info"]:
                    node_content = text_meta_data["obs_nodes_info"][action["element_id"]]["text"]
                    node_content = " ".join(node_content.split()[1:])
                    action_str = action2str(action, action_set_tag, node_content)
                else:
                    action_str = f"Attempt to perfom \"{action_name}\" on element \"[{action['element_id']}]\" but no matching element found. Please check the observation more carefully."
                    action["invalid"] = True
            else:
                if action["action_type"] == ActionTypes.NONE and prompt_constructor is not None:
                    action_splitter = prompt_constructor.instruction["meta_data"]["action_splitter"]
                    action_str = f'The previous prediction you issued was "{action["raw_prediction"]}". However, the format was incorrect. Ensure that the action is wrapped inside a pair of {action_splitter} and enclose arguments within [] as follows: {action_splitter}action [arg] {action_splitter}.'
                    action["invalid"] = True
                else:
                    action_str = action2str(action, action_set_tag, "")

        case "som":
            meta_data = observation_metadata["image"]
            if action["action_type"] in [
                ActionTypes.CLICK,
                ActionTypes.HOVER,
                ActionTypes.TYPE,
            ]:
                action_name = str(action["action_type"]).split(".")[1].lower()
                if action["element_id"] in meta_data["obs_nodes_semantic_info"]:
                    node_content = meta_data["obs_nodes_semantic_info"][action["element_id"]]
                    action_str = action2str(action, action_set_tag, node_content)
                else:
                    action_str = f"Attempt to perfom \"{action_name}\" on element \"[{action['element_id']}]\" but no matching element found. Please check the observation more carefully."
                    action["invalid"] = True
            else:
                if action["action_type"] == ActionTypes.NONE and prompt_constructor is not None:
                    action_splitter = prompt_constructor.instruction["meta_data"]["action_splitter"]
                    action_str = f'The previous prediction you issued was "{action["raw_prediction"]}". However, the format was incorrect. Ensure that the action is wrapped inside a pair of {action_splitter} and enclose arguments within [] as follows: {action_splitter}action [arg] {action_splitter}.'
                    action["invalid"] = True
                else:
                    action_str = action2str(action, action_set_tag, "")

        case "playwright":
            action_str = action["pw_code"]

        case _:
            raise ValueError(f"Unknown action type {action['action_type']}")

    return action_str


class RenderHelper(object):
    """Helper class to render text and image observations and meta data in the trajectory"""

    def __init__(self, config_file: str, result_dir: str, action_set_tag: str, prefix: str = "") -> None:
        with open(config_file, "r") as f:
            _config = json.load(f)
            _config_str = ""
            for k, v in _config.items():
                _config_str += f"{k}: {v}\n"
            _config_str = f"<pre>{_config_str}</pre>\n"
            task_id = _config["task_id"]

        self.action_set_tag = action_set_tag

        if prefix:
            self.render_file = open(Path(result_dir) / f"render_{prefix}_{task_id}.html", "a+", encoding="utf-8")
        else:
            self.render_file = open(Path(result_dir) / f"render_{task_id}.html", "a+", encoding="utf-8")
        self.render_file.truncate(0)
        # write init template
        self.render_file.write(HTML_TEMPLATE.format(body=f"{_config_str}"))
        self.render_file.read()
        self.render_file.flush()

    def render(
        self,
        action: Action,
        state_info: StateInfo,
        meta_data: dict[str, Any],
        render_screenshot: bool = False,
        additional_text: list[str] | None = None,
    ) -> None:
        """Render the trajectory"""

        # Deep copy the action
        action_copy = copy.deepcopy(action)

        # text observation
        observation = state_info["observation"]
        text_obs = observation["text"]
        info = state_info["info"]
        new_content = f"<h2>New Page</h2>\n"
        new_content += (
            f"<h3 class='url'><a href={state_info['info']['page'].url}>URL: {state_info['info']['page'].url}</a></h3>\n"
        )

        text_obs_escaped = html.escape(text_obs)

        new_content += f"<div class='state_obv'><pre>{text_obs_escaped}</pre><div>\n"

        if render_screenshot:
            # image observation
            img_obs = observation["image"]
            image = Image.fromarray(img_obs)
            byte_io = io.BytesIO()
            image.save(byte_io, format="PNG")
            byte_io.seek(0)
            image_bytes = base64.b64encode(byte_io.read())
            image_str = image_bytes.decode("utf-8")
            new_content += f"<img src='data:image/png;base64,{image_str}' style='width:50vw; height:auto;'/>\n"

        # meta data
        new_content += (
            f"<div class='prev_action' style='background-color:pink'>{meta_data['action_str_history'][-1]}</div>\n"
        )

        if "early_stop" in action_copy:
            new_content += (
                f"<div class='early_stop' style='background-color:red'>Early stop: {action_copy['early_stop']}</div>\n"
            )

        # additional text
        if additional_text:
            for text_i, text in enumerate(additional_text):
                # Alternate background color between light green and light blue
                if text_i % 2 == 0:
                    bg_color = "#87CEFA"
                else:
                    bg_color = "#98FB98"
                new_content += (
                    f"<div class='additional_text' style='background-color: {bg_color}'>#{text_i+1}: {text}</div>\n"
                )

        # If critique_executor_loop_utterances is not empty, extract it and delete the key from Action (to render the action object using less space)
        executor_utterances, critique_utterances = None, None
        if "critique_executor_loop_utterances" in action_copy:
            executor_utterances = action_copy["critique_executor_loop_utterances"]["executor"]
            critique_utterances = action_copy["critique_executor_loop_utterances"]["critique"]
            del action_copy["critique_executor_loop_utterances"]

        # Add raw utterance and action object
        action_str = get_render_action(
            action_copy,
            info["observation_metadata"],
            action_set_tag=self.action_set_tag,
        )
        # Add string representation of action with yellow background
        action_str = f"<div class='predict_action'>{action_str}</div>"
        new_content += f"{action_str}\n"

        # Add critique-executor loop
        if executor_utterances is not None and critique_utterances is not None:
            new_content += "<div class='executor_critique_loop'>"
            new_content += "<h4>Executor Critique Loop</h4>"
            for i in range(len(executor_utterances)):
                executor_text = executor_utterances[i].replace("\n", "<br>")
                new_content += f"<div class='executor_utterance' style='background-color: #D8CEF6'><strong>Executor:</strong> {executor_text}</div>\n"
                if i < len(critique_utterances):
                    critique_text = critique_utterances[i].replace("\n", "<br>")
                    new_content += f"<div class='critique_utterance' style='background-color: #F0E6FF'><strong>Critique:</strong> {critique_text}</div>\n"
            new_content += "</div>"

        # add new content
        self.render_file.seek(0)
        html_content = self.render_file.read()
        soup = BeautifulSoup(html_content, "html.parser")
        body_tag = soup.body

        if body_tag:
            html_body = str(body_tag)
            html_body += new_content
        else:
            # Handle the case when <body> tag is not found
            html_body = new_content

        html_content = HTML_TEMPLATE.format(body=html_body)
        self.render_file.seek(0)
        self.render_file.truncate()
        self.render_file.write(html_content)
        self.render_file.flush()

    def close(self) -> None:
        self.render_file.close()
