components/editor/ask-ai/index.tsx

"use client";
/* eslint-disable @typescript-eslint/no-explicit-any */
import { useState, useRef, useMemo } from "react";
import classNames from "classnames";
import { toast } from "sonner";
import { useLocalStorage, useUpdateEffect } from "react-use";
import { ArrowUp, ChevronDown, Crosshair } from "lucide-react";
import { FaStopCircle } from "react-icons/fa";

import ProModal from "@/components/pro-modal";
import { Button } from "@/components/ui/button";
import { MODELS } from "@/lib/providers";
import { HtmlHistory } from "@/types";
import { InviteFriends } from "@/components/invite-friends";
import { Settings } from "@/components/editor/ask-ai/settings";
import { LoginModal } from "@/components/login-modal";
import { ReImagine } from "@/components/editor/ask-ai/re-imagine";
import Loading from "@/components/loading";
import { Checkbox } from "@/components/ui/checkbox";
import { Tooltip, TooltipTrigger } from "@/components/ui/tooltip";
import { TooltipContent } from "@radix-ui/react-tooltip";
import { SelectedHtmlElement } from "./selected-html-element";
import { FollowUpTooltip } from "./follow-up-tooltip";
import { isTheSameHtml } from "@/lib/compare-html-diff";

export function AskAI({
  html,
  setHtml,
  onScrollToBottom,
  isAiWorking,
  setisAiWorking,
  isEditableModeEnabled = false,
  selectedElement,
  setSelectedElement,
  setIsEditableModeEnabled,
  onNewPrompt,
  onSuccess,
}: {
  html: string;
  setHtml: (html: string) => void;
  onScrollToBottom: () => void;
  isAiWorking: boolean;
  onNewPrompt: (prompt: string) => void;
  htmlHistory?: HtmlHistory[];
  setisAiWorking: React.Dispatch<React.SetStateAction<boolean>>;
  onSuccess: (h: string, p: string, n?: number[][]) => void;
  isEditableModeEnabled: boolean;
  setIsEditableModeEnabled: React.Dispatch<React.SetStateAction<boolean>>;
  selectedElement?: HTMLElement | null;
  setSelectedElement: React.Dispatch<React.SetStateAction<HTMLElement | null>>;
}) {
  const [openrouterApiKey] = useLocalStorage<string>("openrouter-api-key", "");
  
  // Check if OpenRouter API key is properly loaded from localStorage
  const isOpenRouterApiKeyValid = useMemo(() => {
    return openrouterApiKey && openrouterApiKey.trim().length > 0;
  }, [openrouterApiKey]);
  const refThink = useRef<HTMLDivElement | null>(null);
  const audio = useRef<HTMLAudioElement | null>(null);

  const [open, setOpen] = useState(false);
  const [prompt, setPrompt] = useState("");
  const [hasAsked, setHasAsked] = useState(false);
  const [previousPrompt, setPreviousPrompt] = useState("");
  const [provider, setProvider] = useLocalStorage("provider", "auto");
  const [model, setModel] = useLocalStorage("model", MODELS[0].value);
  const [openProvider, setOpenProvider] = useState(false);
  const [providerError, setProviderError] = useState("");
  const [openProModal, setOpenProModal] = useState(false);
  const [think, setThink] = useState<string | undefined>(undefined);
  const [openThink, setOpenThink] = useState(false);
  const [isThinking, setIsThinking] = useState(true);
  const [controller, setController] = useState<AbortController | null>(null);
  const [isFollowUp, setIsFollowUp] = useState(true);

  const callAi = async (redesignMarkdown?: string) => {
    if (isAiWorking) return;
    if (!redesignMarkdown && !prompt.trim()) return;
    setisAiWorking(true);
    setProviderError("");
    setThink("");
    setOpenThink(false);
    setIsThinking(true);

    let contentResponse = "";
    let thinkResponse = "";
    let lastRenderTime = 0;

    const abortController = new AbortController();
    setController(abortController);
    try {
      onNewPrompt(prompt);
      if (isFollowUp && !redesignMarkdown && !isSameHtml) {
        console.log('🔄 DIFF-PATCH MODE: Frontend sending PUT request', {
          isFollowUp,
          redesignMarkdown: !!redesignMarkdown,
          isSameHtml,
          hasSelectedElement: !!selectedElement,
          htmlLength: html.length,
          provider,
          model
        });
        
        const selectedElementHtml = selectedElement
          ? selectedElement.outerHTML
          : "";
          
        console.log('📤 PUT request payload:', {
          method: 'PUT',
          promptLength: prompt.length,
          hasSelectedElement: !!selectedElementHtml,
          selectedElementLength: selectedElementHtml.length,
          totalHtmlLength: html.length,
          hasOpenRouterApiKey: isOpenRouterApiKeyValid
        });
        
        // Validate OpenRouter API key if using OpenRouter
        if (provider === "openrouter" && !isOpenRouterApiKeyValid) {
          toast.error("OpenRouter API key is required for this model");
          setOpenProvider(true);
          setProviderError("OpenRouter API key is required for this model");
          setisAiWorking(false);
          return;
        }
        const request = await fetch("/api/ask-ai", {
          method: "PUT",
          body: JSON.stringify({
            prompt,
            provider,
            previousPrompt,
            model,
            html,
            selectedElementHtml,
            openrouterApiKey: openrouterApiKey?.trim() || undefined,
          }),
          headers: {
            "Content-Type": "application/json",
            "x-forwarded-for": window.location.hostname,
          },
          signal: abortController.signal,
        });
        if (request && request.body) {
          const res = await request.json();
          if (!request.ok) {
            if (res.openLogin) {
              setOpen(true);
            } else if (res.openSelectProvider) {
              setOpenProvider(true);
              setProviderError(res.message);
            } else if (res.openProModal) {
              setOpenProModal(true);
            } else {
              toast.error(res.message);
            }
            setisAiWorking(false);
            return;
          }
          setHtml(res.html);
          toast.success("AI responded successfully");
          setPreviousPrompt(prompt);
          setPrompt("");
          setisAiWorking(false);
          onSuccess(res.html, prompt, res.updatedLines);
          if (audio.current) audio.current.play();
        }
      } else {
        console.log('🆕 INITIAL REQUEST: Frontend sending POST request', {
          isFollowUp,
          redesignMarkdown: !!redesignMarkdown,
          isSameHtml,
          reason: !isFollowUp ? 'isFollowUp=false' : redesignMarkdown ? 'has redesignMarkdown' : 'isSameHtml=true',
          provider,
          model,
          htmlLength: html.length
        });
        
        // Validate OpenRouter API key if using OpenRouter
        if (provider === "openrouter" && !isOpenRouterApiKeyValid) {
          toast.error("OpenRouter API key is required for this model");
          setOpenProvider(true);
          setProviderError("OpenRouter API key is required for this model");
          setisAiWorking(false);
          return;
        }
        
        const request = await fetch("/api/ask-ai", {
          method: "POST",
          body: JSON.stringify({
            prompt,
            provider,
            model,
            html: isSameHtml ? "" : html,
            redesignMarkdown,
            openrouterApiKey: openrouterApiKey?.trim() || undefined,
          }),
          headers: {
            "Content-Type": "application/json",
            "x-forwarded-for": window.location.hostname,
          },
          signal: abortController.signal,
        });
        if (request && request.body) {
          // ROBUST ERROR HANDLING: Handle both JSON and HTML error responses
          if (!request.ok) {
            console.error('❌ Request failed:', {
              status: request.status,
              statusText: request.statusText,
              url: request.url,
              headers: Object.fromEntries(request.headers.entries())
            });
            
            try {
              // Get the full response as text first
              const responseText = await request.text();
              console.log('📄 Raw error response:', {
                length: responseText.length,
                contentType: request.headers.get('content-type'),
                preview: responseText.substring(0, 500) + (responseText.length > 500 ? '...' : ''),
                isHtml: responseText.trim().toLowerCase().startsWith('<!doctype html') || responseText.trim().toLowerCase().startsWith('<html')
              });
              
              let res;
              
              // Check if response is HTML (common for server errors)
              if (responseText.trim().toLowerCase().startsWith('<!doctype html') || 
                  responseText.trim().toLowerCase().startsWith('<html') ||
                  responseText.includes('<title>') || 
                  responseText.includes('Internal Server Error')) {
                
                console.error('❌ HTML error page received instead of JSON');
                
                // Extract error info from HTML if possible
                let errorMessage = 'Server error occurred';
                const titleMatch = responseText.match(/<title[^>]*>([^<]+)<\/title>/i);
                if (titleMatch && titleMatch[1]) {
                  errorMessage = titleMatch[1].trim();
                }
                
                // Look for common error patterns
                if (responseText.includes('Internal Server Error')) {
                  errorMessage = 'Internal Server Error - Please try again';
                } else if (responseText.includes('Service Unavailable')) {
                  errorMessage = 'Service temporarily unavailable';
                } else if (responseText.includes('Bad Gateway')) {
                  errorMessage = 'Gateway error - Please try again';
                }
                
                res = {
                  ok: false,
                  error: errorMessage,
                  message: errorMessage,
                  isHtmlError: true
                };
              } else {
                // Try to parse as JSON
                try {
                  res = JSON.parse(responseText);
                  console.error('❌ Error response (JSON):', res);
                } catch (jsonError) {
                  console.error('❌ Failed to parse JSON, treating as text:', {
                    jsonError: (jsonError as Error).message || String(jsonError),
                    responseLength: responseText.length
                  });
                  
                  // Create error object from text response
                  let errorMessage = responseText.trim();
                  
                  // Clean up common error patterns
                  if (errorMessage.length > 200) {
                    errorMessage = errorMessage.substring(0, 200) + '...';
                  }
                  
                  if (!errorMessage) {
                    errorMessage = `HTTP ${request.status}: ${request.statusText}`;
                  }
                  
                  res = {
                    ok: false,
                    error: errorMessage,
                    message: errorMessage,
                    isTextError: true
                  };
                }
              }
              
              // Handle different error types
              if (res.openLogin) {
                setOpen(true);
              } else if (res.openSelectProvider) {
                setOpenProvider(true);
                setProviderError(res.message || res.error || 'Provider error');
              } else if (res.openProModal) {
                setOpenProModal(true);
              } else {
                const errorMsg = res.message || res.error || 'Unknown error occurred';
                console.error('❌ Showing error to user:', errorMsg);
                toast.error(errorMsg);
              }
            } catch (parseError) {
              console.error('❌ Complete failure to parse error response:', {
                parseError: (parseError as Error).message || String(parseError),
                status: request.status,
                statusText: request.statusText
              });
              toast.error(`Request failed: ${request.status} ${request.statusText}`);
            }
            setisAiWorking(false);
            return;
          }
          
          console.log('✅ Request successful, processing stream...', {
            status: request.status,
            headers: Object.fromEntries(request.headers.entries())
          });
          
          const reader = request.body.getReader();
          const decoder = new TextDecoder("utf-8");
          const selectedModel = MODELS.find(
            (m: { value: string }) => m.value === model
          );
          let contentThink: string | undefined = undefined;
          let chunkCount = 0;
          
          const read = async () => {
            const { done, value } = await reader.read();
            if (done) {
              console.log('✅ Stream completed:', {
                totalChunks: chunkCount,
                finalContentLength: contentResponse.length,
                hasDoctype: contentResponse.includes('<!DOCTYPE html>'),
                hasHtmlEnd: contentResponse.includes('</html>')
              });
              
              toast.success("AI responded successfully");
              setPreviousPrompt(prompt);
              setPrompt("");
              setisAiWorking(false);
              setHasAsked(true);
              // Note: Removed automatic model reset to preserve user's selection
              if (audio.current) audio.current.play();

              // Now we have the complete HTML including </html>, so set it to be sure
              const finalDoc = contentResponse.match(
                /<!DOCTYPE html>[\s\S]*<\/html>/
              )?.[0];
              if (finalDoc) {
                setHtml(finalDoc);
              } else {
                console.warn('⚠️ No complete HTML document found in response');
              }
              onSuccess(finalDoc ?? contentResponse, prompt);

              return;
            }

            chunkCount++;
            const chunk = decoder.decode(value, { stream: true });
            
            console.log(`📦 Frontend chunk ${chunkCount}:`, {
              chunkLength: chunk.length,
              hasJsonError: chunk.trim().startsWith('{'),
              preview: chunk.substring(0, 200) + (chunk.length > 200 ? '...' : '')
            });
            
            // Check if this chunk looks like a JSON error response
            if (chunk.trim().startsWith('{') && chunk.trim().endsWith('}')) {
              try {
                const res = JSON.parse(chunk);
                // Only treat as error if it has error indicators
                if (res.ok === false || res.error || res.openLogin || res.openSelectProvider || res.openProModal) {
                  console.error('❌ Error response received:', res);
                  
                  if (res.openLogin) {
                    setOpen(true);
                  } else if (res.openSelectProvider) {
                    setOpenProvider(true);
                    setProviderError(res.message || res.error);
                  } else if (res.openProModal) {
                    setOpenProModal(true);
                  } else {
                    toast.error(res.message || res.error || 'Unknown error occurred');
                  }
                  setisAiWorking(false);
                  return;
                }
              } catch {
                // If it looks like JSON but can't be parsed, treat as content
                console.log('⚠️ Chunk looks like JSON but failed to parse, treating as content');
              }
            }
            
            // Treat as normal content
            thinkResponse += chunk;
            if (selectedModel?.isThinker) {
              const thinkMatch = thinkResponse.match(/<think>[\s\S]*/)?.[0];
              if (thinkMatch && !thinkResponse?.includes("</think>")) {
                if ((contentThink?.length ?? 0) < 3) {
                  setOpenThink(true);
                }
                setThink(thinkMatch.replace("<think>", "").trim());
                contentThink += chunk;
                return read();
              }
            }

            contentResponse += chunk;

            const newHtml = contentResponse.match(
              /<!DOCTYPE html>[\s\S]*/
            )?.[0];
            if (newHtml) {
              setIsThinking(false);
              let partialDoc = newHtml;
              if (
                partialDoc.includes("<head>") &&
                !partialDoc.includes("</head>")
              ) {
                partialDoc += "\n</head>";
              }
              if (
                partialDoc.includes("<body") &&
                !partialDoc.includes("</body>")
              ) {
                partialDoc += "\n</body>";
              }
              if (!partialDoc.includes("</html>")) {
                partialDoc += "\n</html>";
              }

              // Ultra-conservative throttling to eliminate all flashing
              const now = Date.now();
              const isCompleteDocument = partialDoc.includes('</html>');
              const htmlLength = partialDoc.length;
              const timeSinceLastUpdate = lastRenderTime === 0 ? 0 : now - lastRenderTime;
              
              // ZERO-FLASH throttling - more responsive since we have seamless injection
              let throttleDelay = 1000; // Default: responsive for zero-flash system
              if (isCompleteDocument) {
                throttleDelay = 200; // Fast completion
              } else if (htmlLength < 1000) {
                throttleDelay = 1500; // Moderate for small content
              } else if (htmlLength > 8000) {
                throttleDelay = 800; // Faster for large content since no flash
              } else if (timeSinceLastUpdate > 3000) {
                throttleDelay = 500; // Faster if it's been long
              }
              
              // Only update if enough time has passed OR it's completion
              const shouldUpdate = (now - lastRenderTime > throttleDelay) || isCompleteDocument;
              
              if (shouldUpdate) {
                setHtml(partialDoc);
                lastRenderTime = now;
                console.log('� Frontend: Zero-flash HTML update', {
                  htmlLength: partialDoc.length,
                  isComplete: isCompleteDocument,
                  chunkNumber: chunkCount,
                  throttleDelay,
                  timeSinceLastUpdate,
                  updateReason: isCompleteDocument ? 'completion' : 'time-based'
                });
              } else {
                console.log('⏳ Frontend: Throttling update for optimal UX', {
                  htmlLength: partialDoc.length,
                  timeSinceLastUpdate,
                  requiredDelay: throttleDelay,
                  chunkNumber: chunkCount
                });
              }

              if (partialDoc.length > 200) {
                onScrollToBottom();
              }
            } else {
              // Still thinking/no HTML yet
              console.log('🤔 Still waiting for HTML to start...');
            }
            read();
          };

          read();
        }
      }
    } catch (error: any) {
      console.error('❌ Frontend error in callAi:', {
        errorMessage: error.message,
        errorStack: error.stack,
        model,
        provider,
        isAborted: error.name === 'AbortError'
      });
      
      setisAiWorking(false);
      
      if (error.name === 'AbortError') {
        toast.error('Request was cancelled');
      } else {
        toast.error(error.message || 'An unexpected error occurred');
      }
      
      if (error.openLogin) {
        setOpen(true);
      }
    }
  };

  const stopController = () => {
    if (controller) {
      controller.abort();
      setController(null);
      setisAiWorking(false);
      setThink("");
      setOpenThink(false);
      setIsThinking(false);
    }
  };

  useUpdateEffect(() => {
    if (refThink.current) {
      refThink.current.scrollTop = refThink.current.scrollHeight;
    }
  }, [think]);

  useUpdateEffect(() => {
    if (!isThinking) {
      setOpenThink(false);
    }
  }, [isThinking]);

  const isSameHtml = useMemo(() => {
    return isTheSameHtml(html);
  }, [html]);

  return (
    <div className="px-3">
      <div className="relative bg-neutral-800 border border-neutral-700 rounded-2xl ring-[4px] focus-within:ring-neutral-500/30 focus-within:border-neutral-600 ring-transparent z-10 w-full group">
        {think && (
          <div className="w-full border-b border-neutral-700 relative overflow-hidden">
            <header
              className="flex items-center justify-between px-5 py-2.5 group hover:bg-neutral-600/20 transition-colors duration-200 cursor-pointer"
              onClick={() => {
                setOpenThink(!openThink);
              }}
            >
              <p className="text-sm font-medium text-neutral-300 group-hover:text-neutral-200 transition-colors duration-200">
                {isThinking ? "DeepSite is thinking..." : "DeepSite's plan"}
              </p>
              <ChevronDown
                className={classNames(
                  "size-4 text-neutral-400 group-hover:text-neutral-300 transition-all duration-200",
                  {
                    "rotate-180": openThink,
                  }
                )}
              />
            </header>
            <main
              ref={refThink}
              className={classNames(
                "overflow-y-auto transition-all duration-200 ease-in-out",
                {
                  "max-h-[0px]": !openThink,
                  "min-h-[250px] max-h-[250px] border-t border-neutral-700":
                    openThink,
                }
              )}
            >
              <p className="text-[13px] text-neutral-400 whitespace-pre-line px-5 pb-4 pt-3">
                {think}
              </p>
            </main>
          </div>
        )}
        {selectedElement && (
          <div className="px-4 pt-3">
            <SelectedHtmlElement
              element={selectedElement}
              isAiWorking={isAiWorking}
              onDelete={() => setSelectedElement(null)}
            />
          </div>
        )}
        <div className="w-full relative flex items-center justify-between">
          {isAiWorking && (
            <div className="absolute bg-neutral-800 rounded-lg bottom-0 left-4 w-[calc(100%-30px)] h-full z-1 flex items-center justify-between max-lg:text-sm">
              <div className="flex items-center justify-start gap-2">
                <Loading overlay={false} className="!size-4" />
                <p className="text-neutral-400 text-sm">
                  AI is {isThinking ? "thinking" : "coding"}...{" "}
                </p>
              </div>
              <div
                className="text-xs text-neutral-400 px-1 py-0.5 rounded-md border border-neutral-600 flex items-center justify-center gap-1.5 bg-neutral-800 hover:brightness-110 transition-all duration-200 cursor-pointer"
                onClick={stopController}
              >
                <FaStopCircle />
                Stop generation
              </div>
            </div>
          )}
          <input
            type="text"
            disabled={isAiWorking}
            className={classNames(
              "w-full bg-transparent text-sm outline-none text-white placeholder:text-neutral-400 p-4",
              {
                "!pt-2.5": selectedElement && !isAiWorking,
              }
            )}
            placeholder={
              selectedElement
                ? `Ask DeepSite about ${selectedElement.tagName.toLowerCase()}...`
                : hasAsked
                ? "Ask DeepSite for edits"
                : "Ask DeepSite anything..."
            }
            value={prompt}
            onChange={(e) => setPrompt(e.target.value)}
            onKeyDown={(e) => {
              if (e.key === "Enter" && !e.shiftKey) {
                callAi();
              }
            }}
          />
        </div>
        <div className="flex items-center justify-between gap-2 px-4 pb-3">
          <div className="flex-1 flex items-center justify-start gap-1.5">
            <ReImagine onRedesign={(md) => callAi(md)} />
            {!isSameHtml && (
              <Tooltip>
                <TooltipTrigger asChild>
                  <Button
                    size="xs"
                    variant={isEditableModeEnabled ? "default" : "outline"}
                    onClick={() => {
                      console.log("🎯 Edit button clicked:", {
                        currentState: isEditableModeEnabled,
                        newState: !isEditableModeEnabled,
                        hasSetFunction: !!setIsEditableModeEnabled
                      });
                      setIsEditableModeEnabled?.(!isEditableModeEnabled);
                    }}
                    className={classNames("h-[28px]", {
                      "!text-neutral-400 hover:!text-neutral-200 !border-neutral-600 !hover:!border-neutral-500":
                        !isEditableModeEnabled,
                    })}
                  >
                    <Crosshair className="size-4" />
                    Edit
                  </Button>
                </TooltipTrigger>
                <TooltipContent
                  align="start"
                  className="bg-neutral-950 text-xs text-neutral-200 py-1 px-2 rounded-md -translate-y-0.5"
                >
                  Select an element on the page to ask DeepSite edit it
                  directly.
                </TooltipContent>
              </Tooltip>
            )}
            <InviteFriends />
          </div>
          <div className="flex items-center justify-end gap-2">
            <Settings
              provider={provider as string}
              model={model as string}
              onChange={setProvider}
              onModelChange={setModel}
              open={openProvider}
              error={providerError}
              isFollowUp={!isSameHtml}
              onClose={setOpenProvider}
            />
            <Button
              size="iconXs"
              disabled={isAiWorking || !prompt.trim()}
              onClick={() => callAi()}
            >
              <ArrowUp className="size-4" />
            </Button>
          </div>
        </div>
        <LoginModal open={open} onClose={() => setOpen(false)} html={html} />
        <ProModal
          html={html}
          open={openProModal}
          onClose={() => setOpenProModal(false)}
        />
        {!isSameHtml && (
          <div className="absolute top-0 right-0 -translate-y-[calc(100%+8px)] select-none text-xs text-neutral-400 flex items-center justify-center gap-2 bg-neutral-800 border border-neutral-700 rounded-md p-1 pr-2.5">
            <label
              htmlFor="diff-patch-checkbox"
              className="flex items-center gap-1.5 cursor-pointer"
            >
              <Checkbox
                id="diff-patch-checkbox"
                checked={isFollowUp}
                onCheckedChange={(e: boolean) => {
                  setIsFollowUp(e === true);
                }}
              />
              Diff-Patch Update
            </label>
            <FollowUpTooltip />
          </div>
        )}
      </div>
      <audio ref={audio} id="audio" className="hidden">
        <source src="/success.mp3" type="audio/mpeg" />
        Your browser does not support the audio element.
      </audio>
    </div>
  );
}