ubiquitous-invention/apps/web/components/editor/extensions/ai-block.tsx

"use client";

import { Node, mergeAttributes, nodeInputRule } from "@tiptap/react";
import {
  NodeViewWrapper,
  ReactNodeViewRenderer,
  type NodeViewProps,
} from "@tiptap/react";
import { streamText } from "ai";
import {
  Check,
  Languages,
  Loader2,
  Minimize2,
  RefreshCw,
  Sparkles,
  TextQuote,
  Trash2,
  Wand2,
  X,
} from "lucide-react";
import * as React from "react";

import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { cn } from "@/lib/utils";

import {
  createOpenAIClient,
  EDITOR_SYSTEM_PROMPT,
  expand,
  rewrite,
  selectOpenAIModel,
  summarize,
  translate,
} from "../../../../../packages/ai/src";

type AiStatus = "idle" | "loading" | "done" | "error";

function escapeHtml(s: string) {
  return s
    .replace(/&/g, "&")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/"/g, "&quot;");
}

function plainTextToHtml(text: string) {
  const trimmed = text.trim();
  if (!trimmed) return "<p></p>";
  return trimmed
    .split(/\n\n+/)
    .map((block) => `<p>${escapeHtml(block).replace(/\n/g, "<br/>")}</p>`)
    .join("");
}

async function runEditorGeneration(prompt: string, signal: AbortSignal) {
  const client = createOpenAIClient();
  if (!client.ok) {
    throw new Error(
      "OpenAI API key is not configured. Set OPENAI_API_KEY or NEXT_PUBLIC_OPENAI_API_KEY.",
    );
  }
  const model = selectOpenAIModel(client.provider, "gpt-4o-mini");
  const result = streamText({
    model,
    abortSignal: signal,
    system: EDITOR_SYSTEM_PROMPT,
    prompt: prompt.trim(),
  });
  let out = "";
  for await (const chunk of result.textStream) {
    out += chunk;
  }
  return out;
}

async function runActionPair(
  systemPrompt: string,
  userPrompt: string,
  signal: AbortSignal,
) {
  const client = createOpenAIClient();
  if (!client.ok) {
    throw new Error(
      "OpenAI API key is not configured. Set OPENAI_API_KEY or NEXT_PUBLIC_OPENAI_API_KEY.",
    );
  }
  const model = selectOpenAIModel(client.provider, "gpt-4o-mini");
  const result = streamText({
    model,
    abortSignal: signal,
    system: systemPrompt,
    prompt: userPrompt,
  });
  let out = "";
  for await (const chunk of result.textStream) {
    out += chunk;
  }
  return out;
}

function LoadingDots() {
  return (
    <span className="inline-flex items-center gap-0.5 px-1" aria-hidden>
      {[0, 1, 2].map((i) => (
        <span
          key={i}
          className="inline-block size-1.5 animate-pulse rounded-full bg-primary/70"
          style={{ animationDelay: `${i * 120}ms` }}
        />
      ))}
    </span>
  );
}

