rose-ash/shared/sx/jinja_bridge.py

"""
Jinja ↔ s-expression bridge.

Provides two-way integration so s-expression components and Jinja templates
can coexist during incremental migration:

**Jinja → s-expression** (use s-expression components inside Jinja templates)::

    {{ sx('(~link-card :slug "apple" :title "Apple")') | safe }}

**S-expression → Jinja** (embed Jinja output inside s-expressions)::

    (raw! (jinja "fragments/link_card.html" :slug "apple" :title "Apple"))

Setup::

    from shared.sx.jinja_bridge import setup_sx_bridge
    setup_sx_bridge(app)  # call after setup_jinja(app)
"""

from __future__ import annotations

import glob
import hashlib
import os
from typing import Any

from .types import NIL, Component, Island, Keyword, Macro, Symbol
from .parser import parse
import os as _os
if _os.environ.get("SX_USE_REF") == "1":
    from .ref.sx_ref import render as html_render, render_html_component as _render_component
else:
    from .html import render as html_render, _render_component


# ---------------------------------------------------------------------------
# Shared component environment
# ---------------------------------------------------------------------------

# Global component registry — populated at app startup by loading component
# definition files or calling register_components().
_COMPONENT_ENV: dict[str, Any] = {}

# SHA-256 hash (12 hex chars) of all component definitions — used for
# client-side localStorage caching.
_COMPONENT_HASH: str = ""


def get_component_env() -> dict[str, Any]:
    """Return the shared component environment."""
    return _COMPONENT_ENV


def get_component_hash() -> str:
    """Return the current component definitions hash."""
    return _COMPONENT_HASH


def _compute_component_hash() -> None:
    """Recompute _COMPONENT_HASH from all registered Component and Macro definitions."""
    global _COMPONENT_HASH
    from .parser import serialize
    parts = []
    for key in sorted(_COMPONENT_ENV):
        val = _COMPONENT_ENV[key]
        if isinstance(val, Island):
            param_strs = ["&key"] + list(val.params)
            if val.has_children:
                param_strs.extend(["&rest", "children"])
            params_sx = "(" + " ".join(param_strs) + ")"
            body_sx = serialize(val.body)
            parts.append(f"(defisland ~{val.name} {params_sx} {body_sx})")
        elif isinstance(val, Component):
            param_strs = ["&key"] + list(val.params)
            if val.has_children:
                param_strs.extend(["&rest", "children"])
            params_sx = "(" + " ".join(param_strs) + ")"
            body_sx = serialize(val.body)
            parts.append(f"(defcomp ~{val.name} {params_sx} {body_sx})")
        elif isinstance(val, Macro):
            param_strs = list(val.params)
            if val.rest_param:
                param_strs.extend(["&rest", val.rest_param])
            params_sx = "(" + " ".join(param_strs) + ")"
            body_sx = serialize(val.body)
            parts.append(f"(defmacro {val.name} {params_sx} {body_sx})")
    if parts:
        digest = hashlib.sha256("\n".join(parts).encode()).hexdigest()[:12]
        _COMPONENT_HASH = digest
    else:
        _COMPONENT_HASH = ""


def load_sx_dir(directory: str) -> None:
    """Load all .sx files from a directory and register components.

    Skips boundary.sx — those are parsed separately by the boundary validator.
    """
    for filepath in sorted(
        glob.glob(os.path.join(directory, "**", "*.sx"), recursive=True)
    ):
        if os.path.basename(filepath) == "boundary.sx":
            continue
        with open(filepath, encoding="utf-8") as f:
            register_components(f.read())


# ---------------------------------------------------------------------------
# Dev-mode auto-reload of sx templates
# ---------------------------------------------------------------------------

_watched_dirs: list[str] = []
_file_mtimes: dict[str, float] = {}
_reload_callbacks: list[Any] = []


def register_reload_callback(fn: Any) -> None:
    """Register a function to call after hot-reload clears and reloads components."""
    _reload_callbacks.append(fn)


def watch_sx_dir(directory: str) -> None:
    """Register a directory for dev-mode file watching."""
    _watched_dirs.append(directory)
    # Seed mtimes
    for fp in sorted(
        glob.glob(os.path.join(directory, "**", "*.sx"), recursive=True)
    ):
        _file_mtimes[fp] = os.path.getmtime(fp)


