"""
Logtree: Scope-based logging library for creating nested HTML reports.

This module provides a context-based API for generating structured HTML logs
that reflect the call tree of your code. Ideal for logging RL rollouts,
model evaluations, and other hierarchical computations.

Example usage:
    import logtree

    async def train_iteration():
        with logtree.init_trace("Training Iteration 1", path="output.html"):
            with logtree.scope_header("Sampling"):
                logtree.log_text("Generated 100 samples")
                await sample_data()

            with logtree.scope_header("Training"):
                logtree.log_text("Loss: 0.42")
                await train_model()
"""

import functools
import html as html_module
import inspect
import json
import os
import traceback
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Iterator, Mapping, Protocol, Sequence, TypeVar, cast, overload

# Context variables for task-local state
_current_trace: ContextVar["Trace | None"] = ContextVar("lt_current_trace", default=None)
_container_stack: ContextVar["tuple[Node, ...]"] = ContextVar("lt_container_stack", default=())
_header_depth: ContextVar["tuple[int, ...]"] = ContextVar("lt_header_depth", default=())
_logging_disabled: ContextVar[bool] = ContextVar("lt_logging_disabled", default=False)


class Formatter(Protocol):
    """Protocol for objects that can format themselves as HTML with CSS.

    Optionally implement ``to_data`` to attach structured data to the JSON
    export (avoids consumers having to parse raw HTML).
    """

    def to_html(self) -> str:
        """Generate HTML representation of this object."""
        ...

    def get_css(self) -> str:
        """Get CSS needed to style this object's HTML."""
        ...

    def to_data(self) -> dict[str, Any] | None:
        """Return structured data for JSON export, or None."""
        return None


@dataclass
class Node:
    """Represents an HTML element in the tree."""

    tag: str
    attrs: dict[str, str] = field(default_factory=dict)
    children: list["Node | str"] = field(default_factory=list)
    # Optional structured data attached by formatters.
    # Included in JSON export so consumers can extract typed content
    # (e.g., conversation messages) without parsing raw HTML.
    data: dict[str, Any] | None = field(default=None, repr=False)

    def to_html(self, indent: int = 0) -> str:
        """Convert node to HTML string."""
        ind = "  " * indent
        attrs_str = "".join(
            f' {k}="{html_module.escape(v, quote=True)}"' for k, v in self.attrs.items()
        )

        if not self.children:
            return f"{ind}<{self.tag}{attrs_str}></{self.tag}>\n"

        # Keep simple text-only nodes on one line to avoid extra rendered
        # whitespace when CSS uses `white-space: pre-wrap` (e.g., lt-p).
        if all(isinstance(child, str) for child in self.children):
            text = "".join(child for child in self.children if isinstance(child, str))
            return f"{ind}<{self.tag}{attrs_str}>{text}</{self.tag}>\n"

        lines = [f"{ind}<{self.tag}{attrs_str}>\n"]
        for child in self.children:
            if isinstance(child, str):
                lines.append(child)
            else:
                lines.append(child.to_html(indent + 1))
        lines.append(f"{ind}</{self.tag}>\n")
        return "".join(lines)

    def to_dict(self) -> dict[str, Any]:
        """Convert node to a JSON-serializable dictionary.

        When ``data`` is present, raw HTML string children are omitted from the
        JSON — consumers should use ``data`` instead.
        """
        if self.data is not None:
            children = [child.to_dict() for child in self.children if isinstance(child, Node)]
        else:
            children = [
                child if isinstance(child, str) else child.to_dict() for child in self.children
            ]
        d: dict[str, Any] = {
            "tag": self.tag,
            "attrs": dict(self.attrs),
            "children": children,
        }
        if self.data is not None:
            d["data"] = self.data
        return d


@dataclass
class Theme:
    """Theme configuration for HTML output."""

    css_text: str | None = None  # Custom CSS; if None, use built-in
    css_urls: list[str] = field(default_factory=list)
    css_vars: dict[str, str] = field(default_factory=dict)  # CSS custom properties


