#!/usr/bin/env python3
"""
A script to collect all trajectory.json files and generate an HTML representation.

For each trajectory, the generated HTML shows:
  - A header containing the task id, observed score, domain and objective.
  - The initial screenshot (if available).
  - A sequence of steps, where each step shows:
      • The last screenshot (if more than one screenshot exists, pick the last).
      • The generated text (e.g., the Thought/Action).
      • (Optional) a verifier loop if the key "critique_executor_loop_utterances" or "verifier_loop" exists.

This is inspired by the render() method in browser_env/helper_functions.py.
"""

import concurrent.futures
import json
import os
import sys
from pathlib import Path

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__ = "osw_utils"


# Updated HTML template with escaped curly braces in the CSS.
HTML_TEMPLATE = """
<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>Trajectories</title>
    <style>
        pre {{
            white-space: pre-wrap;
            word-wrap: break-word;
        }}
        .trajectory {{
            border: 1px solid #ccc;
            margin: 20px;
            padding: 10px;
        }}
        .trajectory h2 {{
            margin-top: 0;
        }}
        .step {{
            border: 1px solid #ddd;
            margin: 10px 0;
            padding: 10px;
        }}
        .executor_utterance {{
            background-color: #D8CEF6;
            padding: 5px;
            margin: 5px 0;
        }}
        .critique_utterance {{
            background-color: #F0E6FF;
            padding: 5px;
            margin: 5px 0;
        }}
        .executor_critique_loop {{
            border: 1px dashed #aaa;
            margin: 10px 0;
            padding: 5px;
        }}
        .verifier_loop {{
            border: 1px dashed #ff9900;
            margin: 10px 0;
            padding: 5px;
        }}
        .verifier_output {{
            background-color: #ffe6cc;
            padding: 5px;
            margin: 5px 0;
        }}
    </style>
</head>
<body>
{body}
</body>
</html>
"""

from osw_utils.osw_utils import TrajectoryView
from utils.file_utils import find_files
from utils.image_utils import any_to_b64


def trajectory_to_html(traj_file: Path) -> str:
    try:
        trajectory_view = TrajectoryView(trajectory_path=traj_file, to_english=False)
    except Exception as e:
        print(f"Error reading {traj_file}: {e}")
        return ""

    # Gather header info.
    task_id = trajectory_view.task_id
    score = trajectory_view.score
    domain = trajectory_view.domain
    objective = trajectory_view.objective

    # Compute number of verifier loops by checking each state's keys.
    num_verifier_loops = sum(1 for state in trajectory_view.states if "verifier_loop" in state)

    html_parts = []
    # Start a single container for both header and step details.
    html_parts.append("<div class='trajectory'>")

    # Header information including the verifier loops count.
    html_parts.append(f"<p><strong>Task id:</strong> {task_id}</p>")
    html_parts.append(f"<p><strong>Score:</strong> {score}</p>")
    html_parts.append(f"<p><strong>Domain:</strong> {domain}</p>")
    html_parts.append(f"<p><strong>Objective:</strong> {objective}</p>")
    html_parts.append(f"<p><strong>Verifier Loops:</strong> {num_verifier_loops}</p>")

    # Process each step and generate a block for each.
    for t, (state, action) in enumerate(zip(trajectory_view.states, trajectory_view.actions)):
        html_parts.append("<div class='step'>")
        html_parts.append(f"<h3>State {t}</h3>")

        # For the screenshot, check the 'screenshots' list if available.
        screenshots = state["observation"]["images"]
        if screenshots:
            # Pick the last screenshot.
            last_ss = screenshots[-1]
            html_parts.append(f"<img src='{any_to_b64(last_ss)}' style='max-width:50vw; height:auto;'/><br>")
        else:
            html_parts.append("<p>[No screenshot provided]</p>")

        # Add the generated text.
        generated_text = action["texts"][-1]
        html_parts.append(f"<pre>{generated_text}</pre>")

        # Check if a verifier loop exists in the 'verifier_loop' key.
        if "verifier_loop" in state:
            loop = state["verifier_loop"]
            html_parts.append("<div class='verifier_loop'>")
            html_parts.append("<h4>Verifier Loop</h4>")
            for item in loop:
                verifier_text = item.get("verifier", "")
                generator_text = item.get("agent", "")
                if verifier_text:
                    verifier_text_html = verifier_text.replace("\n", "<br>")
                    html_parts.append(f"<div class='verifier_output'><pre>{verifier_text_html}</pre></div>")
                if generator_text:
                    generator_text_html = generator_text.replace("\n", "<br>")
                    html_parts.append(f"<div class='generator_output'><pre>{generator_text_html}</pre></div>")
            html_parts.append("</div>")
        html_parts.append("</div>")
    html_parts.append("</div>")  # end .trajectory

    return "\n".join(html_parts)


def process_trajectory(traj_file: Path | str) -> str:
    """
    Process a single trajectory file: convert it to HTML and write it to disk.
    Returns the path of the output HTML file as a string if successful, or an empty string if not.
    """
    traj_file = Path(traj_file)
    html_fragment = trajectory_to_html(traj_file)
    if not html_fragment:
        return ""
    final_html = HTML_TEMPLATE.format(body=html_fragment)
    output_file = traj_file.with_suffix(".html")
    with open(output_file, "w", encoding="utf-8") as f:
        f.write(final_html)
    return str(output_file)


def main():
    # Use the first argument as the base directory; default to "trace_osworld"
    if len(sys.argv) > 1:
        base_dir = Path(sys.argv[1])
    else:
        base_dir = Path("trace_osworld")

    if not base_dir.exists():
        print(f"Error: Base directory '{base_dir}' does not exist.")
        sys.exit(1)

    traj_files = find_files(base_dir, "trajectory.json")

    if not traj_files:
        print(f"No trajectory.json files found under '{base_dir}'.")
        sys.exit(0)

    print(f"Found {len(traj_files)} trajectory file(s). Processing...")

    # Parallelize processing using ProcessPoolExecutor.
    with concurrent.futures.ProcessPoolExecutor() as executor:
        # Submit all trajectory files for processing.
        futures = {executor.submit(process_trajectory, traj_file): traj_file for traj_file in traj_files}
        for future in concurrent.futures.as_completed(futures):
            result = future.result()
            if result:
                print(f"HTML trace written to: {result}")


if __name__ == "__main__":
    main()
