import argparse
import concurrent.futures
import html
import json
import os
import sys
import traceback
from pathlib import Path

import pandas as pd

from trajectory_utils.trajectory import VWA_Trajectory
from utils.file_utils import find_files
from utils.image_utils import any_to_b64

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


# <!DOCTYPE html>
# <head>
#     <style>
#         pre {{
#             white-space: pre-wrap;
#             word-wrap: break-word;
#         }}
#     </style>
# </head>
# <html>
#     <body>
#      {body}
#     </body>
# </html>
# """


HTML_TEMPLATE = """
<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>Trajectories</title>
    <style>
        body {{
            font-family: monospace;
            background-color: #ffffff;
            color: #000000;
        }}
        pre {{
            white-space: pre-wrap;
            word-wrap: break-word;
        }}
        .executor_utterance {{
            background-color: #D8CEF6;
        }}
        .verifier_loop {{
            border: 1px dashed #ff9900;
        }}
        .verifier_output {{
            background-color: #ffe6cc;
        }}
        .predict_action {{
            background-color: yellow;
        }}
        .raw_parsed_prediction {{
            background-color: #e0e0e0;
        }}
        .action_object {{
            background-color: #e0e0e0;
        }}
        .url a {{
            color: #0645ad;
            text-decoration: none;
        }}
    </style>
</head>
<body>
{body}
</body>
</html>
"""


def format_evaluator_html(evaluator):
    if not evaluator:
        return "<p><strong>Evaluator:</strong> <em>Not found</em></p>"
    html = ["<div class='evaluator'><h4>Evaluator</h4>"]
    html.append(f"<b>Function:</b> {evaluator.get('func', '')}<br>")
    if "result" in evaluator:
        html.append(f"<b>Result:</b> <pre>{json.dumps(evaluator['result'], indent=2)}</pre>")
    if "expected" in evaluator:
        html.append(f"<b>Expected:</b> <pre>{json.dumps(evaluator['expected'], indent=2)}</pre>")
    html.append("</div>")
    return "\n".join(html)


