import json
from agent import *
from argparse import ArgumentParser
import os

with open("./_generate_task_prompt.txt", "r") as f:
    system_prompt = f.read()


class Instruction(BaseModel):
    content: str = Field(description="the instruction for the task")
    degreeOfDetail: int = Field(description="the degree of detail for the instruction, from 1 to 10")
    armMention: bool = Field(description="whether the instruction mentions arm, whether by schema or by fixed text")
    numOfWords: int = Field(description="the number of words in the instruction")


class InstructionFormat(BaseModel):
    stepsOfTask: List[str] = Field(
        description=
        "split the task into small steps, and make sure each step is explicitly or implicitly mentioned in each of the instructions.Avoid using adjectives in it!"
    )
    instructions: List[Instruction] = Field(
        description="several different text instructions describing this same task here")


def make_prompt_generate(detailed_task, preferences, schema, instruction_num):
    system_prompt_schema = ""
    if schema:
        with open("./_generate_task_prompt_schema.txt", "r") as f:
            system_prompt_schema = f.read()
    messages = [
        {
            "role": "system",
            "content": system_prompt + "\n" + system_prompt_schema
        },
        {
            "role":
            "user",
            "content": [
                # {"type":"image_url","image_url":{"url":f"data:image/png;base64,{imgStr}"},
                {
                    "type": "text",
                    "text": f"The detailed task description for you to abstract is {detailed_task}",
                },
                {
                    "type": "text",
                    "text": f"For each instruction, you should follow the preference: {preferences}",
                },
                {
                    "type": "text",
                    "text": f"Generate {instruction_num} alternative descriptions based on the input.",
                },
            ],
        },
    ]
    if schema:
        messages[1]["content"].append({
            "type": "text",
            "text": f"The object schema for you to abstract is {schema}",
        })
    result = generate(messages, InstructionFormat)
    result_dict = result.model_dump()
    print(json.dumps(result_dict, indent=2, ensure_ascii=False))
    insList = []
    for ins in result.instructions:
        insList.append(ins.content)
    return insList


def generate_task_description(task_name, instruction_num):
    with open(f"./task_instruction/{task_name}.json", "r") as f:
        task_info_json = f.read()
    # print(task_info_json)
    task_info = json.loads(task_info_json)
    if "seen" not in task_info.keys():
        task_info["seen"] = []
    if "unseen" not in task_info.keys():
        task_info["unseen"] = []
    for required_keys in [
            "full_description",
            "preference",
    ]:  # schema can be empty to disable it
        if (not task_info.get(required_keys, "") or task_info.get(required_keys, "") == ""):
            print(f"{required_keys} is not in the ./task_instruction/{task_name}.json or is empty")
            return
    result = make_prompt_generate(
        task_info["full_description"],
        task_info["preference"],
        task_info["schema"],
        instruction_num,
    )
    print(f'{task_name} generated {len(result)} descriptions with length {len("".join(result))}')
    task_info["seen"].extend(result[2:])
    task_info["unseen"].extend(result[0:2])
    # task_info['seen'] = result[2:]
    # task_info['unseen'] = result[0:2]
    with open(f"./task_instruction/{task_name}.json", "w") as f:
        json.dump(task_info, f, indent=2, ensure_ascii=False)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("task_name", type=str, default="beat_block_hammer")
    parser.add_argument("instruction_num", type=int, default=11)
    usr_args = parser.parse_args()
    task_name = usr_args.task_name
    instruction_num = usr_args.instruction_num
    if instruction_num % 12 != 0:
        print("instruction_num should be divisible by 12")
        exit()
    for i in range(instruction_num // 12):
        generate_task_description(task_name, 12)
