# report/builder.py
import os
from typing import Any

import jinja2

from mujoco.sysid.report.sections.base import ReportSection

# Path to the templates directory
TEMPLATE_DIR = os.path.join(os.path.dirname(__file__), "templates")


class ReportBuilder:
  def __init__(self, title: str, global_context: dict[str, Any] | None = None):
    self._title = title
    self._sections: list[ReportSection] = []
    self._global_context = global_context or {}

    # Setup Jinja to load from files
    self._env = jinja2.Environment(
      loader=jinja2.FileSystemLoader(TEMPLATE_DIR),
      autoescape=jinja2.select_autoescape(["html", "xml"]),
    )

  def add_section(self, section: ReportSection):
    self._sections.append(section)

  def build(self) -> str:
    # Load the main layout
    layout_template = self._env.get_template("layout.html")

    # Render each section individually
    rendered_sections = []
    all_header_includes = set()

    for section in self._sections:
      # Get the specific template for this section
      sec_template = self._env.get_template(section.template_filename)

      # Check if this is a GroupSection (has 'sections' attribute)
      extra_context = {}
      child_sections: list[ReportSection] = getattr(section, "sections", [])
      if child_sections:
        child_sections_content = []
        for child in child_sections:
          child_template = self._env.get_template(child.template_filename)
          child_html = child_template.render(child.get_context())
          all_header_includes.update(child.header_includes())

          # Fallback anchor for child
          child_anchor = child.anchor
          if not child_anchor:
            child_anchor = (
              child.title.lower()
              .replace(" ", "-")
              .replace("[", "")
              .replace("]", "")
              .replace("(", "")
              .replace(")", "")
            )

          child_sections_content.append(
            {"title": child.title, "content": child_html, "anchor": child_anchor}
          )
        extra_context["child_sections"] = child_sections_content

      # Render the section HTML (main wrapper)
      html_content = sec_template.render(section.get_context() | extra_context)
      # Collect header requirements (scripts/css)
      all_header_includes.update(section.header_includes())

      # Fallback anchor generation
      anchor = section.anchor
      if not anchor:
        anchor = (
          section.title.lower()
          .replace(" ", "-")
          .replace("[", "")
          .replace("]", "")
          .replace("(", "")
          .replace(")", "")
        )

      rendered_sections.append(
        {
          "title": section.title,
          "anchor": anchor,
          "collapsible": section.collapsible,
          "is_open": section.is_open,
          "content": html_content,
        }
      )

    # Render final report
    return layout_template.render(
      report_title=self._title,
      sections=rendered_sections,
      header_includes=all_header_includes,
      **self._global_context,
    )

  def save(self, path: str):
    with open(path, "w", encoding="utf-8") as f:
      f.write(self.build())