rose-ash/shared/sx/ocaml_bridge.py

"""
OCaml SX kernel ↔ Python coroutine bridge.

Manages a persistent OCaml subprocess (sx_server) that evaluates SX
expressions.  When the OCaml kernel needs IO (database queries, service
calls), it yields an ``(io-request ...)`` back to Python, which fulfills
it asynchronously and sends an ``(io-response ...)`` back.

Usage::

    bridge = OcamlBridge()
    await bridge.start()
    html = await bridge.render('(div (p "hello"))')
    await bridge.stop()
"""

from __future__ import annotations

import asyncio
import logging
import os
from typing import Any

_logger = logging.getLogger("sx.ocaml")

# Default binary path — can be overridden via SX_OCAML_BIN env var
_DEFAULT_BIN = os.path.join(
    os.path.dirname(__file__),
    "../../hosts/ocaml/_build/default/bin/sx_server.exe",
)


class OcamlBridgeError(Exception):
    """Error from the OCaml SX kernel."""


class OcamlBridge:
    """Async bridge to a persistent OCaml SX subprocess."""

    def __init__(self, binary: str | None = None):
        self._binary = binary or os.environ.get("SX_OCAML_BIN") or _DEFAULT_BIN
        self._proc: asyncio.subprocess.Process | None = None
        self._lock = asyncio.Lock()
        self._started = False
        self._components_loaded = False

    async def start(self) -> None:
        """Launch the OCaml subprocess and wait for (ready)."""
        if self._started:
            return

        bin_path = os.path.abspath(self._binary)
        if not os.path.isfile(bin_path):
            raise FileNotFoundError(
                f"OCaml SX server binary not found: {bin_path}\n"
                f"Build with: cd hosts/ocaml && eval $(opam env) && dune build"
            )

        _logger.info("Starting OCaml SX kernel: %s", bin_path)
        self._proc = await asyncio.create_subprocess_exec(
            bin_path,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        # Wait for (ready)
        line = await self._readline()
        if line != "(ready)":
            raise OcamlBridgeError(f"Expected (ready), got: {line!r}")

        self._started = True

        # Verify engine identity
        self._send("(ping)")
        kind, engine = await self._read_response()
        engine_name = engine if kind == "ok" else "unknown"
        _logger.info("OCaml SX kernel ready (pid=%d, engine=%s)", self._proc.pid, engine_name)

    async def stop(self) -> None:
        """Terminate the subprocess."""
        if self._proc and self._proc.returncode is None:
            self._proc.stdin.close()
            try:
                await asyncio.wait_for(self._proc.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                self._proc.kill()
                await self._proc.wait()
            _logger.info("OCaml SX kernel stopped")
        self._proc = None
        self._started = False

    async def ping(self) -> str:
        """Health check — returns engine name (e.g. 'ocaml-cek')."""
        async with self._lock:
            self._send("(ping)")
            kind, value = await self._read_response()
            return value or "" if kind == "ok" else ""

    async def load(self, path: str) -> int:
        """Load an .sx file for side effects (defcomp, define, defmacro)."""
        async with self._lock:
            self._send(f'(load "{_escape(path)}")')
            value = await self._read_until_ok(ctx=None)
            return int(float(value)) if value else 0

    async def load_source(self, source: str) -> int:
        """Evaluate SX source for side effects."""
        async with self._lock:
            self._send(f'(load-source "{_escape(source)}")')
            value = await self._read_until_ok(ctx=None)
            return int(float(value)) if value else 0

    async def eval(self, source: str, ctx: dict[str, Any] | None = None) -> str:
        """Evaluate SX expression, return serialized result.

        Supports io-requests (helper calls, query, action, etc.) via the
        coroutine bridge, just like render().
        """
        await self._ensure_components()
        async with self._lock:
            self._send(f'(eval "{_escape(source)}")')
            return await self._read_until_ok(ctx)

    async def render(
        self,
        source: str,
        ctx: dict[str, Any] | None = None,
    ) -> str:
        """Render SX to HTML, handling io-requests via Python async IO."""
        await self._ensure_components()
        async with self._lock:
            self._send(f'(render "{_escape(source)}")')
            return await self._read_until_ok(ctx)

    async def aser(self, source: str, ctx: dict[str, Any] | None = None) -> str:
        """Evaluate SX and return SX wire format, handling io-requests."""
        await self._ensure_components()
        async with self._lock:
            self._send(f'(aser "{_escape(source)}")')
            return await self._read_until_ok(ctx)

    async def _ensure_components(self) -> None:
        """Load all .sx source files into the kernel on first use.

        Errors during loading are handled gracefully — IO responses are
        always sent back to keep the pipe clean.
        """
        if self._components_loaded:
            return
        self._components_loaded = True
        try:
            from .jinja_bridge import _watched_dirs, _dirs_from_cache
            import glob

            # Skip patterns — files that use constructs not available in the kernel
            skip_names = {"boundary.sx", "forms.sx"}
            skip_dirs = {"tests"}

            # Collect files to load
            all_files: list[str] = []

            # Spec files needed by aser
            spec_dir = os.path.join(os.path.dirname(__file__), "../../spec")
            for spec_file in ["parser.sx", "render.sx"]:
                path = os.path.normpath(os.path.join(spec_dir, spec_file))
                if os.path.isfile(path):
                    all_files.append(path)

            # All directories loaded into the Python env
            all_dirs = list(set(_watched_dirs) | _dirs_from_cache)

            # Web adapters (aser lives in adapter-sx.sx) — only load specific files
            web_dir = os.path.join(os.path.dirname(__file__), "../../web")
            if os.path.isdir(web_dir):
                for web_file in ["adapter-sx.sx"]:
                    path = os.path.normpath(os.path.join(web_dir, web_file))
                    if os.path.isfile(path):
                        all_files.append(path)

            for directory in sorted(all_dirs):
                files = sorted(
                    glob.glob(os.path.join(directory, "**", "*.sx"), recursive=True)
                )
                for filepath in files:
                    basename = os.path.basename(filepath)
                    # Skip known-bad files
                    if basename in skip_names:
                        continue
                    # Skip test and handler directories
                    parts = filepath.replace("\\", "/").split("/")
                    if any(d in skip_dirs for d in parts):
                        continue
                    all_files.append(filepath)

            # Load all files under a single lock
            count = 0
            skipped = 0
            async with self._lock:
                for filepath in all_files:
                    try:
                        self._send(f'(load "{_escape(filepath)}")')
                        value = await self._read_until_ok(ctx=None)
                        # Response may be a number (count) or a value — just count files
                        count += 1
                    except OcamlBridgeError as e:
                        skipped += 1
                        _logger.warning("OCaml load skipped %s: %s",
                                        filepath, e)
            _logger.info("Loaded %d definitions from .sx files into OCaml kernel (%d skipped)",
                         count, skipped)
        except Exception as e:
            _logger.error("Failed to load .sx files into OCaml kernel: %s", e)
            self._components_loaded = False  # retry next time

    async def reset(self) -> None:
        """Reset the kernel environment to pristine state."""
        async with self._lock:
            self._send("(reset)")
            kind, value = await self._read_response()
            if kind == "error":
                raise OcamlBridgeError(f"reset: {value}")

    # ------------------------------------------------------------------
    # Internal protocol handling
    # ------------------------------------------------------------------

    def _send(self, line: str) -> None:
        """Write a line to the subprocess stdin."""
        assert self._proc and self._proc.stdin
        self._proc.stdin.write((line + "\n").encode())

    async def _readline(self) -> str:
        """Read a line from the subprocess stdout."""
        assert self._proc and self._proc.stdout
        data = await self._proc.stdout.readline()
        if not data:
            # Process died — collect stderr for diagnostics
            stderr = b""
            if self._proc.stderr:
                stderr = await self._proc.stderr.read()
            raise OcamlBridgeError(
                f"OCaml subprocess died unexpectedly. stderr: {stderr.decode(errors='replace')}"
            )
        return data.decode().rstrip("\n")

    async def _read_response(self) -> tuple[str, str | None]:
        """Read a single (ok ...) or (error ...) response.

        Returns (kind, value) where kind is "ok" or "error".
        """
        line = await self._readline()
        return _parse_response(line)

    async def _read_until_ok(
        self,
        ctx: dict[str, Any] | None = None,
    ) -> str:
        """Read lines until (ok ...) or (error ...).

        Handles (io-request ...) by fulfilling IO and sending (io-response ...).
        ALWAYS sends a response to keep the pipe clean, even on error.
        """
        while True:
            line = await self._readline()

            if line.startswith("(io-request "):
                try:
                    result = await self._handle_io_request(line, ctx)
                    self._send(f"(io-response {_serialize_for_ocaml(result)})")
                except Exception as e:
                    # MUST send a response or the pipe desyncs
                    _logger.warning("IO request failed, sending nil: %s", e)
                    self._send("(io-response nil)")
                continue

            kind, value = _parse_response(line)
            if kind == "error":
                raise OcamlBridgeError(value or "Unknown error")
            # kind == "ok"
            return value or ""

    async def _handle_io_request(
        self,
        line: str,
        ctx: dict[str, Any] | None,
    ) -> Any:
        """Dispatch an io-request to the appropriate Python handler."""
        from .parser import parse_all

        # Parse the io-request
        parsed = parse_all(line)
        if not parsed or not isinstance(parsed[0], list):
            raise OcamlBridgeError(f"Malformed io-request: {line}")

        parts = parsed[0]
        # parts = [Symbol("io-request"), name_str, ...args]
        if len(parts) < 2:
            raise OcamlBridgeError(f"Malformed io-request: {line}")

        req_name = _to_str(parts[1])
        args = parts[2:]

        if req_name == "query":
            return await self._io_query(args)
        elif req_name == "action":
            return await self._io_action(args)
        elif req_name == "request-arg":
            return self._io_request_arg(args)
        elif req_name == "request-method":
            return self._io_request_method()
        elif req_name == "ctx":
            return self._io_ctx(args, ctx)
        elif req_name == "helper":
            return await self._io_helper(args, ctx)
        else:
            # Fall back to registered IO handlers (set-response-status, sleep, etc.)
            from .primitives_io import _IO_HANDLERS, RequestContext
            io_handler = _IO_HANDLERS.get(req_name)
            if io_handler is not None:
                helper_args = [_to_python(a) for a in args]
                return await io_handler(helper_args, {}, ctx or RequestContext())
            raise OcamlBridgeError(f"Unknown io-request type: {req_name}")

    async def _io_query(self, args: list) -> Any:
        """Handle (io-request "query" service name params...)."""
        from shared.infrastructure.internal import fetch_data

        service = _to_str(args[0]) if len(args) > 0 else ""
        query = _to_str(args[1]) if len(args) > 1 else ""
        params = _to_dict(args[2]) if len(args) > 2 else {}
        return await fetch_data(service, query, params)

    async def _io_action(self, args: list) -> Any:
        """Handle (io-request "action" service name payload...)."""
        from shared.infrastructure.internal import call_action

        service = _to_str(args[0]) if len(args) > 0 else ""
        action = _to_str(args[1]) if len(args) > 1 else ""
        payload = _to_dict(args[2]) if len(args) > 2 else {}
        return await call_action(service, action, payload)

    def _io_request_arg(self, args: list) -> Any:
        """Handle (io-request "request-arg" name)."""
        try:
            from quart import request
            name = _to_str(args[0]) if args else ""
            return request.args.get(name)
        except RuntimeError:
            return None

    def _io_request_method(self) -> str:
        """Handle (io-request "request-method")."""
        try:
            from quart import request
            return request.method
        except RuntimeError:
            return "GET"

    def _io_ctx(self, args: list, ctx: dict[str, Any] | None) -> Any:
        """Handle (io-request "ctx" key)."""
        if ctx is None:
            return None
        key = _to_str(args[0]) if args else ""
        return ctx.get(key)

    async def _io_helper(self, args: list, ctx: dict[str, Any] | None) -> Any:
        """Handle (io-request "helper" name arg1 arg2 ...).

        Dispatches to registered page helpers — Python functions like
        read-spec-file, bootstrapper-data, etc.  The helper service name
        is passed via ctx["_helper_service"].
        """
        import asyncio
        from .pages import get_page_helpers
        from .primitives_io import _IO_HANDLERS, RequestContext

        name = _to_str(args[0]) if args else ""
        helper_args = [_to_python(a) for a in args[1:]]

        # Check page helpers first (application-level)
        service = (ctx or {}).get("_helper_service", "sx")
        helpers = get_page_helpers(service)
        fn = helpers.get(name)
        if fn is not None:
            result = fn(*helper_args)
            if asyncio.iscoroutine(result):
                result = await result
            return result

        # Fall back to IO primitives (now, state-get, state-set!, etc.)
        io_handler = _IO_HANDLERS.get(name)
        if io_handler is not None:
            return await io_handler(helper_args, {}, RequestContext())

        # Fall back to regular primitives (json-encode, into, etc.)
        from .primitives import get_primitive as _get_prim
        prim = _get_prim(name)
        if prim is not None:
            return prim(*helper_args)

        raise OcamlBridgeError(f"Unknown helper: {name!r}")


# ------------------------------------------------------------------
# Module-level singleton
# ------------------------------------------------------------------

_bridge: OcamlBridge | None = None


async def get_bridge() -> OcamlBridge:
    """Get or create the singleton bridge instance."""
    global _bridge
    if _bridge is None:
        _bridge = OcamlBridge()
    if not _bridge._started:
        await _bridge.start()
    return _bridge


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------

def _escape(s: str) -> str:
    """Escape a string for embedding in an SX string literal."""
    return s.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "\\r").replace("\t", "\\t")


def _parse_response(line: str) -> tuple[str, str | None]:
    """Parse an (ok ...) or (error ...) response line.

    Returns (kind, value) tuple.
    """
    line = line.strip()
    if line == "(ok)":
        return ("ok", None)
    if line.startswith("(ok-raw "):
        # Raw SX wire format — no unescaping needed
        return ("ok", line[8:-1])
    if line.startswith("(ok "):
        value = line[4:-1]  # strip (ok  and )
        # If the value is a quoted string, unquote it
        if value.startswith('"') and value.endswith('"'):
            value = _unescape(value[1:-1])
        return ("ok", value)
    if line.startswith("(error "):
        msg = line[7:-1]
        if msg.startswith('"') and msg.endswith('"'):
            msg = _unescape(msg[1:-1])
        return ("error", msg)
    return ("error", f"Unexpected response: {line}")


def _unescape(s: str) -> str:
    """Unescape an SX string literal."""
    return (
        s.replace("\\n", "\n")
        .replace("\\r", "\r")
        .replace("\\t", "\t")
        .replace('\\"', '"')
        .replace("\\\\", "\\")
    )


def _to_python(val: Any) -> Any:
    """Convert an SX parsed value to a plain Python value."""
    from .types import NIL as _NIL
    if val is None or val is _NIL:
        return None
    if hasattr(val, "name"):  # Symbol or Keyword
        return val.name
    return val


def _to_str(val: Any) -> str:
    """Convert an SX parsed value to a Python string."""
    if isinstance(val, str):
        return val
    if hasattr(val, "name"):
        return val.name
    return str(val)


def _to_dict(val: Any) -> dict:
    """Convert an SX parsed value to a Python dict."""
    if isinstance(val, dict):
        return val
    return {}


def _serialize_for_ocaml(val: Any) -> str:
    """Serialize a Python value to SX text for sending to OCaml."""
    if val is None:
        return "nil"
    if isinstance(val, bool):
        return "true" if val else "false"
    if isinstance(val, (int, float)):
        if isinstance(val, float) and val == int(val):
            return str(int(val))
        return str(val)
    if isinstance(val, str):
        return f'"{_escape(val)}"'
    if isinstance(val, (list, tuple)):
        items = " ".join(_serialize_for_ocaml(v) for v in val)
        return f"(list {items})"
    if isinstance(val, dict):
        pairs = " ".join(
            f":{k} {_serialize_for_ocaml(v)}" for k, v in val.items()
        )
        return "{" + pairs + "}"
    return f'"{_escape(str(val))}"'