_tab.py

from __future__ import annotations

import base64
import json as _stdlib_json
from typing import TYPE_CHECKING, Any

import logistro

try:
    import orjson
except ImportError:  # pragma: no cover - exercised only when orjson is absent
    orjson = None  # type: ignore[assignment]

from . import _devtools_utils as _dtools
from . import _js_logger
from ._errors import _raise_error

if TYPE_CHECKING:
    import asyncio
    from pathlib import Path

    import choreographer as choreo

    from kaleido._utils import fig_tools


_TEXT_FORMATS = ("svg", "json")  # eps
_CHUNK_SIZE = 10 * 1024 * 1024  # 10 MB

_logger = logistro.getLogger(__name__)


def _orjson_default(obj):
    """Fallback for types orjson can't handle natively (e.g. NumPy string arrays)."""
    if hasattr(obj, "tolist"):
        return obj.tolist()
    raise TypeError(f"Type is not JSON serializable: {type(obj).__name__}")


class _StdlibJSONEncoder(_stdlib_json.JSONEncoder):
    """
    Encoder used when ``orjson`` is unavailable; mirrors ``_orjson_default``.

    Reproduces the ``orjson.OPT_SERIALIZE_NUMPY`` behavior via the standard
    ``.tolist()`` round-trip so callers see the same output regardless of
    whether ``orjson`` is installed.
    """

    def default(self, o: Any) -> Any:
        if hasattr(o, "tolist"):
            return o.tolist()
        return super().default(o)


def _serialize_spec(spec: Any) -> str:
    """
    Serialize a figure spec to a JSON string.

    Uses :mod:`orjson` when available (fast path with native NumPy support);
    falls back to the standard-library :mod:`json` module otherwise.
    """
    if orjson is not None:
        return orjson.dumps(
            spec,
            default=_orjson_default,
            option=orjson.OPT_SERIALIZE_NUMPY,
        ).decode()
    return _stdlib_json.dumps(spec, cls=_StdlibJSONEncoder)


def _subscribe_new(tab: choreo.Tab, event: str) -> asyncio.Future:
    """Create subscription to tab clearing old ones first: helper function."""
    new_future = tab.subscribe_once(event)
    while new_future.done():
        _logger.debug2(f"Clearing an old {event}")
        new_future = tab.subscribe_once(event)
    return new_future