def trajectory_to_html(traj_file: Path, render_text: bool = False, ann_types: list[str] = []) -> str:
    try:
        trajectory = VWA_Trajectory.from_json(traj_file)
    except Exception as e:
        print(f"Error reading {traj_file}: {e}")
        return ""

    # Gather header info.
    task_id = trajectory.task_id
    score = trajectory.cum_reward
    domain = trajectory.domain
    objective_text = ""

    objective_text = trajectory.objective["text"]
    objective_imgs = trajectory.objective["images"]
    objective_img_captions = trajectory.objective.get("img_captions", [])

    with open(trajectory.config_file, "r") as f:
        _config = json.load(f)
        _config_str = ""
        for k, v in _config.items():
            _config_str += f"{k}: {v}\n"
        _config_str = f"<pre>{_config_str}</pre>\n"
        task_id = _config["task_id"]

    # Compute number of verifier loops by checking each state's keys.
    num_verifier_loops = 0
    no_verifier_score = None
    retrieved_knowledge = ""
    for state in trajectory.states:
        if state.get("scores_per_round"):
            num_verifier_loops += 1
            if "score" in state["scores_per_round"][0] and no_verifier_score is None:
                no_verifier_score = state["scores_per_round"][0]["score"]
        if "retrieved_knowledge" in state and not retrieved_knowledge:
            retrieved_knowledge = state["retrieved_knowledge"].replace("\n", "<br>")

    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 style='background-color:#ffcccc; padding:4px 6px; border-radius:4px;'><strong>Score:</strong> {score}</p>"
    )
    if no_verifier_score is not None:
        html_parts.append(f"<p><strong>No Verifier Score:</strong> {no_verifier_score}</p>")

    html_parts.append(
        f"<p style='background-color:#ffcccc; padding:4px 6px; border-radius:4px;'><strong>Task:</strong> {domain}-{task_id}</p>"
    )
    html_parts.append(
        f"<p style='background-color:#ffcccc; padding:4px 6px; border-radius:4px;'><strong>Objective:</strong> {objective_text}</p>"
    )
    if objective_imgs:
        for img, caption in zip(objective_imgs, objective_img_captions):
            html_parts.append(f"<img src='{any_to_b64(img)}' style='width:20vw; height:auto;'/><br>")
            html_parts.append(f"<p><strong>Caption:</strong> {caption}</p>")

    html_parts.append(f"<p><strong>Verifier Loops:</strong> {num_verifier_loops}</p>")
    html_parts.append(f"<p><strong>Config:</strong> {_config_str.strip()}</p>")

    # Try to get retrieved_knowledge
    if retrieved_knowledge:
        html_parts.append(f"<p><strong>Retrieved Knowledge:</strong> {retrieved_knowledge}</p>")

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

        # URL of the page
        url = state.get("url", "")
        if url:
            html_parts.append(f"<h4 class='url'><a href='{url}'>URL: {url}</a></h4>")

        # Last executor atomic action
        executor_action = action["agents"]["executor"][-1]
        generated_text = executor_action["generated_text"]

        # Render text observation optionally
        if render_text:
            text_obs = state.get("text", "")
            if text_obs:
                text_obs_escaped = html.escape(text_obs)
                html_parts.append(f"<div class='state_obv'><pre>{text_obs_escaped}</pre></div>")

        # For the screenshot, check the 'screenshots' list if available.
        images = state.get("images", []) or []
        if len(images) > 0:
            trajectory.annotate_images(state, action, ann_types=ann_types)
            if state.get("images"):
                last_ss = state["images"][-1]
            else:
                last_ss = images[-1]
            html_parts.append(f"<img src='{any_to_b64(last_ss)}' style='width:30vw; height:auto;'/><br>")

        # Add the generated text.
        if generated_text:
            html_parts.append(f"<div class='raw_parsed_prediction'><pre>{generated_text}</pre></div>")

        # Check if a verifier loop exists in the 'verifier_loop' key.
        if state.get("scores_per_round", None):
            loop = state["scores_per_round"]
            html_parts.append("<div class='verifier_loop'>")
            html_parts.append("<h4>Verifier Loop</h4>")
            for item in loop:
                verifier_text = item.get("verifier_raw_response", "")
                generator_text = item.get("raw_prediction", "")
                if generator_text:
                    generator_text_html = generator_text.replace("\n", "<br>")
                    html_parts.append(
                        f"<div class='executor_utterance' style='background-color: #D8CEF6'><strong>Executor:</strong> {generator_text_html}</div>\n"
                    )
                if "score" in item:
                    html_parts.append(
                        f"<div class='score' style='background-color: #ffffe0'>Score: {item['score']}</div>"
                    )

                if verifier_text:
                    verifier_text_html = verifier_text.replace("\n", "<br>")
                    html_parts.append(
                        f"<div class='verifier_utterance' style='background-color: #F0E6FF'><strong>Verifier:</strong> {verifier_text_html}</div>\n"
                    )

            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, render_text: bool = False, ann_types: list[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.
    """
    try:
        traj_file = Path(traj_file)
        html_fragment = trajectory_to_html(traj_file, render_text=render_text, ann_types=ann_types)
        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)
            print(f"{output_file}")
        return str(output_file)
    except Exception as e:
        print(f"Error processing {traj_file}: {e}, {traceback.format_exc()}")
        return ""


def main(args):
    # Use the first argument as the base directory; default to "trace_osworld"

    base_dir = Path(args.base_dir)
    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",
        must_include_strs=args.must_include_strs,
        must_exclude_strs=args.must_exclude_strs,
    )

    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.
    if sys.gettrace():
        futures = [process_trajectory(traj_file, args.render_text, args.ann_types) for traj_file in traj_files]
    else:
        with concurrent.futures.ProcessPoolExecutor(max_workers=16) as executor:
            # Submit all trajectory files for processing.
            futures = {
                executor.submit(process_trajectory, traj_file, (args.render_text), ann_types=args.ann_types): traj_file
                for traj_file in traj_files
            }
            for future in concurrent.futures.as_completed(futures):
                result = future.result()


parser = argparse.ArgumentParser()
parser.add_argument("base_dir", nargs="?", default="", help="Base directory to search for trajectories")
parser.add_argument("--render_text", action="store_true", help="Render text observations")
parser.add_argument("--ann_types", type=str, default="", help="Annotation types to render")
parser.add_argument("--must_include_strs", type=list[str], default=[], help="Must include strings")
parser.add_argument("--must_exclude_strs", type=list[str], default=[], help="Must exclude strings")
args = parser.parse_args()

# args.base_dir = "experiments/gemini-2.5-flash-001/debug_reflexion-4/"
# args.ann_types = ["bbox", "coordinates"]

main(args)
