vibn-agent-runner/vibn-frontend/components/project/preview-bridge-context.tsx

"use client";

import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useRef,
  useState,
  type ReactNode,
} from "react";
import { toast } from "sonner";
import { previewFrameTrustsMessage } from "@/components/project/preview-frame-trust";
import { dashboardBridgeScriptUrl } from "@/lib/dashboard-bridge-url";

/** Must match public/vibn-preview-bridge.js */
export const PREVIEW_BRIDGE_SOURCE = "vibn-preview";

export type PreviewPickPayload = {
  selector: string;
  tagName: string;
  textSnippet: string;
  outerHtmlSnippet?: string;
};

/**
 * Set by PreviewBridgeProvider. Unified project shell uses this to prepend the
 * latest preview element pick to the outbound chat message (then clears the pick).
 */
export const previewMessagePrepRef: { current: ((text: string) => string) | null } = {
  current: null,
};

type PreviewBridgeContextValue = {
  selectMode: boolean;
  setSelectMode: (v: boolean) => void;
  picked: PreviewPickPayload | null;
  clearPick: () => void;
  /** Preview iframe calls this when URL / element changes */
  registerPreviewIframe: (el: HTMLIFrameElement | null, src: string | null) => void;
  /** After iframe load so CMD enable/disable is delivered to the bridge */
  notifyPreviewIframeLoaded: () => void;
};

const PreviewBridgeContext = createContext<PreviewBridgeContextValue | null>(null);

/**
 * Pick targetOrigin for parent → iframe postMessage.
 *
 * - Same-origin or redirected frames: use the live document origin when href is real.
 * - Cross-origin tunnel: cannot read location → '*' (safe here: we only message this iframe's contentWindow).
 * - about:blank / srcdoc: never trust iframe.src's origin (often mismatches inherited / opaque origin).
 */
function postCmdTargetOrigin(iframe: HTMLIFrameElement): string {
  try {
    const cw = iframe.contentWindow;
    if (!cw) return "*";
    const href = cw.location.href;
    if (href && href !== "about:blank" && href !== "about:srcdoc") {
      return cw.location.origin;
    }
  } catch {
    /* Cross-origin — location not readable */
  }
  return "*";
}

function postCmdToIframe(win: Window, iframe: HTMLIFrameElement, msg: unknown) {
  const targetOrigin = postCmdTargetOrigin(iframe);
  try {
    win.postMessage(msg, targetOrigin);
  } catch {
    try {
      win.postMessage(msg, "*");
    } catch {
      /* ignore — iframe may not be ready */
    }
  }
}

export function PreviewBridgeProvider({ children }: { children: ReactNode }) {
  const [selectMode, setSelectMode] = useState(false);
  const [picked, setPicked] = useState<PreviewPickPayload | null>(null);
  const iframeRef = useRef<HTMLIFrameElement | null>(null);
  const [targetOrigin, setTargetOrigin] = useState<string | null>(null);
  const [iframeLoadGen, setIframeLoadGen] = useState(0);
  const bridgeReadyRef = useRef(false);
  /** Avoid stacking duplicate toasts (Strict Mode / rapid toggles). */
  const noBridgeWarnedRef = useRef(false);

  const registerPreviewIframe = useCallback((el: HTMLIFrameElement | null, src: string | null) => {
    iframeRef.current = el;
    bridgeReadyRef.current = false;
    // New frame — allow one no-bridge toast again if Select stays on.
    if (el) noBridgeWarnedRef.current = false;
    // React passes ref(el=null) on unmount but `src` is still the last URL — must clear
    // origin so we don't pair a stale targetOrigin with a detached iframe (breaks soft nav;
    // hard refresh remounts the provider so it masks the bug).
    if (!el) {
      setTargetOrigin(null);
      return;
    }
    if (!src) {
      setTargetOrigin(null);
      return;
    }
    try {
      setTargetOrigin(new URL(src, window.location.href).origin);
    } catch {
      setTargetOrigin(null);
    }
  }, []);

  const notifyPreviewIframeLoaded = useCallback(() => {
    setIframeLoadGen((n) => n + 1);
  }, []);

  useEffect(() => {
    const onMsg = (ev: MessageEvent) => {
      const d = ev.data;
      if (!d || d.source !== PREVIEW_BRIDGE_SOURCE) return;
      const iframeSrc = iframeRef.current?.src;

      if (d.type === "BRIDGE_READY") {
        if (!previewFrameTrustsMessage(ev.origin, targetOrigin, iframeSrc)) return;
        bridgeReadyRef.current = true;
        noBridgeWarnedRef.current = false;
        toast.dismiss("vibn-preview-no-bridge");
        return;
      }

      if (d.type !== "PICK") return;
      if (!previewFrameTrustsMessage(ev.origin, targetOrigin, iframeSrc)) return;
      setPicked(d.payload as PreviewPickPayload);
      setSelectMode(false);
    };
    window.addEventListener("message", onMsg);
    return () => window.removeEventListener("message", onMsg);
  }, [targetOrigin]);

  useEffect(() => {
    if (!selectMode) {
      noBridgeWarnedRef.current = false;
      toast.dismiss("vibn-preview-no-bridge");
      return;
    }
    const t = window.setTimeout(() => {
      if (!bridgeReadyRef.current && !noBridgeWarnedRef.current) {
        noBridgeWarnedRef.current = true;
        const scriptUrl = dashboardBridgeScriptUrl();
        const explicit =
          (process.env.NEXT_PUBLIC_VIBN_BRIDGE_URL ?? "").trim().length > 0;
        const hint = explicit
          ? `Load this script in the preview app root layout: ${scriptUrl}`
          : `Load this script in the preview app root layout (served from this dashboard): ${scriptUrl}`;
        toast.dismiss("vibn-preview-no-bridge");
        toast.error("Preview picker isn't running inside this frame", {
          id: "vibn-preview-no-bridge",
          description: hint,
          duration: 12_000,
        });
      }
    }, 1600);
    return () => window.clearTimeout(t);
  }, [selectMode]);

  useEffect(() => {
    const iframe = iframeRef.current;
    const win = iframe?.contentWindow;
    if (!win || !iframe) return;
    const msg = {
      source: PREVIEW_BRIDGE_SOURCE,
      type: "CMD",
      cmd: selectMode ? "enable-select" : "disable-select",
    };
    postCmdToIframe(win, iframe, msg);
  }, [selectMode, targetOrigin, iframeLoadGen]);

  const clearPick = useCallback(() => setPicked(null), []);

  useEffect(() => {
    previewMessagePrepRef.current = (text: string) => {
      if (!picked) return text;
      const block = [
        "[Preview selection]",
        `Selector: ${picked.selector}`,
        `Tag: ${picked.tagName}`,
        ...(picked.textSnippet ? [`Text: ${picked.textSnippet}`] : []),
        "---",
      ].join("\n");
      setPicked(null);
      return `${block}\n\n${text}`;
    };
    return () => {
      previewMessagePrepRef.current = null;
    };
  }, [picked]);

  const value: PreviewBridgeContextValue = {
    selectMode,
    setSelectMode,
    picked,
    clearPick,
    registerPreviewIframe,
    notifyPreviewIframeLoaded,
  };

  return <PreviewBridgeContext.Provider value={value}>{children}</PreviewBridgeContext.Provider>;
}

export function usePreviewBridge(): PreviewBridgeContextValue | null {
  return useContext(PreviewBridgeContext);
}