class Trace:
    """Root trace object representing an HTML document."""

    def __init__(self, title: str, path: str | os.PathLike | None, write_on_error: bool):
        self.title = title
        self.path = Path(path) if path is not None else None
        self.write_on_error = write_on_error
        self.started_at = datetime.now()
        self.root = Node("body", {"class": "lt-root"})
        self._formatter_css: set[str] = set()  # Deduplicated CSS from formatters

    def _register_formatter_css(self, css: str) -> None:
        """Register CSS from a formatter (deduplicated per trace)."""
        if css:
            self._formatter_css.add(css)

    def body_html(self, wrap_body: bool = True) -> str:
        """Get the body HTML."""
        inner = self.root.to_html(indent=0)
        if wrap_body:
            return inner
        else:
            # Return just the inner content
            return "\n".join(
                line for line in inner.split("\n") if "<body" not in line and "</body>" not in line
            )

    def get_html(self) -> str:
        """Alias for body_html()."""
        return self.body_html(wrap_body=True)

    def head_html(
        self, theme: Theme | None = None, title: str | None = None, extra_head: str | None = None
    ) -> str:
        """Generate the <head> section of the HTML document."""
        if theme is None:
            theme = Theme()

        parts = []
        parts.append(f"<title>{html_module.escape(title or self.title)}</title>")
        parts.append('<meta charset="UTF-8">')
        parts.append('<meta name="viewport" content="width=device-width, initial-scale=1.0">')

        # External CSS
        for url in theme.css_urls:
            parts.append(f'<link rel="stylesheet" href="{html_module.escape(url, quote=True)}">')

        # Inline CSS
        css = theme.css_text if theme.css_text is not None else _DEFAULT_CSS
        if css or self._formatter_css:
            parts.append("<style>")
            if css:
                parts.append(css)
            # CSS custom properties
            if theme.css_vars:
                parts.append(":root {")
                for key, value in theme.css_vars.items():
                    parts.append(f"  {key}: {value};")
                parts.append("}")
            # Formatter CSS
            if self._formatter_css:
                parts.append("\n/* Formatter CSS */")
                for formatter_css in self._formatter_css:
                    parts.append(formatter_css)
            parts.append("</style>")

        if extra_head:
            parts.append(extra_head)

        return "\n".join(parts)

    def to_dict(self) -> dict[str, Any]:
        """Convert the trace to a JSON-serializable dictionary."""
        return {
            "title": self.title,
            "started_at": self.started_at.isoformat(),
            "path": str(self.path) if self.path is not None else None,
            "root": self.root.to_dict(),
        }


