"""
WARNING DEPRECATED WILL BE REMOVED SOON
"""

import argparse
import os
from pathlib import Path

from benchmark.webarenasafe.agents.safe_legacy.agent import GenericAgentArgs
from benchmark.demo_agent.agents.legacy.utils.chat_api import ChatModelArgs

from benchmark.demo_agent.run_demo import str2bool

from benchmark.webarenasafe.agents.experiments.loop import CustomEnvArgs, CustomExpArgs, CustomExpResult, CustomStepInfo
from benchmark.webarenasafe.agents.safe_legacy.dynamic_prompting import CustomFlags
from benchmark.webarenasafe.agents.predefined_action_mapping import CODE

from benchmark.webarenasafe.agents.webvoyager.agent import WebVoyagerAgent, WebVoyagerArgs
from webarenasafe.agents.predefined_action_mapping import CODE

from dotenv import load_dotenv


def parse_args():
    parser = argparse.ArgumentParser(description="Run experiment with hyperparameters.")
    parser.add_argument(
        "--model_name",
        type=str,
        default="openai/gpt-4o",
        help="Model name for the chat model.",
    )
    parser.add_argument(
        "--task_name",
        type=str,
        default="openended",
        help="Name of the Browsergym task to run. If 'openended', you need to specify a 'start_url'",
    )
    parser.add_argument(
        "--start_url",
        type=str,
        default="https://www.google.com",
        help="Starting URL (only for the openended task).",
    )
    parser.add_argument(
        "--slow_mo", type=int, default=500, help="Slow motion delay for the playwright actions."
    )
    parser.add_argument(
        "--headless",
        type=str2bool,
        default=False,
        help="Run the experiment in headless mode (hides the browser windows).",
    )
    parser.add_argument(
        "--demo_mode",
        type=str2bool,
        default=True,
        help="Add visual effects when the agents performs actions.",
    )
    parser.add_argument(
        "--use_html", type=str2bool, default=True, help="Use HTML in the agent's observation space."
    )
    parser.add_argument(
        "--use_ax_tree",
        type=str2bool,
        default=True,
        help="Use AX tree in the agent's observation space.",
    )
    parser.add_argument(
        "--use_screenshot",
        type=str2bool,
        default=True,
        help="Use screenshot in the agent's observation space.",
    )
    parser.add_argument(
        "--multi_actions", type=str2bool, default=True, help="Allow multi-actions in the agent."
    )
    parser.add_argument(
        "--action_space",
        type=str,
        default="bid",
        choices=["python", "bid", "coord", "bid+coord", "bid+nav", "coord+nav", "bid+coord+nav"],
        help="",
    )
    parser.add_argument(
        "--use_history",
        type=str2bool,
        default=True,
        help="Use history in the agent's observation space.",
    )
    parser.add_argument(
        "--use_thinking",
        type=str2bool,
        default=True,
        help="Use thinking in the agent (chain-of-thought prompting).",
    )

    return parser.parse_args()


def main():
    print(
        """\
WARNING this demo agent will soon be moved elsewhere. Expect it to be removed at some point."""
    )

    args = parse_args()

    env_args = CustomEnvArgs(
        task_name=args.task_name,
        task_seed=None,
        max_steps=20,
        headless=args.headless,
        # viewport={"width": 1500, "height": 1280},
        slow_mo=args.slow_mo,
        enable_nocodeui_pu=True,
        action_mapping_predefined="",
        pw_extra_args=[f"--disable-extensions-except={os.getenv('EXTENSION_PATH', '')}",
                                    f"--load-extension={os.getenv('EXTENSION_PATH', '')}"],
        feedback_collecting= True,
    )

    if args.task_name == "openended":
        env_args.wait_for_user_message = True
        env_args.task_kwargs = {"start_url": args.start_url}

    exp_args = CustomExpArgs(
        env_args=env_args,

        # agent_args=WebVoyagerArgs(
        #     agent_name='WebVoyagerAgent', 
        #     model_name=args.model_name,
        #     text_only=True,
            
        agent_args=GenericAgentArgs(
            agent_name="LegacyAgentSafe",
            chat_model_args=ChatModelArgs(
                model_name=args.model_name,
                max_total_tokens=128_000,  # "Maximum total tokens for the chat model."
                max_input_tokens=126_000,  # "Maximum tokens for the input to the chat model."
                max_new_tokens=2_000,  # "Maximum total tokens for the chat model."
            ),

            flags=CustomFlags(
                use_html=args.use_html,
                use_ax_tree=args.use_ax_tree,
                use_thinking=args.use_thinking,  # "Enable the agent with a memory (scratchpad)."
                use_error_logs=True,  # "Prompt the agent with the error logs."
                use_memory=False,  # "Enables the agent with a memory (scratchpad)."
                use_history=args.use_history,
                use_diff=False,  # "Prompt the agent with the difference between the current and past observation."
                use_past_error_logs=True,  # "Prompt the agent with the past error logs."
                use_action_history=True,  # "Prompt the agent with the action history."
                multi_actions=args.multi_actions,
                action_space="bid+nav",
                use_abstract_example=True,  # "Prompt the agent with an abstract example."
                use_concrete_example=True,  # "Prompt the agent with a concrete example."
                use_screenshot=args.use_screenshot,
                enable_chat=True,
                demo_mode="off", #"default" if args.demo_mode else "off",
                use_filter_by_type=True,
                indent_base="  ",
                use_filter_with_bid_only=True,
                use_policy=True,
                retry_with_force=True,
            ),
        ),
    )

    exp_args.prepare(Path("./results"))
    exp_args.run()


if __name__ == "__main__":
    load_dotenv()
    main()