class _KaleidoTab:
    """
    A Kaleido tab is a wrapped choreographer tab providing the functions we need.

    The choreographer tab can be accessed through the `self.tab` attribute.
    """

    tab: choreo.Tab
    """The underlying choreographer tab."""
    js_logger: _js_logger.JavascriptLogger
    """A log for recording javascript."""

    def __init__(self, tab, *, headers: dict[str, str] | None = None):
        """
        Create a new _KaleidoTab.

        Args:
            tab: the choreographer tab to wrap.

            headers (dict[str, str] | None, optional):
                Extra HTTP headers to send with every request made by the
                browser tab. Defaults to None.

        """
        self.tab = tab
        self._headers = headers
        self.js_logger = _js_logger.JavascriptLogger(self.tab)

    async def navigate(self, url: str | Path = ""):
        """
        Navigate to the kaleidofier script. This is effectively the real initialization.

        Args:
            url: Override the location of the kaleidofier script if necessary.

        """
        # Subscribe to event which will contain javascript engine ID (need it
        # for calling javascript functions)
        javascript_ready = _subscribe_new(self.tab, "Runtime.executionContextCreated")

        # Subscribe to event indicating page ready.
        page_ready = _subscribe_new(self.tab, "Page.loadEventFired")

        # Apply headers if they exist
        await self._apply_headers()

        # Navigating page. This will trigger the above events.
        _logger.debug2(f"Calling Page.navigate on {self.tab}")
        _raise_error(await self.tab.send_command("Page.navigate", params={"url": url}))

        # Enabling page events (for page_ready- like all events, if already
        # ready, the latest will fire immediately)
        _logger.debug2(f"Calling Page.enable on {self.tab}")
        _raise_error(await self.tab.send_command("Page.enable"))

        # Enabling javascript events (for javascript_ready)
        _logger.debug2(f"Calling Runtime.enable on {self.tab}")
        _raise_error(await self.tab.send_command("Runtime.enable"))

        self._current_js_id = _dtools.get_js_id(await javascript_ready)

        await page_ready  # don't care result, ready is ready

        # this runs *after* page load because running it first thing
        # requires a couple extra lines
        self.js_logger.reset()

    # reload is truly so close to navigate
    async def reload(self):
        """Reload the tab, and set the javascript runtime id."""
        _logger.debug(f"Reloading tab {self.tab} with javascript.")

        javascript_ready = _subscribe_new(self.tab, "Runtime.executionContextCreated")

        page_ready = _subscribe_new(self.tab, "Page.loadEventFired")

        _logger.debug2(f"Calling Page.reload on {self.tab}")
        _raise_error(await self.tab.send_command("Page.reload"))

        self._current_js_id = _dtools.get_js_id(await javascript_ready)

        await page_ready

        self.js_logger.reset()

    async def _apply_headers(self):
        """Apply extra HTTP headers to the tab if configured."""
        if self._headers:
            _logger.debug(f"Setting extra HTTP headers on {self.tab}")
            _logger.debug2(f"Extra headers are: {self._headers}")
            _raise_error(await self.tab.send_command("Network.enable"))
            _raise_error(
                await self.tab.send_command(
                    "Network.setExtraHTTPHeaders",
                    params={"headers": self._headers},
                )
            )

    async def _calc_fig(
        self,
        spec: fig_tools.Spec,
        *,
        topojson: str | None,
        render_prof,
        stepper,
    ) -> bytes:
        render_prof.profile_log.tick("serializing spec")
        spec_str = _serialize_spec(spec)
        render_prof.profile_log.tick("spec serialized")

        render_prof.profile_log.tick("sending javascript")
        if len(spec_str) <= _CHUNK_SIZE:
            kaleido_js_fn = (
                r"function(specStr, ...args)"
                r"{"
                r"return kaleido_scopes"
                r".plotly(JSON.parse(specStr), ...args)"
                r".then(JSON.stringify);"
                r"}"
            )
            result = await _dtools.exec_js_fn(
                self.tab,
                self._current_js_id,
                kaleido_js_fn,
                spec_str,
                topojson,
                stepper,
            )
        else:
            result = await self._calc_fig_chunked(
                spec_str,
                topojson=topojson,
                stepper=stepper,
            )
        _raise_error(result)
        render_prof.profile_log.tick("javascript sent")

        _logger.debug2(f"Result of function call: {result}")
        js_response = _dtools.check_kaleido_js_response(result)

        if (response_format := js_response.get("format")) == "pdf":
            render_prof.profile_log.tick("printing pdf")
            img_raw = await _dtools.print_pdf(self.tab)
            render_prof.profile_log.tick("pdf printed")
        else:
            img_raw = js_response["result"]

        if response_format not in _TEXT_FORMATS:
            res = base64.b64decode(img_raw)
        else:
            res = str.encode(img_raw)

        render_prof.data_out_size = len(res)
        render_prof.js_log = self.js_logger.log
        return res

    async def _calc_fig_chunked(
        self,
        spec_str: str,
        *,
        topojson: str | None,
        stepper,
    ):
        _raise_error(
            await _dtools.exec_js_fn(
                self.tab,
                self._current_js_id,
                r"function() { window.__kaleido_chunks = []; }",
            )
        )

        for i in range(0, len(spec_str), _CHUNK_SIZE):
            chunk = spec_str[i : i + _CHUNK_SIZE]
            _raise_error(
                await _dtools.exec_js_fn(
                    self.tab,
                    self._current_js_id,
                    r"function(c) { window.__kaleido_chunks.push(c); }",
                    chunk,
                )
            )

        kaleido_js_fn = (
            r"function(...args)"
            r"{"
            r"var spec = JSON.parse(window.__kaleido_chunks.join(''));"
            r"delete window.__kaleido_chunks;"
            r"return kaleido_scopes.plotly(spec, ...args).then(JSON.stringify);"
            r"}"
        )
        return await _dtools.exec_js_fn(
            self.tab,
            self._current_js_id,
            kaleido_js_fn,
            topojson,
            stepper,
        )