# fmt: off
import os
import sys

from llms.llm_utils import call_llm, visualize_prompt
from llms.prompt_utils import get_messages
from offline_experiments.vwa_specific import (get_intent_message_vwa,
                                              get_prompt_constructor,
                                              get_trace_data_vwa)
from utils.string_utils import safe_format

if __name__ == "__main__" and not __package__:  # @debug
    # Insert the parent directory into sys.path so that the package can be found
    parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    sys.path.insert(0, parent_dir)
    # Manually set the package name so that relative imports work
    __package__ = "offline_experiments"

# ===============================================
# LINK: Prompt parts
# ===============================================
# --- System prompt ---

sys_prompt_critic_base = f"""You are an intelligent agent tasked with supervising an assistant navigating a web browser to accomplish a web-based task. Your job is to evaluate the assistant's work, and provide feedback so it can progress towards the objective.
\n
## Here's the information you'll have:
### The objective: This is the task the assistant is trying to complete.
### Webpage screenshots: These are screenshots of the webpage, with each interactable element assigned a unique numerical id. Each bounding box and its respective id shares the same color.
{{trace_info}}{{summ_info}}{{k_info}}
\n
## Assistant's capabilities: To effectively analyze the assistant's work, consider the actions it can perform. These actions fall into the following categories:
### Page Operation Actions:
```click [id]```: Click on an element with a specific id on the webpage.
```type [id] [content] [enter_after]```: Type the content into the field with id. If `enter_after` is 0, the "Enter" key is not pressed after typing; otherwise, the "Enter" key is automatically pressed.
```hover [id]```: Hover over an element with id.
```press [key_comb]```: Press a keyboard key or key combination (e.g., delete, ctrl+a).
```scroll [down]``` or ```scroll [up]```: Scroll the webpage up or down.

### Tab Management Actions:
```new_tab```: Open a new, empty browser tab.
```tab_focus [tab_index]```: Switch the browser's focus to a specific tab using its index.
```close_tab```: Close the currently active tab.

### URL Navigation Actions:
```goto [url]```: Navigate to a specific URL.
```go_back```: Navigate to the previously viewed page.
```go_forward```: Navigate to the next page (if a previous 'go_back' action was performed).

### Completion Action:
```stop [answer]```: This action is issued when the task is believed to be complete. If the task requires a text-based answer, the answer is provided within the brackets. If the task is deemed impossible, this action is issued optionally with a reason why.

{{privileged_hint}}\n
## To be successful, it is very important to follow the following rules:
{{rule}}
"""

trace_info_actions = "\n### The execution trace: This is a sequence of webpage screenshots paired with the assistant's actions, detailing the web navigation so far."
k_info = "\n### General web knowledge: This is a general description of how tasks like this are typically accomplished on the web."

k_rule_efficiency_vision_v2 = """
1. You must not assume the assistant's work is correct or incorrect beforehand. You should come up with your own opinion based on the information provided.
2. You must connect the dots between the objective and the information provided.
3. Your evaluation should be based on task completion; don't worry about efficiency for now.
4. Give utmost importance to the visual information provided, especially the screenshots in the execution trace."""
# 5. Use the General web knowledge to guide your judgment, but also consider the context of the specific task given to you."""


# For example, if we are giving the retrieved knowledge, we create the final sys_prompt as follows:
sys_prompt = safe_format(sys_prompt_critic_base, rule=k_rule_efficiency_vision_v2, trace_info=trace_info_actions, k_info=k_info)
sys_prompt_noK = safe_format(sys_prompt_critic_base, rule="", trace_info=trace_info_actions)
# This will fill the above prompt with a third rule at the end. Print it out to see.


eval_criteria_ternary = """Now please provide your response.
\n
## Here is the evaluation criteria:
SUCCESS: The agent executed **all** that's necessary to complete the objective. The task is fully accomplished.
PARTIAL SUCCESS: The agent executed **part of** what's necessary to complete the objective. The task is partially accomplished.
FAILURE: The agent executed **mostly incorrect** steps. The task is not accomplished, and major revisions are needed."""

response_format_basicCoT = """## Provide your response as follows:
REASONING: [Comprehensive step-by-step reasoning to come up with your evaluation and feedback]
EVALUATION: [Your evaluation following the evaluation criteria]
FEEDBACK: [Feedback so the assistant can progress towards the objective]"""


prompt_constructor = get_prompt_constructor()
domain, task_id = "shopping", 402
html_path = f"./experiments/gemini-2.5-flash-001/no-verifier-2025-06-19/{domain}/htmls/render_{task_id}.html"
trace_data = get_trace_data_vwa(html_path, task_id)


prompt_constructor.instruction["meta_data"]["use_assistant_utterance"] = False
prompt_constructor.instruction["meta_data"]["use_low_level_actions_env_parsed"] = True
traj_actions = prompt_constructor.build_interaction_history(trace_data["trajectory"], trace_data["meta_data"], idxs_history=list(range(len(trace_data["trajectory"].states))))

prompt_constructor.instruction["meta_data"]["use_assistant_utterance"] = False
prompt_constructor.instruction["meta_data"]["use_low_level_actions_env_parsed"] = False
traj_no_agent = prompt_constructor.build_interaction_history(trace_data["trajectory"], trace_data["meta_data"], idxs_history=list(range(len(trace_data["trajectory"].states))))

