"""Raw-binary hoisting helpers for encoded response payloads."""
from __future__ import annotations

import base64
import re

DATA_URL_RE = re.compile(r"^data:([^;,]+);base64,(.+)$", re.S)
B64_MARKER = "__b64__"

def hoist_screenshot(obj, command: str | None):
  """Replace a screenshot data URL with raw bytes so msgpack ships it unencoded.

  Gated to ``tabs.screenshot`` so arbitrary page-derived data is never touched.
  """
  if command != "tabs.screenshot" or not isinstance(obj, dict):
    return obj
  data = obj.get("data")
  if not isinstance(data, dict):
    return obj
  url = data.get("dataUrl")
  if not isinstance(url, str):
    return obj
  match = DATA_URL_RE.match(url)
  if not match:
    return obj
  try:
    raw = base64.b64decode(match.group(2))
  except Exception:
    return obj
  new_data = dict(data)
  new_data["dataUrl"] = {B64_MARKER: True, "mime": match.group(1), "raw": raw}
  return {**obj, "data": new_data}

def unhoist_binary(obj):
  """Rebuild any hoisted data URL so callers see the original string again."""
  if isinstance(obj, dict):
    raw = obj.get("raw")
    if obj.get(B64_MARKER) and isinstance(raw, (bytes, bytearray)):
      mime = obj.get("mime") or "application/octet-stream"
      return f"data:{mime};base64," + base64.b64encode(bytes(raw)).decode("ascii")
    return {key: unhoist_binary(value) for key, value in obj.items()}
  if isinstance(obj, list):
    return [unhoist_binary(value) for value in obj]
  return obj