function AiBlockView(props: NodeViewProps) {
  const { node, editor, getPos, updateAttributes, deleteNode } = props;
  const status = (node.attrs.status as AiStatus) ?? "idle";
  const prompt = (node.attrs.prompt as string) ?? "";
  const resultHtml = (node.attrs.resultHtml as string) ?? "";
  const errorMessage = (node.attrs.errorMessage as string) ?? "";

  const inputRef = React.useRef<HTMLInputElement>(null);
  const abortRef = React.useRef<AbortController | null>(null);

  React.useEffect(() => {
    if (status === "idle" && !resultHtml) {
      inputRef.current?.focus();
    }
  }, [status, resultHtml]);

  const [selTick, setSelTick] = React.useState(0);
  React.useEffect(() => {
    const bump = () => setSelTick((n) => n + 1);
    editor.on("selectionUpdate", bump);
    editor.on("transaction", bump);
    return () => {
      editor.off("selectionUpdate", bump);
      editor.off("transaction", bump);
    };
  }, [editor]);

  const selectionText = React.useMemo(() => {
    const { from, to } = editor.state.selection;
    if (from === to) return "";
    return editor.state.doc.textBetween(from, to, "\n");
    // eslint-disable-next-line react-hooks/exhaustive-deps -- sync selection text when editor updates
  }, [editor, selTick]);

  const stopGeneration = React.useCallback(() => {
    abortRef.current?.abort();
    abortRef.current = null;
  }, []);

  const runPrompt = React.useCallback(
    async (text: string) => {
      const trimmed = text.trim();
      if (!trimmed) return;
      stopGeneration();
      const ac = new AbortController();
      abortRef.current = ac;
      updateAttributes({
        prompt: trimmed,
        status: "loading",
        resultHtml: "",
        errorMessage: "",
      });
      try {
        const out = await runEditorGeneration(trimmed, ac.signal);
        updateAttributes({
          status: "done",
          resultHtml: plainTextToHtml(out),
          errorMessage: "",
        });
      } catch (e) {
        if ((e as Error).name === "AbortError") return;
        updateAttributes({
          status: "error",
          errorMessage: (e as Error).message ?? "Something went wrong.",
        });
      } finally {
        abortRef.current = null;
      }
    },
    [stopGeneration, updateAttributes],
  );

  const runQuick = React.useCallback(
    async (kind: "summarize" | "expand" | "simplify" | "translate") => {
      const base = selectionText.trim() || prompt.trim();
      if (!base) {
        updateAttributes({
          status: "error",
          errorMessage: "Select text in the editor or enter a prompt first.",
        });
        return;
      }
      stopGeneration();
      const ac = new AbortController();
      abortRef.current = ac;
      updateAttributes({
        status: "loading",
        resultHtml: "",
        errorMessage: "",
      });
      try {
        let pair;
        if (kind === "summarize") pair = summarize(base);
        else if (kind === "expand") pair = expand(base);
        else if (kind === "simplify") pair = rewrite(base, "concise");
        else pair = translate(base, "Spanish");

        const out = await runActionPair(
          pair.systemPrompt,
          pair.userPrompt,
          ac.signal,
        );
        updateAttributes({
          status: "done",
          prompt: prompt || `[${kind}]`,
          resultHtml: plainTextToHtml(out),
          errorMessage: "",
        });
      } catch (e) {
        if ((e as Error).name === "AbortError") return;
        updateAttributes({
          status: "error",
          errorMessage: (e as Error).message ?? "Something went wrong.",
        });
      } finally {
        abortRef.current = null;
      }
    },
    [prompt, selectionText, stopGeneration, updateAttributes],
  );

  const accept = React.useCallback(() => {
    const pos = getPos();
    if (typeof pos !== "number") return;
    const html = resultHtml || "<p></p>";
    editor
      .chain()
      .focus()
      .deleteRange({ from: pos, to: pos + node.nodeSize })
      .insertContentAt(pos, html)
      .run();
  }, [editor, getPos, node.nodeSize, resultHtml]);

  const discard = React.useCallback(() => {
    deleteNode();
  }, [deleteNode]);

  const regenerate = React.useCallback(() => {
    if (prompt.trim()) void runPrompt(prompt);
  }, [prompt, runPrompt]);

  const onKeyDownInput = (e: React.KeyboardEvent<HTMLInputElement>) => {
    if (e.key === "Enter") {
      e.preventDefault();
      void runPrompt(inputRef.current?.value ?? "");
    }
    if (e.key === "Escape") {
      e.preventDefault();
      discard();
    }
  };

  return (
    <NodeViewWrapper
      className={cn(
        "ai-block my-4 rounded-xl border border-border/80 bg-gradient-to-br from-muted/40 via-background to-muted/30 p-4 shadow-sm ring-1 ring-border/50",
        props.selected && "ring-2 ring-primary/35",
      )}
      data-type="ai-block"
    >
      <div className="mb-3 flex items-center gap-2">
        <span className="flex size-8 items-center justify-center rounded-lg border border-border bg-background shadow-sm">
          <Sparkles className="size-4 text-primary" />
        </span>
        <div className="min-w-0 flex-1">
          <div className="flex items-center gap-2">
            <span className="text-sm font-semibold tracking-tight">
              AI assistant
            </span>
            <Badge variant="muted" className="text-[10px] uppercase">
              /ai
            </Badge>
          </div>
          <p className="text-xs text-muted-foreground">
            Ask in natural language, then accept or regenerate.
          </p>
        </div>
      </div>

      {status !== "loading" && !(status === "done" && resultHtml) ? (
        <div className="space-y-2">
          <Input
            ref={inputRef}
            className="h-11 border-dashed bg-background/80 text-[15px]"
            placeholder='e.g. "Write a short PRD for mobile offline mode"'
            defaultValue={prompt}
            key={prompt}
            onKeyDown={onKeyDownInput}
          />
          <div className="flex flex-wrap gap-1.5">
            <QuickBtn
              icon={TextQuote}
              label="Summarize"
              onClick={() => void runQuick("summarize")}
            />
            <QuickBtn
              icon={Wand2}
              label="Expand"
              onClick={() => void runQuick("expand")}
            />
            <QuickBtn
              icon={Minimize2}
              label="Simplify"
              onClick={() => void runQuick("simplify")}
            />
            <QuickBtn
              icon={Languages}
              label="Translate"
              onClick={() => void runQuick("translate")}
            />
          </div>
          <p className="text-[11px] text-muted-foreground">
            Enter to run · Esc to remove block · Quick actions use selected text
            or your prompt
          </p>
        </div>
      ) : null}

      {status === "loading" ? (
        <div className="flex min-h-[52px] items-center gap-2 rounded-lg border border-dashed border-border bg-muted/30 px-3 py-2 text-sm text-muted-foreground">
          <Loader2 className="size-4 shrink-0 animate-spin text-primary" />
          <span>Generating</span>
          <LoadingDots />
        </div>
      ) : null}

      {status === "error" && errorMessage ? (
        <div className="rounded-lg border border-destructive/40 bg-destructive/10 px-3 py-2 text-sm text-destructive">
          {errorMessage}
        </div>
      ) : null}

      {status === "done" && resultHtml ? (
        <div className="space-y-3">
          <div
            className="editor-prose max-h-[min(360px,50vh)] overflow-y-auto rounded-lg border border-border bg-card px-3 py-2.5 text-[15px] leading-relaxed [&_p]:my-2 [&_p:first-child]:mt-0 [&_p:last-child]:mb-0"
            // eslint-disable-next-line react/no-danger
            dangerouslySetInnerHTML={{ __html: resultHtml }}
          />
          <div className="flex flex-wrap gap-2">
            <Button type="button" size="sm" onClick={accept}>
              <Check className="mr-1.5 size-3.5" />
              Accept
            </Button>
            <Button type="button" size="sm" variant="secondary" onClick={regenerate}>
              <RefreshCw className="mr-1.5 size-3.5" />
              Regenerate
            </Button>
            <Button type="button" size="sm" variant="ghost" onClick={discard}>
              <Trash2 className="mr-1.5 size-3.5" />
              Discard
            </Button>
          </div>
        </div>
      ) : null}

      {status === "idle" ||
      status === "loading" ||
      status === "error" ? (
        <button
          type="button"
          className="mt-2 flex items-center gap-1 text-[11px] text-muted-foreground hover:text-foreground"
          contentEditable={false}
          onClick={discard}
        >
          <X className="size-3" />
          Cancel
        </button>
      ) : null}
    </NodeViewWrapper>
  );
}