def reload_if_changed() -> None:
    """Re-read sx files if any have changed on disk. Called per-request in dev."""
    import logging
    import time
    _logger = logging.getLogger("sx.reload")

    changed_files = []
    for directory in _watched_dirs:
        for fp in sorted(
            glob.glob(os.path.join(directory, "**", "*.sx"), recursive=True)
        ):
            mtime = os.path.getmtime(fp)
            if fp not in _file_mtimes or _file_mtimes[fp] != mtime:
                _file_mtimes[fp] = mtime
                changed_files.append(fp)
    if changed_files:
        for fp in changed_files:
            _logger.info("Changed: %s", fp)
        t0 = time.monotonic()
        _COMPONENT_ENV.clear()
        # Reload SX libraries first (e.g. z3.sx) so reader macros resolve
        for cb in _reload_callbacks:
            cb()
        for directory in _watched_dirs:
            load_sx_dir(directory)
        t1 = time.monotonic()
        _logger.info("Reloaded %d file(s), components in %.1fms",
                     len(changed_files), (t1 - t0) * 1000)

        # Recompute render plans for all services that have pages
        from .pages import _PAGE_REGISTRY, compute_page_render_plans
        for svc in _PAGE_REGISTRY:
            t2 = time.monotonic()
            compute_page_render_plans(svc)
            _logger.info("Render plans for %s in %.1fms", svc, (time.monotonic() - t2) * 1000)


def load_service_components(service_dir: str, service_name: str | None = None) -> None:
    """Load service-specific s-expression components and handlers.

    Components from ``{service_dir}/sx/`` and handlers from
    ``{service_dir}/sx/handlers/`` or ``{service_dir}/sx/handlers.sx``.
    """
    sx_dir = os.path.join(service_dir, "sx")
    if os.path.isdir(sx_dir):
        load_sx_dir(sx_dir)
        watch_sx_dir(sx_dir)

    # Load handler definitions if service_name is provided
    if service_name:
        load_handler_dir(os.path.join(sx_dir, "handlers"), service_name)
        # Also check for a single handlers.sx file
        handlers_file = os.path.join(sx_dir, "handlers.sx")
        if os.path.isfile(handlers_file):
            from .handlers import load_handler_file
            load_handler_file(handlers_file, service_name)


def load_handler_dir(directory: str, service_name: str) -> None:
    """Load handler .sx files from a directory if it exists."""
    if os.path.isdir(directory):
        from .handlers import load_handler_dir as _load
        _load(directory, service_name)


def register_components(sx_source: str) -> None:
    """Parse and evaluate s-expression component definitions into the
    shared environment.

    Typically called at app startup::

        register_components('''
            (defcomp ~link-card (&key link title image icon)
              (a :href link :class "block rounded ..."
                (div :class "flex ..."
                  (if image
                    (img :src image :class "...")
                    (div :class "..." (i :class icon)))
                  (div :class "..." (div :class "..." title)))))
        ''')
    """
    if _os.environ.get("SX_USE_REF") == "1":
        from .ref.sx_ref import eval_expr as _raw_eval, trampoline as _trampoline
    else:
        from .evaluator import _eval as _raw_eval, _trampoline
    _eval = lambda expr, env: _trampoline(_raw_eval(expr, env))
    from .parser import parse_all
    from .css_registry import scan_classes_from_sx

    # Snapshot existing component names before eval
    existing = set(_COMPONENT_ENV.keys())

    exprs = parse_all(sx_source)
    for expr in exprs:
        _eval(expr, _COMPONENT_ENV)

    # Pre-scan CSS classes for newly registered components.
    # Scan the full source once — components from the same file share the set.
    # Slightly over-counts per component but safe and avoids re-scanning at request time.
    all_classes: set[str] | None = None
    for key, val in _COMPONENT_ENV.items():
        if key not in existing and isinstance(val, (Component, Island)):
            if all_classes is None:
                all_classes = scan_classes_from_sx(sx_source)
            val.css_classes = set(all_classes)

    # Recompute transitive deps for all components (cheap — just AST walking)
    from .deps import compute_all_deps, compute_all_io_refs, get_all_io_names
    compute_all_deps(_COMPONENT_ENV)
    compute_all_io_refs(_COMPONENT_ENV, get_all_io_names())

    _compute_component_hash()


# ---------------------------------------------------------------------------
# sx() — render s-expression from Jinja template
# ---------------------------------------------------------------------------

