Merge pull request #2 from novitalabs/refine-render

src/components/preview.tsx

@@ -1,10 +1,11 @@
 "use client"
 
 import { useState, forwardRef, useImperativeHandle, useEffect, useRef } from "react"
-import { DEFAULT_HTML, PARTIAL_GENERATING_HTML_TEMPLATE } from "@/lib/constants"
+import { DEFAULT_HTML } from "@/lib/constants"
 import { PreviewRef, Version } from "@/lib/types"
 import { MinimizeIcon, MaximizeIcon, DownloadIcon } from "./ui/fullscreen-icons"
 import { useModel } from "@/lib/contexts/model-context"
+import { Loader2 } from "lucide-react"
 import { cn } from "@/lib/utils"
 
 interface PreviewProps {
@@ -27,6 +28,8 @@ export const Preview = forwardRef<PreviewRef, PreviewProps>(function Preview(
   const [error, setError] = useState<string | null>(null);
   const [showAuthError, setShowAuthError] = useState(false);
   const { selectedModelId } = useModel();
+  const renderCount = useRef(0);
+  const headUpdated = useRef(false);
 
   // Update html when initialHtml changes
   useEffect(() => {
@@ -55,26 +58,22 @@ export const Preview = forwardRef<PreviewRef, PreviewProps>(function Preview(
   const partialUpdate = (htmlStr: string) => {
     const parser = new DOMParser();
     const partialDoc = parser.parseFromString(htmlStr, 'text/html');
-
     const iframe = document.querySelector('iframe');
     if (!iframe || !iframe.contentDocument) return;
-
-    // Check if container exists, if not initialize with previewHtml directly
-    const iframeContainer = iframe.contentDocument.querySelector('#container');
-    if (!iframeContainer) {
-      const initDoc = parser.parseFromString(PARTIAL_GENERATING_HTML_TEMPLATE, 'text/html');
-      if (iframe.contentDocument.documentElement) {
-        iframe.contentDocument.documentElement.innerHTML = initDoc.documentElement.innerHTML;
-      }
-      setHtml(PARTIAL_GENERATING_HTML_TEMPLATE);
-      return;
+    
+    const iframeContainer = iframe.contentDocument;
+    if (iframeContainer?.body && iframeContainer) {
+      iframeContainer.body.innerHTML = partialDoc.body?.innerHTML;
     }
-
-    // Extract the container content
-    const containerContent = partialDoc.body?.innerHTML;
-    if (containerContent && iframeContainer) {
-      iframeContainer.innerHTML = containerContent;
+    if (renderCount.current % 10 === 0 && !headUpdated.current) {
+      setHtml(htmlStr);
+      if (htmlStr.includes('</head>')) {
+        setTimeout(() => {
+          headUpdated.current = true;
+        }, 1000);
+      }
     }
+    renderCount.current++;
   }
 
   const downloadHtml = () => {
@@ -106,6 +105,8 @@ export const Preview = forwardRef<PreviewRef, PreviewProps>(function Preview(
 
   const generateCode = async (prompt: string, colors: string[] = [], previousPrompt?: string) => {
     setLoading(true);
+    renderCount.current = 0;
+    headUpdated.current = false;
     if (onLoadingChange) {
       onLoadingChange(true);
     }
@@ -191,8 +192,8 @@ export const Preview = forwardRef<PreviewRef, PreviewProps>(function Preview(
                 }, 50);
               }
             }
-            break;
             setIsPartialGenerating(false);
+            break;
           } else {
             setIsPartialGenerating(true);
           }
@@ -252,7 +253,17 @@ export const Preview = forwardRef<PreviewRef, PreviewProps>(function Preview(
 
   return (
     <div className={`${isFullscreen ? 'fixed inset-0 z-10 bg-novita-dark' : 'h-full'} p-4`}>
-      <div className="bg-white text-black rounded-md h-full overflow-hidden relative isolation-auto">
+      { isPartialGenerating && (
+        <div className="w-full bg-slate-50 border-b border-slate-200 py-2 px-4">
+          <div className="container mx-auto flex items-center justify-center">
+          <div className="flex items-center space-x-2 text-slate-700">
+            <Loader2 className="h-4 w-4 animate-spin" />
+            <span className="text-sm font-medium">building...</span>
+          </div>
+          </div>
+        </div>
+      )}
+      <div className="bg-white text-black h-full overflow-hidden relative isolation-auto">
         <div className="absolute top-3 right-3 flex gap-2 z-[100]">
           <button
             onClick={downloadHtml}

src/lib/constants.ts

@@ -88,25 +88,3 @@ export const DEFAULT_HTML = `<!DOCTYPE html>
   </div>
 </body>
 </html>`; 
-
-export const PARTIAL_GENERATING_HTML_TEMPLATE = `
-<!DOCTYPE html>
-<html lang="en">
-  <head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>Novita AnySite</title>
-    <script src="https://cdn.tailwindcss.com"></script>
-  </head>
-  <body>
-    <div id="container" class="py-10">
-    </div>
-    <div class="fixed inset-0 bg-gray-300 bg-opacity-25 flex items-center justify-center z-50">
-      <div class="flex flex-col items-center">
-        <div class="animate-spin rounded-full h-8 w-8 border-b-2 border-gray-600"></div>
-        <p class="mt-3 text-gray-600 text-sm">Loading...</p>
-      </div>
-    </div>
-  </body>
-</html>
-`;