"""The visualizer load will render the intent, past action, next action and the correponding HTML"""
import json
import os
import uuid
from dataclasses import dataclass, field

from bs4 import BeautifulSoup


def add_html_wrapper(html: str) -> str:
    """Add head and tail to the html"""
    html = f"""<!DOCTYPE html>
<html>
<body>
{html}
</body>
</html>"""
    return html


class Step:
    def __init__(
        self,
        observation: str,
        action: str,
        history: list[str],
        id: str = None,
        save_path: str = "",
    ):
        self.observation = observation
        self.action = action
        self.history = history
        self.id = uuid.uuid4().hex if not id else id
        self.save_path = save_path if save_path else os.getcwd()

        # highlight the observation with

    def to_html(self) -> str:
        """Visualize the step"""
        id = self.id
        history_str = "\n".join(self.history)
        div = f"""
<iframe src="{self.save_path}/{id}.html" id="{id}" width=100% height=500></iframe>
<div class="history" id="{id}"><pre>{history_str}</pre></div>
<div class="action" id="{id}"><pre>{self.action}</pre></div>"""
        return div

    @property
    def html_complexity(self) -> int:
        """Return the number of nodes in the HTML"""
        soup = BeautifulSoup(self.observation, "html.parser")
        return len(soup.find_all())


class Trajectory:
    def __init__(self, intent: str):
        self.intent = intent
        self.id = uuid.uuid4().hex
        self.step_list: list[Step] = []
        self.step_ids: list[str] = []

    def to_html(self, add_head: bool = True) -> str:
        """Visualize the example"""
        html = f'<div class="intent"><h2>{self.intent}</h2></div>'
        for step in self.step_list:
            html += step.to_html()
        html = f'<div id="{self.id}" class="trajectory">{html}</div>'
        # add head and tail
        if add_head:
            html = add_html_wrapper(html)

        return html

    def add_step(self, step: Step) -> None:
        self.step_list.append(step)

    def __len__(self):
        return len(self.step_list)
