"""Rerun helpers — init + jpeg-encode for fast image logging.

Two functions. ``init_web`` starts the gRPC + web viewer (raiden's
two-call pattern; ``localhost`` fails on russet — use ``127.0.0.1``
in SSH tunnels). ``to_jpeg`` resizes + JPEG-encodes before logging —
sending compressed bytes over gRPC is ~50× lighter than raw HD1080
RGB and prevents the viewer-backlog warnings.
"""
from __future__ import annotations

import cv2
import numpy as np
import rerun as rr


def init_web(name: str, port: int = 9092, send_default_blueprint: bool = True):
    """Initialize rerun + serve web. ``port`` = web UI, ``port+1`` = grpc.

    With ``send_default_blueprint=True`` (deploy default), pushes a
    programmatic ``rrb.Blueprint`` containing every panel deploy.py logs.
    This guarantees the panels appear on every browser connect, regardless
    of stale localStorage from previous runs (which is what hides
    ``pred/z_grid`` / ``pred/rot_grid`` / ``pred/pca_dino_mlp`` after a
    hard-refresh otherwise).
    """
    grpc_port = port + 1
    rr.init(name, spawn=False)
    server_uri = rr.serve_grpc(grpc_port=grpc_port)
    rr.serve_web_viewer(web_port=port, open_browser=False)
    if send_default_blueprint:
        try:
            import rerun.blueprint as rrb
            bp = rrb.Blueprint(
                rrb.Grid(
                    rrb.Spatial2DView(name="scene RGB",     origin="scene/rgb"),
                    rrb.Spatial2DView(name="scene composite", origin="scene/composite"),
                    rrb.Spatial2DView(name="scene mask",    origin="scene/mask_overlay"),
                    rrb.Spatial2DView(name="wrist RGB",      origin="wrist/rgb"),
                    rrb.Spatial2DView(name="pred kp",        origin="pred/kp"),
                    rrb.Spatial2DView(name="pred kp scene",  origin="pred/kp_scene"),
                    rrb.Spatial2DView(name="pred kp wrist",  origin="pred/kp_wrist"),
                    rrb.Spatial2DView(name="pred heatmap",   origin="pred/heatmap_grid"),
                    rrb.Spatial2DView(name="pred heatmap raw", origin="pred/heatmap_grid_raw"),
                    rrb.Spatial2DView(name="pred heatmap scene",     origin="pred/heatmap_scene"),
                    rrb.Spatial2DView(name="pred heatmap wrist",     origin="pred/heatmap_wrist"),
                    rrb.Spatial2DView(name="pred heatmap raw scene", origin="pred/heatmap_raw_scene"),
                    rrb.Spatial2DView(name="pred heatmap raw wrist", origin="pred/heatmap_raw_wrist"),
                    rrb.Spatial2DView(name="pred PCA",        origin="pred/pca_dino_mlp"),
                    rrb.Spatial2DView(name="pred PCA scene",  origin="pred/pca_scene"),
                    rrb.Spatial2DView(name="pred PCA wrist",  origin="pred/pca_wrist"),
                    rrb.Spatial2DView(name="pred z grid",     origin="pred/z_grid"),
                    rrb.Spatial2DView(name="pred grip grid",  origin="pred/grip_grid"),
                    rrb.Spatial2DView(name="pred rot grid",   origin="pred/rot_grid"),
                    rrb.TextDocumentView(name="status",       origin="status"),
                ),
                collapse_panels=True,
            )
            rr.send_blueprint(bp)
            print("  ✓ sent default blueprint (all panels forced visible)")
        except Exception as e:
            print(f"  blueprint send failed ({e}); rely on browser layout cache")
    print(f"rerun: http://localhost:{port}"
          f"?url=rerun%2Bhttp%3A%2F%2F127.0.0.1%3A{grpc_port}%2Fproxy")
    return server_uri


def to_jpeg(img_bgr: np.ndarray, target_w: int | None = 960, quality: int = 75):
    """Resize → JPEG-encode → return ``rr.EncodedImage`` ready for ``rr.log``.

    Input is BGR (the cv2 convention). Pass to ``rr.log("path", to_jpeg(img))``.
    Pass ``target_w=None`` to skip the resize (preserve native resolution).
    """
    h, w = img_bgr.shape[:2]
    if target_w is not None and w != target_w:
        scale = target_w / w
        img_bgr = cv2.resize(img_bgr, (target_w, int(h * scale)),
                             interpolation=cv2.INTER_AREA)
    ok, buf = cv2.imencode(".jpg", img_bgr, [cv2.IMWRITE_JPEG_QUALITY, quality])
    if not ok:
        return rr.Image(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
    return rr.EncodedImage(contents=bytes(buf), media_type="image/jpeg")