vscode-bridge.utils.ts

import type {
  AppMessage,
  HostMessage,
  PayloadOf,
} from '@lemoncode/quickmock-bridge-protocol';
import { isVSCodeEnv } from './env.utils';

type HandlerFor<T extends HostMessage['type']> = (
  payload: PayloadOf<HostMessage, T>
) => void;

type AnyHandler = (payload: unknown) => void;

const handlers = new Map<string, Set<AnyHandler>>();

const resolveParentOrigin = (): string => {
  if (!document.referrer) return '*';
  try {
    return new URL(document.referrer).origin;
  } catch {
    return '*';
  }
};

const parentOrigin = resolveParentOrigin();

export const sendToExtension = (msg: AppMessage): void => {
  if (!isVSCodeEnv()) return;
  window.parent.postMessage(msg, parentOrigin);
};

export const onMessage = <T extends HostMessage['type']>(
  type: T,
  handler: HandlerFor<T>
): (() => void) => {
  if (!isVSCodeEnv()) return () => {};

  const existing = handlers.get(type) ?? new Set<AnyHandler>();
  existing.add(handler as AnyHandler);
  handlers.set(type, existing);

  return () => {
    const set = handlers.get(type);
    if (!set) return;
    set.delete(handler as AnyHandler);
    if (set.size === 0) handlers.delete(type);
  };
};

if (isVSCodeEnv()) {
  window.addEventListener('message', (event: MessageEvent) => {
    if (event.source !== window.parent) return;

    const msg = event.data as Partial<HostMessage> | undefined;
    if (!msg?.type) return;

    const set = handlers.get(msg.type);
    if (!set) return;

    const payload = (msg as { payload?: unknown }).payload;
    for (const handler of set) handler(payload);
  });
}