def sx(source: str, **kwargs: Any) -> str:
    """Render an s-expression string to HTML.

    Keyword arguments are merged into the evaluation environment,
    so Jinja context variables can be passed through::

        {{ sx('(~link-card :title title :slug slug)',
                title=post.title, slug=post.slug) | safe }}

    This is a synchronous function — suitable for Jinja globals.
    For async resolution (with I/O primitives), use ``sx_async()``.
    """
    env = dict(_COMPONENT_ENV)
    env.update(kwargs)
    expr = parse(source)
    return html_render(expr, env)


def render(component_name: str, **kwargs: Any) -> str:
    """Call a registered component by name with Python kwargs.

    Automatically converts Python snake_case to sx kebab-case.
    No sx strings needed — just a function call.
    """
    name = component_name if component_name.startswith("~") else f"~{component_name}"
    comp = _COMPONENT_ENV.get(name)
    if not isinstance(comp, Component):
        raise ValueError(f"Unknown component: {name}")

    env = dict(_COMPONENT_ENV)
    args: list[Any] = []
    for key, val in kwargs.items():
        kw_name = key.replace("_", "-")
        args.append(Keyword(kw_name))
        args.append(val)
        env[kw_name] = val

    return _render_component(comp, args, env)


async def sx_async(source: str, **kwargs: Any) -> str:
    """Async version of ``sx()`` — resolves I/O primitives (frag, query)
    before rendering.

    Use when the s-expression contains I/O nodes::

        {{ sx_async('(frag "blog" "card" :slug "apple")') | safe }}
    """
    from .resolver import resolve, RequestContext

    env = dict(_COMPONENT_ENV)
    env.update(kwargs)
    expr = parse(source)

    # Try to get request context from Quart
    ctx = _get_request_context()
    return await resolve(expr, ctx=ctx, env=env)


def _get_request_context():
    """Build RequestContext from current Quart request, if available."""
    from .primitives_io import RequestContext
    try:
        from quart import g, request
        user = getattr(g, "user", None)
        is_htmx = bool(request.headers.get("SX-Request") or request.headers.get("HX-Request"))
        return RequestContext(user=user, is_htmx=is_htmx)
    except Exception:
        return RequestContext()


# ---------------------------------------------------------------------------
# Quart integration
# ---------------------------------------------------------------------------

def client_components_tag(*names: str) -> str:
    """Emit a <script type="text/sx"> tag with component definitions.

    Reads the source definitions from loaded .sx files and sends them
    to the client so sx.js can render them identically.

    Usage in Python::

        body_end_html = client_components_tag("test-filter-card", "test-row")

    Or send all loaded components::

        body_end_html = client_components_tag()
    """
    from .parser import serialize
    parts = []
    for key, val in _COMPONENT_ENV.items():
        if isinstance(val, Island):
            if names and val.name not in names and key.lstrip("~") not in names:
                continue
            param_strs = ["&key"] + list(val.params)
            if val.has_children:
                param_strs.extend(["&rest", "children"])
            params_sx = "(" + " ".join(param_strs) + ")"
            body_sx = serialize(val.body, pretty=True)
            parts.append(f"(defisland ~{val.name} {params_sx} {body_sx})")
        elif isinstance(val, Component):
            if names and val.name not in names and key.lstrip("~") not in names:
                continue
            # Reconstruct defcomp source from the Component object
            param_strs = ["&key"] + list(val.params)
            if val.has_children:
                param_strs.extend(["&rest", "children"])
            params_sx = "(" + " ".join(param_strs) + ")"
            body_sx = serialize(val.body, pretty=True)
            parts.append(f"(defcomp ~{val.name} {params_sx} {body_sx})")
        elif isinstance(val, Macro):
            if names and val.name not in names:
                continue
            param_strs = list(val.params)
            if val.rest_param:
                param_strs.extend(["&rest", val.rest_param])
            params_sx = "(" + " ".join(param_strs) + ")"
            body_sx = serialize(val.body, pretty=True)
            parts.append(f"(defmacro {val.name} {params_sx} {body_sx})")
    if not parts:
        return ""
    source = "\n".join(parts)
    return f'<script type="text/sx" data-components>{source}</script>'


