import { ClientWidgetApi } from 'matrix-widget-api';

type MessageLike = Pick<MessageEvent, 'source' | 'origin'>;

/**
 * True when a message came from the widget's own frame, at the origin the
 * widget was loaded from.
 */
export const isFromWidgetFrame = (
  ev: MessageLike,
  frameWindow: Window | null | undefined,
  widgetOrigin: string,
): boolean => !!frameWindow && ev.source === frameWindow && ev.origin === widgetOrigin;

type InboundTransport = {
  handleMessage: (ev: MessageEvent) => void;
};

/**
 * matrix-widget-api's host transport accepts a message from ANY window on the
 * page as long as it carries the widget's id (its `strictOriginCheck` only
 * compares against the host's own origin, and is off by default). The call
 * widget's id is fixed ('call-embed'), so any other frame (a room widget, a
 * URL-preview embed) could post fromWidget actions as the call: e.g. a fake
 * push-to-talk keydown turned the user's mic on.
 *
 * Swap the transport's listener for one that only lets through messages from
 * this widget's iframe and origin. `stop()` removes `handleMessage`, which is
 * the guarded one after this. Call right after `new ClientWidgetApi(...)`,
 * which has already started the transport.
 */
export const restrictWidgetMessages = (
  api: ClientWidgetApi,
  iframe: HTMLIFrameElement,
  widgetOrigin: string,
  inbound: Pick<Window, 'addEventListener' | 'removeEventListener'> = window,
): void => {
  const transport = api.transport as unknown as InboundTransport;
  const original = transport.handleMessage;
  const guarded = (ev: MessageEvent) => {
    if (!isFromWidgetFrame(ev, iframe.contentWindow, widgetOrigin)) return;
    original(ev);
  };
  inbound.removeEventListener('message', original as EventListener);
  transport.handleMessage = guarded;
  inbound.addEventListener('message', guarded as EventListener);
};