function QuickBtn({
  icon: Icon,
  label,
  onClick,
}: {
  icon: React.ComponentType<{ className?: string }>;
  label: string;
  onClick: () => void;
}) {
  return (
    <Button
      type="button"
      variant="outline"
      size="sm"
      className="h-8 gap-1.5 text-xs"
      onClick={onClick}
    >
      <Icon className="size-3.5 opacity-80" />
      {label}
    </Button>
  );
}

export const AiBlock = Node.create({
  name: "aiBlock",
  group: "block",
  atom: true,
  draggable: true,

  addAttributes() {
    return {
      prompt: { default: "" },
      status: { default: "idle" },
      resultHtml: { default: "" },
      errorMessage: { default: "" },
    };
  },

  parseHTML() {
    return [{ tag: 'div[data-type="ai-block"]' }];
  },

  renderHTML({
    HTMLAttributes,
  }: {
    HTMLAttributes: Record<string, unknown>;
  }) {
    return [
      "div",
      mergeAttributes(HTMLAttributes, { "data-type": "ai-block" }),
    ];
  },

  addNodeView() {
    return ReactNodeViewRenderer(AiBlockView);
  },

  addInputRules() {
    return [
      nodeInputRule({
        find: /(^|\s)\/ai$/,
        type: this.type,
        getAttributes: () => ({
          prompt: "",
          status: "idle",
          resultHtml: "",
          errorMessage: "",
        }),
      }),
    ];
  },
});