# Default CSS styling
_DEFAULT_CSS = """
body {
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
    line-height: 1.45;
    max-width: 1200px;
    margin: 0 auto;
    padding: 14px;
    background: var(--lt-bg, #f5f5f5);
    color: var(--lt-text, #333);
}

.lt-root {
    background: var(--lt-card, white);
    padding: 1.2rem 1.4rem;
    border-radius: 8px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

.lt-title {
    margin: 0 0 0.5rem 0;
    color: var(--lt-accent, #2563eb);
    border-bottom: 2px solid var(--lt-border, #e5e7eb);
    padding-bottom: 0.5rem;
}

.lt-subtitle {
    color: var(--lt-sub, #666);
    font-size: 0.875rem;
    margin-bottom: 1.2rem;
}

.lt-section {
    margin: 0.95rem 0;
    padding-left: 0.75rem;
    border-left: 2px solid var(--lt-border, #e5e7eb);
}

.lt-section-body {
    margin-top: 0.12rem;
}

.lt-section h2, .lt-section h3, .lt-section h4, .lt-section h5, .lt-section h6 {
    margin: 0.2rem 0;
    line-height: 1.3;
    color: var(--lt-accent, #2563eb);
}

.lt-h2 { font-size: 1.15rem; }
.lt-h3 { font-size: 1.05rem; }
.lt-h4 { font-size: 0.98rem; }
.lt-h5 { font-size: 0.95rem; }
.lt-h6 { font-size: 0.92rem; }

.lt-p {
    margin: 0.2rem 0;
    white-space: pre-wrap;
}

.lt-details {
    margin: 0.35rem 0;
    border: 1px solid var(--lt-border, #e5e7eb);
    border-radius: 4px;
    padding: 0.35rem 0.45rem;
}

.lt-details summary {
    cursor: pointer;
    font-weight: 600;
    user-select: none;
}

.lt-details-body {
    margin-top: 0.25rem;
    padding: 0.35rem 0.45rem;
    background: var(--lt-bg, #f5f5f5);
    border-radius: 4px;
    overflow-x: auto;
}

.lt-details-body pre {
    margin: 0;
    font-family: var(--lt-mono, "Courier New", monospace);
    font-size: 0.875rem;
    white-space: pre-wrap;
}

.lt-table {
    border-collapse: separate;
    border-spacing: 0;
    width: 100%;
    margin: 0.6rem 0;
    font-size: 0.875rem;
    border: 1px solid var(--lt-border, #d5dbe3);
    border-radius: 6px;
    background: #fff;
    overflow: hidden;
}

.lt-table th {
    background: var(--lt-table-head-bg, #eef2f7);
    color: var(--lt-text, #1f2937);
    padding: 0.5rem;
    text-align: left;
    font-weight: 600;
    border-bottom: 1px solid var(--lt-border, #d5dbe3);
}

.lt-table td {
    padding: 0.5rem;
    border-bottom: 1px solid var(--lt-border, #e5e7eb);
}

.lt-table tr:nth-child(even) {
    background: #f8fafc;
}

.lt-table-caption {
    font-weight: 600;
    margin-bottom: 0.5rem;
    color: var(--lt-text, #333);
}

.lt-exc {
    background: #fee;
    border: 2px solid #c00;
    border-radius: 4px;
    padding: 0.75rem;
    margin: 0.5rem 0;
}

.lt-exc summary {
    color: #c00;
    font-weight: 700;
    cursor: pointer;
}

.lt-exc pre {
    margin-top: 0.5rem;
    font-family: var(--lt-mono, "Courier New", monospace);
    font-size: 0.875rem;
    overflow-x: auto;
}

.answer, .reward {
    font-weight: 600;
    padding: 0.25rem 0.5rem;
    border-radius: 4px;
    display: inline-block;
    margin: 0.25rem 0;
}

.answer {
    background: #dbeafe;
    color: #1e40af;
}

.reward {
    background: #dcfce7;
    color: #166534;
}
"""


# Helper functions


def _normalize_attrs(**attrs: Any) -> dict[str, str]:
    """Normalize attribute names (class_ -> class, data__foo -> data-foo)."""
    result = {}
    for key, value in attrs.items():
        if key == "class_":
            key = "class"
        elif key.startswith("data__"):
            key = key.replace("__", "-", 1)
        result[key] = str(value)
    return result


def _append(node: Node) -> None:
    """Append a node to the current container."""
    stack = _container_stack.get()
    if not stack:
        raise RuntimeError("No active container to append to")
    stack[-1].children.append(node)


def _next_header_level() -> int:
    """Get the next header level based on current depth."""
    depth = _header_depth.get()
    current = depth[-1] if depth else 1
    return min(6, current + 1)


def _is_logging_enabled() -> bool:
    """Check if logging is currently enabled."""
    return _current_trace.get() is not None and not _logging_disabled.get()


@contextmanager
def _in_container(node: Node) -> Iterator[None]:
    """Context manager to push/pop a container."""
    token = _container_stack.set(_container_stack.get() + (node,))
    try:
        yield
    finally:
        _container_stack.reset(token)


def _exception_block(exc: BaseException) -> Node:
    """Create a details block for an exception."""
    tb_str = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
    details_node = Node("details", {"class": "lt-exc", "open": "open"})
    details_node.children.append(Node("summary", {}, [f"Exception: {type(exc).__name__}: {exc}"]))
    pre_node = Node("pre", {})
    pre_node.children.append(html_module.escape(tb_str))
    details_node.children.append(pre_node)
    return details_node


def _write_trace(trace: Trace, theme: Theme | None = None) -> None:
    """Write the trace to disk."""
    if trace.path is None:
        return

    trace.path.parent.mkdir(parents=True, exist_ok=True)

    with open(trace.path, "w") as f:
        f.write("<!doctype html>\n")
        f.write('<html lang="en">\n')
        f.write("<head>\n")
        f.write(trace.head_html(theme=theme))
        f.write("</head>\n")
        f.write(trace.body_html(wrap_body=True))
        f.write("</html>\n")


