import argparse
import json
import re

from palm2 import generate_from_palm2_chat_completion


def main(args):
    with open(args.input_file, "r") as file:
        data = json.load(file)
    i_dx = 0
    for entry in data:
        intent = entry["conversations"][0]["value"]["intent"]
        obs = entry["conversations"][0]["value"]["obs"]
        action = entry["conversations"][1]["value"]
        if "TYPE" not in action:
            continue

        i_dx += 1
        if i_dx == 1:
            continue

        prompt = """Summarize in a brief sentence what the action is doing on the web page: """
        prompt += "action=" + action + "."
        prompt += "webpage=" + obs + "."

        messages = {}
        messages[
            "context"
        ] = """You are a helpful assistant. You always think before you respond, and you give succient, accurate, easy to understand answers."""
        messages["example"] = [
            """''Summarize in a brief sentence what the action is doing on the web page: action=CLICK [132].webpage=[1] RootWebArea \'Home | MTA\' focused: True\n\\t[78] region \'Main menu region\'checkbox \'Search\' checked: false\n\\t[121] img \'Search\'\n\\t[128] link \'MTA\'\n\\t[131] link \'Schedules\'\n\\t[132] link \'Maps\'\n\\t[133] link \'Fares & Tolls\'\n\\t[134] link \'Planned Work\'\n\\t[12] main \'\'\n\\t\\t[63] complementary \'\'\n\\t\\t\\t[461] heading \'Plan a Trip\'\n\\t\\t\\t[462] region \' Plan a Trip \'\n\\t\\t\\t\\t[945]""",
            "click on the maps",
            """'Summarize in a brief sentence what the action is doing on the web page: action=TYPE [109] [kiss].webpage={\'text\': \'Tab 0: Playwright Trace Viewer | Tab 1 (current): ULTIMATE GUITAR TABS - 1,100,000 songs catalog with free Chords, Guitar Tabs, Bass Tabs, Ukulele Chords and Guitar Pro Tabs!\n\n[1] RootWebArea \\\'ULTIMATE GUITAR TABS - 1,100,000 songs catalog with free Chords, Guitar Tabs, Bass Tabs, Ukulele Chords and Guitar Pro Tabs!\\\' focused: True\n\\t[18] link \\\'Create your Account and get \\\\xa0Pro Access 80% OFF 0 days : 21 hrs : 35 min : 15 sec SIGN UP\\\'\n\\t\\t[51] button \\\'SIGN UP\\\'\n\\t[19] button \\\'\\\'\n\\t[109] textbox \\\'Enter artist name or song title\\\' required: True\n\\t[39] button \\\'SEARCH\\\'\n\\t[16] main \\\'\\\'\n\\t\\t[131] article \\\'\\\'\n\\t\\t\\t[157] HeaderAsNonLandmark \\\'\\\'\n\\t\\t\\t\\t[356] button \\\'SIGN UP\\\'\n\\t\\t\\t\\t[357] button \\\'LOG IN\\\'\n\\t\\t[158] HeaderAsNonLandmark \\\'\\\'\n\\t\\t\\t[160] StaticText \\\'Users online: 5,097\\\'\n\\t\\t[163] StaticText \\\'4,625 members\\\'\n\\t\\t[164] StaticText \\\'472 guests\\\'\n\\t\\t""",
            'enter the phrase "kiss" into artist name or song title',
        ]
        messages["messages"] = [prompt]

        if args.model == "palm2":
            history = generate_from_palm2_chat_completion(
                messages=messages,
                model="models/chat-bison-001",
                temperature=0.8,
                max_tokens=-1,
                top_p=0.5,
                context_length=8000,
            )
            print(history)
        assert False


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--model", type=str, default=None, help="palm2/gpt/llama.."
    )

    parser.add_argument(
        "--input-file", type=str, default=None, help="input-file"
    )
    args = parser.parse_args()

    main(args)