prompt_constructor.instruction["meta_data"]["use_assistant_utterance"] = True
prompt_constructor.instruction["meta_data"]["use_low_level_actions_env_parsed"] = False
traj_utterances = prompt_constructor.build_interaction_history(trace_data["trajectory"], trace_data["meta_data"], idxs_history=list(range(len(trace_data["trajectory"].states))))

visualize_prompt(traj_utterances, "./vis.html", verbose=True)

#=========== K retrieval ===============================================
sys_prompt_k_expert_v4 = f"""You are an expert on web navigation. Your job is to provide a comprehensive and **detailed** description of how tasks like the ones provided to you are typically accomplished on the web.
\n
## Here's the information you'll have:
### Objective: This is an english description of the task, possibly accompanied by images that provide more context on what must be accomplished.
### Screenshots: These are screenshots of the webpage, giving you context of the state of the navigation for the task at hand."""

k_retrieval_expert_question = """Now please provide your response: how tasks like these are typically accomplished on the web?"""

msg_intent = get_intent_message_vwa(trace_data, add_state_idxs=[0], state_img_intros=["Description: Initial webpage screenshot"])
prompt = [
    {"role": "system", "content": sys_prompt_k_expert_v4},
   msg_intent,
   k_retrieval_expert_question,
]
# visualize_prompt(prompt, "./vis.html", verbose=True)

gen_config = {
    "model": "gemini-2.5-flash",
    "temperature": 1,
    "top_p": 0.001,
    "top_k": 64,
    "max_tokens": 8192,
    "num_generations": 1,
    "thinking_budget": 0,
}

# visualize_prompt(prompt, "./vis.html", verbose=True)

# api_responses, model_messages = call_llm(
#     gen_config,
#     prompt,
# )
# k = model_messages[0].text()

# # 1) Describe the navigation, only images.
# sys_prompt_vis = f"""You are an intelligent and helpful assistant.
# You will be given a sequence of screenshots representing a navigation on the web. Your job is to carefully understand each screenshot and comprehensively describe the navigation so far.
# \n
# ## Here's the information you'll have:
# ### Webpage screenshots: These are screenshots of the webpage, with each interactable element assigned a unique numerical id. Each bounding box and its respective id shares the same color.
# """


# vis_retrieval = """Now please provide your response:"""

# vis_injection = """EXECUTION TRACE DESCRIPTION:
# {k}"""


# # prompt = get_messages(inputs=[traj_no_agent, vis_retrieval], sys_prompt=sys_prompt_vis)
# prompt = get_messages(inputs=[traj_actions, vis_retrieval], sys_prompt=sys_prompt_vis)
# visualize_prompt(prompt, "./vis.html", verbose=True)


# _, model_messages = call_llm(gen_config, prompt)
# exec_trace_description = model_messages[0].text()
msg_intent_no_state = get_intent_message_vwa(trace_data)

# k_injection = """### General web knowledge:
# {k}"""

# vis_injection = """### Execution trace description:
# {k}"""

# final_k = k_injection.format(k=k)
# final_vis = vis_injection.format(k=exec_trace_description)

# all_ks_injection = """
# ## Supplemental information:
# {final_k}
# \n
# {final_vis}
# """

# response_format_no_cot = """## Provide your response as follows:
# EVALUATION: [Your evaluation following the evaluation criteria]
# FEEDBACK: [Feedback so the assistant can progress towards the objective]"""


# response_format_desc_cot = """## Provide your response as follows:
# [Carefully understand each screenshot and comprehensively describe the navigation so far.]
# EVALUATION: [Your evaluation following the evaluation criteria]
# FEEDBACK: [Feedback so the assistant can progress towards the objective]"""


# response_format_desc_cot_v2 = """## Provide your response as follows:
# EXECUTION TRACE DESCRIPTION: [Carefully understand and describe the assistant's work]
# COMPARISON: [Compare the assistant's work with the general web knowledge. Make a step by step comparison of what the assistant did and what is expected]
# RELEVANCE: [Which of the steps missing are relevant in this context? If the steps are not relevant, explain why]
# CONCLUSION: [Based on the assistant's work, comparison to general web knowledge, and the objective, what are your conclusions?]
# EVALUATION: [Your evaluation following the evaluation criteria]
# FEEDBACK: [Feedback so the assistant can progress towards the objective]"""



prompt = get_messages(inputs=[
    # traj_no_agent,
    msg_intent_no_state,
    traj_actions,
    # traj_utterances,
    # safe_format(all_ks_injection, final_k=final_k, final_vis=final_vis),
    # safe_format(all_ks_injection, final_vis=final_vis),
    # safe_format(all_ks_injection, final_k=final_k),
    # [eval_criteria_ternary, response_format_desc_cot_v2],
    [eval_criteria_ternary, response_format_basicCoT],
    # [eval_criteria_ternary],
], sys_prompt=sys_prompt)

# gen_config["temperature"] = 1
# gen_config["top_p"] = .8
# gen_config["top_k"] = 64
# gen_config["num_generations"] = 5


visualize_prompt(prompt, "./vis.html", verbose=True)
_, model_messages = call_llm(gen_config, prompt, conversation_dir="./abias-examples", call_id=f"{domain}-{task_id}")
print(model_messages[0].text())


for m in model_messages:
    print(m.text())
    print("-"*100)