# Public API: Trace lifecycle


@contextmanager
def init_trace(
    title: str, path: str | os.PathLike | None = None, *, write_on_error: bool = True
) -> Iterator[Trace]:
    """
    Initialize a new trace context.

    Args:
        title: Title for the HTML document (becomes <h1>)
        path: Path to write HTML file (None = don't write automatically)
        write_on_error: If True, write partial HTML even if exception occurs

    Example:
        with logtree.init_trace("My Report", path="output.html"):
            logtree.log_text("Hello world")
    """
    trace = Trace(title, path, write_on_error=write_on_error)

    tok_t = _current_trace.set(trace)
    tok_s = _container_stack.set((trace.root,))
    tok_h = _header_depth.set((1,))

    # Emit title and subtitle
    _append(Node("h1", {"class": "lt-title"}, [html_module.escape(title)]))
    _append(
        Node(
            "div",
            {"class": "lt-subtitle"},
            [f"Generated {trace.started_at.isoformat(timespec='seconds')}"],
        )
    )

    try:
        yield trace
    except BaseException as e:
        _append(_exception_block(e))
        if write_on_error and trace.path is not None:
            _write_trace(trace)
        raise
    else:
        if trace.path is not None:
            _write_trace(trace)
    finally:
        # Always reset context even on exception
        _header_depth.reset(tok_h)
        _container_stack.reset(tok_s)
        _current_trace.reset(tok_t)


@contextmanager
def scope_header(title: str, **attrs: Any) -> Iterator[None]:
    """
    Open a section with an auto-leveled header.

    Args:
        title: Text for the header
        **attrs: HTML attributes (use class_="foo" for class, data__x="y" for data-x)

    Example:
        with logtree.scope_header("Results", class_="important"):
            logtree.log_text("Success rate: 95%")
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        yield
        return

    section = Node("section", {"class": "lt-section", **_normalize_attrs(**attrs)})
    _append(section)

    tok_h = None
    try:
        with _in_container(section):
            h = _next_header_level()
            _append(Node(f"h{h}", {"class": f"lt-h{h}"}, [html_module.escape(title)]))

            # Push header level for nested scopes
            tok_h = _header_depth.set(_header_depth.get() + (h,))
            body = Node("div", {"class": "lt-section-body"})
            _append(body)
            with _in_container(body):
                yield
    finally:
        # Always reset context even on exception
        if tok_h is not None:
            _header_depth.reset(tok_h)


F = TypeVar("F", bound=Callable[..., Any])


# Overloads the parameterized usage
@overload
def scope_header_decorator(title: str) -> Callable[[F], F]: ...  # String title


# Overloads the bare usage
@overload
def scope_header_decorator(title: F) -> F: ...  # Bare: @scope_header_decorator


def scope_header_decorator(
    title: str | F,
) -> F | Callable[[F], F]:
    """
    Decorator to wrap function in a scope_header.

    Args:
        title: String or function returning string

    Examples:
        @logtree.scope_header_decorator
        async def process_batch():
            ...

        @logtree.scope_header_decorator("Handling item")
        def handle_item():
            ...
    """
    title_str = title if isinstance(title, str) else title.__name__

    def _wrap(fn: F) -> F:
        if inspect.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def aw(*args: Any, **kwargs: Any) -> Any:
                # Graceful degradation: if logging is disabled, just run the function
                if not _is_logging_enabled():
                    return await fn(*args, **kwargs)

                with scope_header(title_str):
                    return await fn(*args, **kwargs)

            return aw  # type: ignore
        else:

            @functools.wraps(fn)
            def w(*args: Any, **kwargs: Any) -> Any:
                # Graceful degradation: if logging is disabled, just run the function
                if not _is_logging_enabled():
                    return fn(*args, **kwargs)

                with scope_header(title_str):
                    return fn(*args, **kwargs)

            return w  # type: ignore

    if isinstance(title, str):
        return _wrap
    else:
        fn = title
        return _wrap(fn)


@contextmanager
def scope_div(**attrs: Any) -> Iterator[None]:
    """
    Open a <div> scope (does not change header level).

    Args:
        **attrs: HTML attributes

    Example:
        with logtree.scope_div(class_="grading"):
            logtree.log_text("Grade: A")
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        yield
        return

    div = Node("div", _normalize_attrs(**attrs))
    _append(div)
    with _in_container(div):
        yield


