import dataclasses
import re
from dotenv import load_dotenv
from browsergym.core.action.highlevel import HighLevelActionSet
from browsergym.experiments import Agent, AbstractAgentArgs
from browsergym.utils.obs import flatten_axtree_to_str

from agentS.consts import MODEL_NAME


class DemoAgent(Agent):
    """A basic agent using OpenAI API, to demonstrate BrowserGym's functionalities."""

    action_set = HighLevelActionSet(
        subsets=["chat", "bid"],  # define a subset of the action space
        # subsets=["chat", "bid", "coord"] # allow the agent to also use x,y coordinates
        strict=False,  # less strict on the parsing of the actions
        multiaction=True,  # enable to agent to take multiple actions at once
        demo_mode="default",  # add visual effects
    )

    # use this instead to allow the agent to directly use Python code
    # action_set = PythonActionSet())

    def obs_preprocessor(self, obs: dict) -> dict:
        return {
            "goal": obs["goal"],
            "axtree_txt": flatten_axtree_to_str(obs["axtree_object"]),
        }

    def __init__(self, model_name) -> None:
        super().__init__()
        self.model_name = model_name

        load_dotenv()

        from openai import OpenAI

        self.openai_client = OpenAI()

    def get_action(self, obs: dict) -> tuple[str, dict]:
        system_msg = f"""\
# Instructions
Review the current state of the page and all other information to find the best
possible next action to accomplish your goal. Your answer will be interpreted
and executed by a program, make sure to follow the formatting instructions.

# Goal:
{obs["goal"]}"""

        prompt = f"""\
# Current Accessibility Tree:
{obs["axtree_txt"]}

# Action Space
{self.action_set.describe(with_long_description=False, with_examples=True)}

Here is an example with chain of thought of a valid action when clicking on a button:
"
In order to accomplish my goal I need to click on the button with bid 12.
```click("12")```
"

Only return one action at a time.
"""

        # query OpenAI model
        response = self.openai_client.chat.completions.create(
            model=self.model_name,
            temperature=1,
            max_tokens=256,
            messages=[
                {"role": "system", "content": system_msg},
                {"role": "user", "content": prompt},
            ],
        )

        def extract_content(text):
            # Use regular expression to find content between the backticks
            matches = re.findall(r'```(.*?)```', text, re.DOTALL)
            return matches[0]

        action = response.choices[0].message.content
        return extract_content(action)


@dataclasses.dataclass
class DemoAgentArgs(AbstractAgentArgs):
    """
    This class is meant to store the arguments that define the agent.

    By isolating them in a dataclass, this ensures serialization without storing
    internal states of the agent.
    """

    def __init__(self, llm_type, sync=True, architecture='general'):
        self.model_name = MODEL_NAME
        self.llm_type = llm_type
        self.sync = sync
        self.architecture = architecture

    def make_agent(self):
        return DemoAgent(model_name=self.model_name)