def components_for_page(page_sx: str, service: str | None = None) -> tuple[str, str]:
    """Return (component_defs_source, page_hash) for a page.

    Scans *page_sx* for component references, computes the transitive
    closure, and returns only the definitions needed for this page.

    When *service* is given, also includes deps for all :data pages
    in that service so the client can render them without a server
    roundtrip on navigation.

    The hash is computed from the page-specific bundle for caching.
    """
    from .deps import components_needed
    from .parser import serialize

    needed = components_needed(page_sx, _COMPONENT_ENV)

    # Include deps for all :data pages so the client can render them.
    # Pages with IO deps use the async render path (Phase 5) — the IO
    # primitives are proxied via /sx/io/<name>.
    if service:
        from .pages import get_all_pages
        for page_def in get_all_pages(service).values():
            if page_def.data_expr is not None and page_def.content_expr is not None:
                content_src = serialize(page_def.content_expr)
                needed |= components_needed(content_src, _COMPONENT_ENV)

    if not needed:
        return "", ""

    # Also include macros — they're needed for client-side expansion
    parts = []
    for key, val in _COMPONENT_ENV.items():
        if isinstance(val, Island):
            if f"~{val.name}" in needed or key in needed:
                param_strs = ["&key"] + list(val.params)
                if val.has_children:
                    param_strs.extend(["&rest", "children"])
                params_sx = "(" + " ".join(param_strs) + ")"
                body_sx = serialize(val.body, pretty=True)
                parts.append(f"(defisland ~{val.name} {params_sx} {body_sx})")
        elif isinstance(val, Component):
            if f"~{val.name}" in needed or key in needed:
                # Skip server-affinity components — they're expanded server-side
                # and the client doesn't have the define values they depend on.
                if val.render_target == "server":
                    continue
                param_strs = ["&key"] + list(val.params)
                if val.has_children:
                    param_strs.extend(["&rest", "children"])
                params_sx = "(" + " ".join(param_strs) + ")"
                body_sx = serialize(val.body, pretty=True)
                parts.append(f"(defcomp ~{val.name} {params_sx} {body_sx})")
        elif isinstance(val, Macro):
            # Include macros that are referenced in needed components' bodies
            # For now, include all macros (they're small and often shared)
            param_strs = list(val.params)
            if val.rest_param:
                param_strs.extend(["&rest", val.rest_param])
            params_sx = "(" + " ".join(param_strs) + ")"
            body_sx = serialize(val.body, pretty=True)
            parts.append(f"(defmacro {val.name} {params_sx} {body_sx})")

    if not parts:
        return "", ""

    source = "\n".join(parts)
    digest = hashlib.sha256(source.encode()).hexdigest()[:12]
    return source, digest


def css_classes_for_page(page_sx: str, service: str | None = None) -> set[str]:
    """Return CSS classes needed for a page's component bundle + page source.

    Instead of unioning ALL component CSS classes, only includes classes
    from components the page actually uses (plus all :data page deps).
    """
    from .deps import components_needed
    from .css_registry import scan_classes_from_sx
    from .parser import serialize

    needed = components_needed(page_sx, _COMPONENT_ENV)

    if service:
        from .pages import get_all_pages
        for page_def in get_all_pages(service).values():
            if page_def.data_expr is not None and page_def.content_expr is not None:
                content_src = serialize(page_def.content_expr)
                needed |= components_needed(content_src, _COMPONENT_ENV)
    classes: set[str] = set()

    for key, val in _COMPONENT_ENV.items():
        if isinstance(val, (Component, Island)):
            if (f"~{val.name}" in needed or key in needed) and val.css_classes:
                classes.update(val.css_classes)

    # Page sx is unique per request — scan it
    classes.update(scan_classes_from_sx(page_sx))
    return classes


def sx_css_all() -> str:
    """Return all CSS rules (preamble + utilities) for Jinja fallback pages."""
    from .css_registry import get_all_css
    return get_all_css()


def setup_sx_bridge(app: Any) -> None:
    """Register s-expression helpers with a Quart app's Jinja environment.

    Call this in your app factory after ``setup_jinja(app)``::

        from shared.sx.jinja_bridge import setup_sx_bridge
        setup_sx_bridge(app)

    This registers:
    - ``sx(source, **kwargs)`` — sync render (components, pure HTML)
    - ``sx_async(source, **kwargs)`` — async render (with I/O resolution)
    - ``sx_css_all()`` — full CSS dump for non-sx pages
    """
    app.jinja_env.globals["sx"] = sx
    app.jinja_env.globals["render"] = render
    app.jinja_env.globals["sx_async"] = sx_async
    app.jinja_env.globals["sx_css_all"] = sx_css_all