@contextmanager
def scope_disable() -> Iterator[None]:
    """
    Disable all logging within this scope.

    Example:
        with scope_header("Group A") if should_log else scope_disable():
            logtree.log_text("Data here")
    """
    token = _logging_disabled.set(True)
    try:
        yield
    finally:
        _logging_disabled.reset(token)


@contextmanager
def optional_enable_logging(enable: bool) -> Iterator[None]:
    """Context manager to optionally enable logging."""
    if enable:
        yield
    else:
        with scope_disable():
            yield


@contextmanager
def scope_details(summary: str) -> Iterator[None]:
    """
    Open a collapsible <details> scope.

    Args:
        summary: Summary text shown when collapsed

    Example:
        with logtree.scope_details("Click to expand"):
            logtree.log_text("Hidden content")
            logtree.log_text("More hidden content")
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        yield
        return

    details_node = Node("details", {"class": "lt-details"})
    details_node.children.append(Node("summary", {}, [html_module.escape(summary)]))

    body_div = Node("div", {"class": "lt-details-body"})
    details_node.children.append(body_div)

    _append(details_node)
    with _in_container(body_div):
        yield


# Public API: Content


def log_text(text: str, *, div_class: str | None = None) -> None:
    """
    Log a text paragraph.

    Args:
        text: Text to log (will be HTML-escaped)
        div_class: If set, wrap in <div class="{div_class}"> instead of <p>

    Example:
        logtree.log_text("Processing complete")
        logtree.log_text("Score: 0.95", div_class="score")
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        return

    escaped = html_module.escape(text)
    if div_class:
        _append(Node("div", {"class": div_class}, [escaped]))
    else:
        _append(Node("p", {"class": "lt-p"}, [escaped]))


def log_html(html: str, *, div_class: str | None = None) -> None:
    """
    Log raw HTML (not escaped).

    Args:
        html: HTML string to insert verbatim
        div_class: If set, wrap in <div class="{div_class}">

    Example:
        logtree.log_html("<strong>Important</strong>")
        logtree.log_html(conversation_html, div_class="conversation")
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        return

    if div_class:
        div = Node("div", {"class": div_class})
        div.children.append(html)
        _append(div)
    else:
        # Create a container node that holds raw HTML
        container = Node("div", {})
        container.children.append(html)
        _append(container)


def log_formatter(formatter: Formatter) -> None:
    """
    Log an object that knows how to format itself as HTML.

    The formatter's CSS will be automatically included in the trace output.
    CSS is deduplicated per trace, so logging multiple objects of the same
    type only includes the CSS once.

    Args:
        formatter: Object implementing the Formatter protocol (to_html() and get_css())

    Example:
        from tinker_cookbook.utils.logtree_formatters import ConversationFormatter

        logtree.log_formatter(ConversationFormatter(messages=[...]))
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        return

    # Register CSS from the formatter (deduplicated)
    trace = _current_trace.get()
    assert trace is not None  # _is_logging_enabled() ensures this
    css = formatter.get_css()
    trace._register_formatter_css(css)

    # Log the HTML, with optional structured data for JSON export
    html = formatter.to_html()
    container = Node("div", {})
    container.children.append(html)
    to_data = cast(
        "Callable[[], dict[str, Any] | None] | None", getattr(formatter, "to_data", None)
    )
    data = to_data() if callable(to_data) else None
    if data is not None:
        container.data = data
    _append(container)


def details(text: str, *, summary: str = "Details", pre: bool = True) -> None:
    """
    Log collapsible details block.

    Args:
        text: Content text
        summary: Summary text (what you see when collapsed)
        pre: If True, use <pre> (preserves whitespace), else <div>

    Example:
        logtree.details(long_chain_of_thought, summary="CoT Reasoning", pre=True)
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        return

    details_node = Node("details", {"class": "lt-details"})
    details_node.children.append(Node("summary", {}, [html_module.escape(summary)]))

    body_node = Node("pre" if pre else "div", {"class": "lt-details-body"})
    body_node.children.append(html_module.escape(text))
    details_node.children.append(body_node)

    _append(details_node)


def header(text: str, *, level: int | None = None) -> None:
    """
    Log an inline header.

    Args:
        text: Header text
        level: Header level (1-6), or None to auto-compute from scope depth

    Example:
        logtree.header("Results")
        logtree.header("Subsection", level=4)
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        return

    h = level if level is not None else _next_header_level()
    h = max(1, min(6, h))
    _append(Node(f"h{h}", {"class": f"lt-h{h}"}, [html_module.escape(text)]))


