Update hooks/useCallAi.ts

hooks/useCallAi.ts

@@ -1,10 +1,7 @@
 import { useState, useRef } from "react";
 import { toast } from "sonner";
+import { MODELS } from "@/lib/providers";
 import { Page } from "@/types";
-import {
-  TITLE_PAGE_START,
-  TITLE_PAGE_END,
-} from "@/lib/prompts";
 
 interface UseCallAiProps {
   onNewPrompt: (prompt: string) => void;
@@ -16,7 +13,6 @@ interface UseCallAiProps {
   pages: Page[];
   isAiWorking: boolean;
   setisAiWorking: React.Dispatch<React.SetStateAction<boolean>>;
-  setAiThinking: React.Dispatch<React.SetStateAction<string | null>>;
 }
 
 export const useCallAi = ({
@@ -25,31 +21,25 @@ export const useCallAi = ({
   onScrollToBottom,
   setPages,
   setCurrentPage,
-  setAiThinking,
+  pages,
   isAiWorking,
   setisAiWorking,
 }: UseCallAiProps) => {
   const audio = useRef<HTMLAudioElement | null>(null);
   const [controller, setController] = useState<AbortController | null>(null);
 
-  const callAiNewProject = async (
-    prompt: string,
-    model: string | undefined,
-    provider: string | undefined,
-    redesignMarkdown?: string
-  ) => {
+  const callAiNewProject = async (prompt: string, model: string | undefined, provider: string | undefined, redesignMarkdown?: string, handleThink?: (think: string) => void, onFinishThink?: () => void) => {
     if (isAiWorking) return;
     if (!redesignMarkdown && !prompt.trim()) return;
-
+    
     setisAiWorking(true);
-    setAiThinking(""); // Clear previous thinking text
-
+    
     const abortController = new AbortController();
     setController(abortController);
-
+    
     try {
       onNewPrompt(prompt);
-
+      
       const request = await fetch("/api/ask-ai", {
         method: "POST",
         body: JSON.stringify({
@@ -65,58 +55,184 @@ export const useCallAi = ({
         signal: abortController.signal,
       });
 
-      if (!request.body) {
-        throw new Error("The response body is empty.");
+      if (request && request.body) {
+        const reader = request.body.getReader();
+        const decoder = new TextDecoder("utf-8");
+        const selectedModel = MODELS.find(
+          (m: { value: string }) => m.value === model
+        );
+        let contentResponse = "";
+
+        const read = async () => {
+          const { done, value } = await reader.read();
+          if (done) {
+            const isJson =
+              contentResponse.trim().startsWith("{") &&
+              contentResponse.trim().endsWith("}");
+            const jsonResponse = isJson ? JSON.parse(contentResponse) : null;
+            
+            if (jsonResponse && !jsonResponse.ok) {
+              if (jsonResponse.openLogin) {
+                // Handle login required
+                return { error: "login_required" };
+              } else if (jsonResponse.openSelectProvider) {
+                // Handle provider selection required
+                return { error: "provider_required", message: jsonResponse.message };
+              } else if (jsonResponse.openProModal) {
+                // Handle pro modal required
+                return { error: "pro_required" };
+              } else {
+                toast.error(jsonResponse.message);
+                setisAiWorking(false);
+                return { error: "api_error", message: jsonResponse.message };
+              }
+            }
+
+            toast.success("AI responded successfully");
+            setisAiWorking(false);
+            
+            if (audio.current) audio.current.play();
+
+            const newPages = formatPages(contentResponse);
+            onSuccess(newPages, prompt);
+
+            return { success: true, pages: newPages };
+          }
+
+          const chunk = decoder.decode(value, { stream: true });
+          contentResponse += chunk;
+          
+          if (selectedModel?.isThinker) {
+            const thinkMatch = contentResponse.match(/<think>[\s\S]*/)?.[0];
+            if (thinkMatch && !contentResponse?.includes("</think>")) {
+              handleThink?.(thinkMatch.replace("<think>", "").trim());
+              return read();
+            }
+          }
+
+          if (contentResponse.includes("</think>")) {
+            onFinishThink?.();
+          }
+
+          formatPages(contentResponse);
+          return read();
+        };
+
+        return await read();
       }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    } catch (error: any) {
+      setisAiWorking(false);
+      toast.error(error.message);
+      if (error.openLogin) {
+        return { error: "login_required" };
+      }
+      return { error: "network_error", message: error.message };
+    }
+  };
+
+  const callAiNewPage = async (prompt: string, model: string | undefined, provider: string | undefined, currentPagePath: string, previousPrompts?: string[]) => {
+    if (isAiWorking) return;
+    if (!prompt.trim()) return;
+    
+    setisAiWorking(true);
+    
+    const abortController = new AbortController();
+    setController(abortController);
+    
+    try {
+      onNewPrompt(prompt);
       
-      const reader = request.body.getReader();
-      const decoder = new TextDecoder("utf-8");
-      let accumulatedResponse = "";
-      let inCodeBlock = false;
-
-      const read = async () => {
-        const { done, value } = await reader.read();
-        if (done) {
-          toast.success("AI has finished generating.");
-          setisAiWorking(false);
-          if (audio.current) audio.current.play();
-          onSuccess(formatPages(accumulatedResponse), prompt);
-          return { success: true };
-        }
+      const request = await fetch("/api/ask-ai", {
+        method: "POST",
+        body: JSON.stringify({
+          prompt,
+          provider,
+          model,
+          pages,
+          previousPrompts,
+        }),
+        headers: {
+          "Content-Type": "application/json",
+          "x-forwarded-for": window.location.hostname,
+        },
+        signal: abortController.signal,
+      });
 
-        const chunk = decoder.decode(value, { stream: true });
-        accumulatedResponse += chunk;
-        
-        // Handle thinking and code separation
-        if (!inCodeBlock) {
-          const codeStartIndex = accumulatedResponse.indexOf(TITLE_PAGE_START);
-          if (codeStartIndex !== -1) {
-            const thinkingText = accumulatedResponse.substring(0, codeStartIndex);
-            setAiThinking(thinkingText);
+      if (request && request.body) {
+        const reader = request.body.getReader();
+        const decoder = new TextDecoder("utf-8");
+        const selectedModel = MODELS.find(
+          (m: { value: string }) => m.value === model
+        );
+        let contentResponse = "";
+
+        const read = async () => {
+          const { done, value } = await reader.read();
+          if (done) {
+            const isJson =
+              contentResponse.trim().startsWith("{") &&
+              contentResponse.trim().endsWith("}");
+            const jsonResponse = isJson ? JSON.parse(contentResponse) : null;
             
-            const codePart = accumulatedResponse.substring(codeStartIndex);
-            formatPages(codePart); // Start rendering code
-            inCodeBlock = true;
-          } else {
-            setAiThinking(accumulatedResponse); // It's still thinking
+            if (jsonResponse && !jsonResponse.ok) {
+              if (jsonResponse.openLogin) {
+                // Handle login required
+                return { error: "login_required" };
+              } else if (jsonResponse.openSelectProvider) {
+                // Handle provider selection required
+                return { error: "provider_required", message: jsonResponse.message };
+              } else if (jsonResponse.openProModal) {
+                // Handle pro modal required
+                return { error: "pro_required" };
+              } else {
+                toast.error(jsonResponse.message);
+                setisAiWorking(false);
+                return { error: "api_error", message: jsonResponse.message };
+              }
+            }
+
+            toast.success("AI responded successfully");
+            setisAiWorking(false);
+            
+            if (selectedModel?.isThinker) {
+              // Reset to default model if using thinker model
+              // Note: You might want to add a callback for this
+            }
+            
+            if (audio.current) audio.current.play();
+
+            const newPage = formatPage(contentResponse, currentPagePath);
+            if (!newPage) { return { error: "api_error", message: "Failed to format page" } }
+            onSuccess([...pages, newPage], prompt);
+
+            return { success: true, pages: [...pages, newPage] };
           }
-        } else {
-          formatPages(accumulatedResponse); // Continue rendering code
-        }
 
-        return read();
-      };
-      
-      return await read();
+          const chunk = decoder.decode(value, { stream: true });
+          contentResponse += chunk;
+          
+          if (selectedModel?.isThinker) {
+            const thinkMatch = contentResponse.match(/<think>[\s\S]*/)?.[0];
+            if (thinkMatch && !contentResponse?.includes("</think>")) {
+              // contentThink += chunk;
+              return read();
+            }
+          }
+
+          formatPage(contentResponse, currentPagePath);
+          return read();
+        };
 
+        return await read();
+      }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     } catch (error: any) {
       setisAiWorking(false);
-      setAiThinking(null); // Clear thinking text on error
-      if (error.name === 'AbortError') {
-        toast.info("AI generation stopped.");
-        return { error: "aborted" };
+      toast.error(error.message);
+      if (error.openLogin) {
+        return { error: "login_required" };
       }
-      toast.error(error.message || "An unexpected error occurred.");
       return { error: "network_error", message: error.message };
     }
   };
@@ -126,7 +242,6 @@ export const useCallAi = ({
     if (!prompt.trim()) return;
     
     setisAiWorking(true);
-    setAiThinking(""); // Clear previous thinking text
     
     const abortController = new AbortController();
     setController(abortController);
@@ -152,39 +267,48 @@ export const useCallAi = ({
         signal: abortController.signal,
       });
 
-      if (request.body) {
-        // Since PUT is not streaming in the original code, we will handle it similarly but add thinking text.
-        // We will assume the non-streaming response for PUT contains the thinking part first.
+      if (request && request.body) {
         const res = await request.json();
         
         if (!request.ok) {
-          setisAiWorking(false);
-          setAiThinking(null);
-          // Handle specific errors like login, pro, etc.
-          toast.error(res.message || "An error occurred during follow-up.");
-          return { error: "api_error", message: res.message };
+          if (res.openLogin) {
+            setisAiWorking(false);
+            return { error: "login_required" };
+          } else if (res.openSelectProvider) {
+            setisAiWorking(false);
+            return { error: "provider_required", message: res.message };
+          } else if (res.openProModal) {
+            setisAiWorking(false);
+            return { error: "pro_required" };
+          } else {
+            toast.error(res.message);
+            setisAiWorking(false);
+            return { error: "api_error", message: res.message };
+          }
         }
-        
-        // Let's assume the response for PUT is not streamed and contains the full text.
-        // We'll have to parse it. For now, we just update pages.
+
         toast.success("AI responded successfully");
         setisAiWorking(false);
+
         setPages(res.pages);
         onSuccess(res.pages, prompt, res.updatedLines);
         
         if (audio.current) audio.current.play();
 
-        return { success: true };
+        return { success: true, html: res.html, updatedLines: res.updatedLines };
       }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     } catch (error: any) {
       setisAiWorking(false);
-      setAiThinking(null);
       toast.error(error.message);
+      if (error.openLogin) {
+        return { error: "login_required" };
+      }
       return { error: "network_error", message: error.message };
     }
   };
 
-
+  // Stop the current AI generation
   const stopController = () => {
     if (controller) {
       controller.abort();
@@ -192,48 +316,144 @@ export const useCallAi = ({
       setisAiWorking(false);
     }
   };
-  
+
   const formatPages = (content: string) => {
     const pages: Page[] = [];
-    const pageRegex = new RegExp(`${TITLE_PAGE_START.replace(/[.*+?^${}()|[]\\]/g, '\\$&')}(.*?)${TITLE_PAGE_END.replace(/[.*+?^${}()|[]\\]/g, '\\$&')}`, "g");
-    
-    let lastIndex = 0;
-    const thinkingText = content.split(pageRegex)[0] || "";
-    if (thinkingText.trim()) {
-        setAiThinking(thinkingText);
+    if (!content.match(/<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/)) {
+      return pages;
     }
-    
-    const matches = Array.from(content.matchAll(pageRegex));
-    matches.forEach((match, i) => {
-        const pagePath = match[1].trim();
-        const startIndex = match.index! + match[0].length;
-        const endIndex = i + 1 < matches.length ? matches[i+1].index! : content.length;
-        const chunk = content.substring(startIndex, endIndex);
-
-        const htmlContent = extractHtmlContent(chunk);
-        if (htmlContent) {
-            pages.push({ path: pagePath, html: htmlContent });
+
+    const cleanedContent = content.replace(
+      /[\s\S]*?<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/,
+      "<<<<<<< START_TITLE $1 >>>>>>> END_TITLE"
+    );
+    const htmlChunks = cleanedContent.split(
+      /<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/
+    );
+    const processedChunks = new Set<number>();
+
+    htmlChunks.forEach((chunk, index) => {
+      if (processedChunks.has(index) || !chunk?.trim()) {
+        return;
+      }
+      const htmlContent = extractHtmlContent(htmlChunks[index + 1]);
+
+      if (htmlContent) {
+        const page: Page = {
+          path: chunk.trim(),
+          html: htmlContent,
+        };
+        pages.push(page);
+
+        if (htmlContent.length > 200) {
+          onScrollToBottom();
         }
+
+        processedChunks.add(index);
+        processedChunks.add(index + 1);
+      }
     });
-    
     if (pages.length > 0) {
       setPages(pages);
       const lastPagePath = pages[pages.length - 1]?.path;
       setCurrentPage(lastPagePath || "index.html");
-      onScrollToBottom();
     }
+
     return pages;
   };
-  
+
+  const formatPage = (content: string, currentPagePath: string) => {
+    if (!content.match(/<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/)) {
+      return null;
+    }
+
+    const cleanedContent = content.replace(
+      /[\s\S]*?<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/,
+      "<<<<<<< START_TITLE $1 >>>>>>> END_TITLE"
+    );
+
+    const htmlChunks = cleanedContent.split(
+      /<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/
+    )?.filter(Boolean);
+
+    const pagePath = htmlChunks[0]?.trim() || "";
+    const htmlContent = extractHtmlContent(htmlChunks[1]);
+
+    if (!pagePath || !htmlContent) {
+      return null;
+    }
+
+    const page: Page = {
+      path: pagePath,
+      html: htmlContent,
+    };
+
+    setPages(prevPages => {
+      const existingPageIndex = prevPages.findIndex(p => p.path === currentPagePath || p.path === pagePath);
+      
+      if (existingPageIndex !== -1) {
+        const updatedPages = [...prevPages];
+        updatedPages[existingPageIndex] = page;
+        return updatedPages;
+      } else {
+        return [...prevPages, page];
+      }
+    });
+
+    setCurrentPage(pagePath);
+
+    if (htmlContent.length > 200) {
+      onScrollToBottom();
+    }
+
+    return page;
+  };
+
+  // Helper function to extract and clean HTML content
   const extractHtmlContent = (chunk: string): string => {
-    const htmlMatch = chunk.match(/``````/);
-    return htmlMatch ? htmlMatch[1].trim() : chunk.trim();
+    if (!chunk) return "";
+
+    // Extract HTML content
+    const htmlMatch = chunk.trim().match(/<!DOCTYPE html>[\s\S]*/);
+    if (!htmlMatch) return "";
+
+    let htmlContent = htmlMatch[0];
+
+    // Ensure proper HTML structure
+    htmlContent = ensureCompleteHtml(htmlContent);
+
+    // Remove markdown code blocks if present
+    htmlContent = htmlContent.replace(/```/g, "");
+
+    return htmlContent;
+  };
+
+  // Helper function to ensure HTML has complete structure
+  const ensureCompleteHtml = (html: string): string => {
+    let completeHtml = html;
+
+    // Add missing head closing tag
+    if (completeHtml.includes("<head>") && !completeHtml.includes("</head>")) {
+      completeHtml += "\n</head>";
+    }
+
+    // Add missing body closing tag
+    if (completeHtml.includes("<body") && !completeHtml.includes("</body>")) {
+      completeHtml += "\n</body>";
+    }
+
+    // Add missing html closing tag
+    if (!completeHtml.includes("</html>")) {
+      completeHtml += "\n</html>";
+    }
+
+    return completeHtml;
   };
 
   return {
     callAiNewProject,
     callAiFollowUp,
-    callAiNewPage: callAiNewProject, // Rerouting for simplicity, can be expanded
+    callAiNewPage,
     stopController,
     controller,
     audio,