# Public API: Tables


def table(obj: Any, *, caption: str | None = None) -> None:
    """
    Log a table from various data types.

    Supports:
    - pandas.DataFrame
    - list[dict]
    - list[list]

    Does NOT support raw dict (use table_from_dict or table_from_dict_of_lists).

    Args:
        obj: Data object
        caption: Optional caption text

    Example:
        logtree.table(df, caption="Results")
        logtree.table([{"name": "Alice", "score": 95}, {"name": "Bob", "score": 87}])
    """
    # Graceful degradation: if logging is disabled, do nothing
    if not _is_logging_enabled():
        return

    if isinstance(obj, dict):
        raise TypeError(
            "table() does not accept dict directly. Use table_from_dict() or table_from_dict_of_lists()."
        )

    if isinstance(obj, list) and not obj:
        if caption:
            _append(Node("div", {"class": "lt-table-caption"}, [html_module.escape(caption)]))
        return

    # Try DataFrame — convert to records so the JSON tree gets structured
    # Nodes (thead/tbody/tr/td) instead of a raw HTML string.
    try:
        import pandas as pd

        if isinstance(obj, pd.DataFrame):
            _table_from_list_of_dicts(obj.to_dict("records"), caption=caption)
            return
    except ImportError:
        pass

    # list[dict]
    if isinstance(obj, list) and obj and isinstance(obj[0], dict):
        _table_from_list_of_dicts(obj, caption=caption)
        return

    # list[list]
    if isinstance(obj, list) and obj and isinstance(obj[0], (list, tuple)):
        _table_from_list_of_lists(obj, caption=caption)
        return

    raise TypeError(f"table() does not support type {type(obj)}")


def table_from_dict(
    data: Mapping[Any, Any],
    *,
    caption: str | None = None,
    key_header: str = "key",
    value_header: str = "value",
    sort_by: str | None = None,
) -> None:
    """
    Log a two-column key-value table from a dict.

    Args:
        data: Dictionary to display
        caption: Optional caption
        key_header: Column header for keys
        value_header: Column header for values
        sort_by: "key", "value", or None

    Example:
        logtree.table_from_dict({"lr": 0.001, "batch_size": 32}, caption="Hyperparams")
    """
    if not _is_logging_enabled():
        return

    items = list(data.items())
    if sort_by == "key":
        items.sort(key=lambda x: x[0])
    elif sort_by == "value":
        items.sort(key=lambda x: x[1])

    rows = [[key_header, value_header]] + [[str(k), str(v)] for k, v in items]
    _table_from_list_of_lists(rows, caption=caption, has_header=True)


def table_from_dict_of_lists(
    columns: Mapping[str, Sequence[Any]],
    *,
    caption: str | None = None,
    order: Sequence[str] | None = None,
) -> None:
    """
    Log a columnar table from dict of lists.

    Args:
        columns: Dict where keys are column names, values are column data
        caption: Optional caption
        order: Column order (if None, use insertion order)

    Example:
        logtree.table_from_dict_of_lists({
            "name": ["Alice", "Bob"],
            "score": [95, 87]
        })
    """
    if not _is_logging_enabled():
        return

    if not columns:
        return

    # Validate equal lengths
    lengths = [len(v) for v in columns.values()]
    if len(set(lengths)) > 1:
        raise ValueError("All columns must have equal length")

    col_names = list(order) if order else list(columns.keys())
    rows = [col_names]
    for i in range(lengths[0]):
        rows.append([str(columns[name][i]) for name in col_names])

    _table_from_list_of_lists(rows, caption=caption, has_header=True)


def _table_from_list_of_dicts(data: list[dict], *, caption: str | None = None) -> None:
    """Helper: create table from list of dicts."""
    if not data:
        return

    keys = list(data[0].keys())
    rows = [keys]
    for item in data:
        rows.append([str(item.get(k, "")) for k in keys])

    _table_from_list_of_lists(rows, caption=caption, has_header=True)


def _table_from_list_of_lists(
    rows: list[list[Any]], *, caption: str | None = None, has_header: bool = False
) -> None:
    """Helper: create HTML table from list of lists."""
    if not rows:
        return

    if caption:
        _append(Node("div", {"class": "lt-table-caption"}, [html_module.escape(caption)]))

    table_node = Node("table", {"class": "lt-table"})

    if has_header:
        thead = Node("thead")
        tr = Node("tr")
        for cell in rows[0]:
            tr.children.append(Node("th", {}, [html_module.escape(str(cell))]))
        thead.children.append(tr)
        table_node.children.append(thead)
        rows = rows[1:]

    tbody = Node("tbody")
    for row in rows:
        tr = Node("tr")
        for cell in row:
            tr.children.append(Node("td", {}, [html_module.escape(str(cell))]))
        tbody.children.append(tr)
    table_node.children.append(tbody)

    _append(table_node)


# Public API: Export & theming


def write_html_with_default_style(
    body_html: str,
    path: str | os.PathLike,
    *,
    title: str = "Trace",
    theme: Theme | None = None,
    lang: str = "en",
    extra_head: str | None = None,
) -> None:
    """
    Write a complete HTML document with default styling.

    Args:
        body_html: Body HTML (with or without <body> tags)
        path: Output file path
        title: Document title
        theme: Optional theme
        lang: HTML lang attribute
        extra_head: Extra content for <head>
    """
    if theme is None:
        theme = Theme()

    # Create a temporary trace just for head generation
    trace = Trace(title, None, False)

    path_obj = Path(path)
    path_obj.parent.mkdir(parents=True, exist_ok=True)

    with open(path_obj, "w") as f:
        f.write(f'<!doctype html>\n<html lang="{html_module.escape(lang)}">\n')
        f.write("<head>\n")
        f.write(trace.head_html(theme=theme, title=title, extra_head=extra_head))
        f.write("</head>\n")
        # Ensure body tags are present
        if "<body" not in body_html:
            f.write("<body>\n")
            f.write(body_html)
            f.write("</body>\n")
        else:
            f.write(body_html)
        f.write("</html>\n")


def write_trace_json(trace: Trace, path: str | os.PathLike) -> None:
    """
    Write the trace structure to JSON.

    Args:
        trace: Trace object to serialize.
        path: Output JSON file path.
    """
    path_obj = Path(path)
    path_obj.parent.mkdir(parents=True, exist_ok=True)
    with open(path_obj, "w") as f:
        json.dump(trace.to_dict(), f, indent=2)


def jinja_context(trace: Trace, **extra: Any) -> dict[str, Any]:
    """
    Create a context dict for Jinja2 templates.

    Args:
        trace: Trace object
        **extra: Additional context variables

    Returns:
        Dict with standard keys: title, generated_at, started_at, body_html, head_html
    """
    return {
        "title": trace.title,
        "generated_at": datetime.now().isoformat(),
        "started_at": trace.started_at.isoformat(),
        "body_html": trace.body_html(),
        "head_html": trace.head_html(),
        **extra,
    }


def render_with_jinja(
    env: Any,
    template_name: str,
    *,
    context: dict[str, Any],
    write_to: str | os.PathLike | None = None,
) -> str:
    """
    Render using Jinja2 (requires jinja2 to be installed).

    Args:
        env: jinja2.Environment instance
        template_name: Template file name
        context: Template context
        write_to: Optional path to write output

    Returns:
        Rendered HTML string
    """
    template = env.get_template(template_name)
    html = template.render(**context)

    if write_to is not None:
        path = Path(write_to)
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            f.write(html)

    return html


def flush_trace() -> bool:
    """
    Flush the current trace to the saved path even if the trace has not been exited.
    This is useful for long-running programs where we wanna inspect some logs early.

    Returns:
        True if the trace was flushed, False otherwise.
    """
    trace = _current_trace.get()
    if trace is not None and trace.path is not None:
        _write_trace(trace)
        return True
    return False
