src/components/ask-ai/ask-ai.tsx

import { useState } from "react";
import { RiSparkling2Fill } from "react-icons/ri";
import { GrSend } from "react-icons/gr";
import classNames from "classnames";
import { toast } from "react-toastify";
// Removed monaco editor import

import Login from "../login/login";
import { defaultHTML } from "../../utils/consts";
import SuccessSound from "./../../assets/success.mp3";

function AskAI({
  html, // Current full HTML content (used for initial request and context)
  setHtml, // Used for updates (both full and diff-based)
  onScrollToBottom, // Used for full updates
  isAiWorking,
  setisAiWorking,
}: {
  html: string;
  setHtml: (html: string) => void;
  onScrollToBottom: () => void;
  isAiWorking: boolean;
  setisAiWorking: React.Dispatch<React.SetStateAction<boolean>>;
  // Removed editorRef prop
}) {
  const [open, setOpen] = useState(false);
  const [prompt, setPrompt] = useState("");
  const [hasAsked, setHasAsked] = useState(false);
  const [previousPrompt, setPreviousPrompt] = useState("");
  const [streamProgress, setStreamProgress] = useState(0); // Track streaming progress
  const [responseMode, setResponseMode] = useState<"full" | "diff" | null>(null); // Track response type
  const audio = new Audio(SuccessSound);
  audio.volume = 0.5;

  // Removed client-side diff parsing/applying logic
  
  // --- Main AI Call Logic ---
  const callAi = async () => {
    if (isAiWorking || !prompt.trim()) return;
    const originalHtml = html; // Store the HTML state at the start of the request
    setisAiWorking(true);
    setStreamProgress(0); // Reset progress indicator
    setResponseMode(null); // Reset response mode

    let fullContentResponse = ""; // Used for full HTML mode
    let accumulatedDiffResponse = ""; // Used for diff mode
    let lastRenderTime = 0; // For throttling full HTML updates
    let totalChunksReceived = 0;

    try {
      const request = await fetch("/api/ask-ai", {
        method: "POST",
        body: JSON.stringify({
          prompt,
          ...(html === defaultHTML ? {} : { html }),
          ...(previousPrompt ? { previousPrompt } : {}),
        }),
        headers: {
          "Content-Type": "application/json",
        },
      });
      if (request && request.body) {
        if (!request.ok) {
          const res = await request.json();
          if (res.openLogin) {
            setOpen(true);
          } else {
            // don't show toast if it's a login error
            toast.error(res.message);
          }
          setisAiWorking(false);
          setStreamProgress(0);
          return;
        }

        const responseType = request.headers.get("X-Response-Type") || "full"; // Default to full if header missing
        setResponseMode(responseType as "full" | "diff");
        console.log(`[AI Response] Type: ${responseType}`);

        const reader = request.body.getReader();
        const decoder = new TextDecoder("utf-8");
        
        console.log("[AI Request] Starting to read response stream");
        
        // Process all chunks until done
        while (true) {
          try {
            const { done, value } = await reader.read();
            if (done) {
              console.log("[AI Response] Stream finished successfully.");
              console.log("[AI Response] Accumulated full content length:", fullContentResponse.length);
              console.log("[AI Response] Accumulated diff content length:", accumulatedDiffResponse.length);
              setStreamProgress(100); // Set progress to 100% when done
              
              // Play success animation
              document.body.classList.add('ai-success');
              setTimeout(() => {
                document.body.classList.remove('ai-success');
              }, 1000);

              // --- Post-stream processing ---
              if (responseType === 'diff') {
                // Apply diffs server-side
                try {
                  console.log("[Diff Apply] Sending original HTML and AI diff response to server...");
                  const applyRequest = await fetch("/api/apply-diffs", {
                    method: "POST",
                    headers: { "Content-Type": "application/json" },
                    body: JSON.stringify({
                      originalHtml: originalHtml, // Send the HTML from the start of the request
                      aiResponseContent: accumulatedDiffResponse,
                    }),
                  });

                  if (!applyRequest.ok) {
                    const errorData = await applyRequest.json();
                    throw new Error(errorData.message || `Server failed to apply diffs (status ${applyRequest.status})`);
                  }

                  const patchedHtml = await applyRequest.text();
                  console.log("[Diff Apply] Received patched HTML from server.");
                  setHtml(patchedHtml); // Update editor with the final result
                  toast.success("AI changes applied");

                } catch (applyError: unknown) {
                  console.error("Error applying diffs server-side:", applyError);
                  const errorMessage = applyError instanceof Error 
                    ? applyError.message 
                    : 'Unknown error applying changes';
                  toast.error(`Failed to apply AI changes: ${errorMessage}`);
                  // Optionally revert to originalHtml? Or leave the editor as is?
                  // setHtml(originalHtml); // Uncomment to revert on failure
                }

              } else {
                // Final update for full HTML mode
                const finalDoc = fullContentResponse.match(/<!DOCTYPE html>[\s\S]*<\/html>/)?.[0];
                if (finalDoc) {
                  console.log("[AI Response] Found complete HTML document");
                  setHtml(finalDoc); // Ensure final complete HTML is set
                } else if (fullContentResponse.includes("<html") && fullContentResponse.includes("</html>")) {
                  // Try to match HTML without DOCTYPE
                  const htmlMatch = fullContentResponse.match(/<html[\s\S]*<\/html>/);
                  if (htmlMatch) {
                    console.log("[AI Response] Found HTML without DOCTYPE, adding DOCTYPE");
                    setHtml(`<!DOCTYPE html>\n${htmlMatch[0]}`);
                  }
                } else if (fullContentResponse.trim()) {
                  console.warn("[AI Response] Final response doesn't contain proper HTML");
                  
                  // Search for any HTML-like content
                  if (fullContentResponse.includes("<body") || 
                      (fullContentResponse.includes("<div") && fullContentResponse.includes("</div>"))) {
                    
                    console.log("[AI Response] Found partial HTML content, wrapping it");
                    // Wrap the content in a basic HTML structure
                    const wrappedContent = `<!DOCTYPE html>
<html>
<head>
  <title>AI Generated Content</title>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
</head>
<body>
  ${fullContentResponse}
</body>
</html>`;
                    setHtml(wrappedContent);
                  } else {
                    // If it's just text, wrap it in pre tag
                    console.log("[AI Response] No HTML found, creating preview with text content");
                    const textContent = `<!DOCTYPE html>
<html>
<head>
  <title>AI Generated Text</title>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <style>
    body { font-family: Arial, sans-serif; padding: 20px; }
    pre { white-space: pre-wrap; word-break: break-word; }
  </style>
</head>
<body>
  <pre>${fullContentResponse}</pre>
</body>
</html>`;
                    setHtml(textContent);
                  }
                }
              }

              toast.success("AI processing complete");
              setPrompt("");
              setPreviousPrompt(prompt);
              setisAiWorking(false);
              setHasAsked(true);
              audio.play();
              break; // Exit the loop
            }

            const chunk = decoder.decode(value, { stream: true });
            totalChunksReceived++;
            
            // Update progress indicator with better estimation
            // First 5 chunks are typically slower (setup)
            let estimatedProgress;
            if (totalChunksReceived <= 5) {
              estimatedProgress = Math.floor((totalChunksReceived / 5) * 25); // First 25%
            } else if (totalChunksReceived <= 15) {
              estimatedProgress = 25 + Math.floor(((totalChunksReceived - 5) / 10) * 50); // Next 50%
            } else {
              estimatedProgress = 75 + Math.floor(((totalChunksReceived - 15) / 10) * 20); // Final 25%
            }
            estimatedProgress = Math.min(95, estimatedProgress);
            setStreamProgress(estimatedProgress);
            
            console.log("[AI Stream] Received chunk of length:", chunk.length);
            if (chunk.length > 0) {
              console.log("[AI Stream] Sample:", chunk.substring(0, 50) + (chunk.length > 50 ? "..." : ""));
            }

            if (responseType === 'diff') {
              // --- Diff Mode ---
              accumulatedDiffResponse += chunk; // Just accumulate the raw response
              console.log("[AI Diff] Accumulated diff length now:", accumulatedDiffResponse.length);
            } else {
              // --- Full HTML Mode ---
              fullContentResponse += chunk; // Accumulate for preview
              // Use regex to find the start of the HTML doc
              const newHtmlMatch = fullContentResponse.match(/<!DOCTYPE html>[\s\S]*/);
              const newHtml = newHtmlMatch ? newHtmlMatch[0] : null;

              if (newHtml) {
                console.log("[AI Full] Found HTML content, length:", newHtml.length);
                // Throttle the re-renders to avoid flashing/flicker - reduced from 300ms to 150ms
                const now = Date.now();
                if (now - lastRenderTime > 150) {
                   // Force-close the HTML tag for preview if needed
                   let partialDoc = newHtml;
                   if (!partialDoc.trim().endsWith("</html>")) {
                       partialDoc += "\n</html>";
                   }
                  console.log("[AI Full] Updating preview with partial HTML");
                  setHtml(partialDoc); // Update the preview iframe content
                  lastRenderTime = now;
                }

                // Scroll editor down if content is long (heuristic)
                if (newHtml.length > 200 && now - lastRenderTime < 50) { // Only scroll if recently rendered
                  onScrollToBottom();
                }
              } else {
                console.log("[AI Full] No HTML content found yet in stream");
              }
            }
          } catch (streamError) {
            console.error("[AI Response] Error processing stream chunk:", streamError);
            // Continue trying to read the stream despite error
          }
        } // end while loop
      } else {
         throw new Error("Response body is null");
      }
    } catch (error: unknown) {
      setisAiWorking(false);
      setStreamProgress(0);
      
      // Handle the error with proper type checking
      if (error instanceof Error) {
        toast.error(error.message);
      } else if (typeof error === 'object' && error !== null && 'openLogin' in error) {
        const loginError = error as { openLogin: boolean, message?: string };
        toast.error(loginError.message || 'Authentication error');
        if (loginError.openLogin) {
          setOpen(true);
        }
      } else {
        toast.error('An unexpected error occurred');
      }
    }
  };

  return (
    <div
      className={`bg-gray-950 rounded-xl py-2 lg:py-2.5 pl-3.5 lg:pl-4 pr-2 lg:pr-2.5 absolute lg:sticky bottom-3 left-3 lg:bottom-4 lg:left-4 w-[calc(100%-1.5rem)] lg:w-[calc(100%-2rem)] z-10 group`}
    >
      {/* Progress indicator */}
      {isAiWorking && (
        <div className="absolute top-0 left-0 h-2 bg-gradient-to-r from-purple-500 via-pink-500 to-blue-500 transition-all duration-200 ease-in-out shadow-lg animate-pulse"
             style={{ width: `${streamProgress}%`, borderTopLeftRadius: '0.75rem', zIndex: 20 }}>
          <div className="absolute right-2 top-1/2 -translate-y-1/2 flex items-center gap-2">
            <span className="text-xs font-bold text-white drop-shadow-lg">
              {streamProgress}%
            </span>
            <span className="inline-block w-3 h-3 border-2 border-white border-t-transparent rounded-full animate-spin"></span>
          </div>
        </div>
      )}
      
      <div className="w-full relative flex items-center justify-between">
        <RiSparkling2Fill className={`text-lg lg:text-xl ${isAiWorking ? 'text-pink-500 animate-pulse' : 'text-gray-500'} group-focus-within:text-pink-500`} />
        <input
          type="text"
          disabled={isAiWorking}
          className="w-full bg-transparent max-lg:text-sm outline-none pl-3 text-white placeholder:text-gray-500 font-code"
          placeholder={
            hasAsked ? "What do you want to ask AI next?" : "Ask AI anything..."
          }
          value={prompt}
          onChange={(e) => setPrompt(e.target.value)}
          onKeyDown={(e) => {
            if (e.key === "Enter") {
              callAi();
            }
          }}
        />
        {isAiWorking ? (
          <div className="flex-none flex items-center justify-center rounded-full text-sm font-semibold size-8 text-center bg-pink-500 text-white shadow-sm dark:shadow-highlight/20">
            <div className="animate-spin h-4 w-4 border-2 border-white border-t-transparent rounded-full mr-1"></div>
            <span className="ml-1 text-xs font-bold animate-ellipsis">{streamProgress}%</span>
          </div>
        ) : (
          <button
            disabled={isAiWorking}
            className="relative overflow-hidden cursor-pointer flex-none flex items-center justify-center rounded-full text-sm font-semibold size-8 text-center bg-pink-500 hover:bg-pink-400 text-white shadow-sm dark:shadow-highlight/20 disabled:bg-gray-300 disabled:text-gray-500 disabled:cursor-not-allowed disabled:hover:bg-gray-300"
            onClick={callAi}
          >
            <GrSend className="-translate-x-[1px]" />
          </button>
        )}
      </div>
      
      {/* Response mode indicator */}
      {isAiWorking && responseMode && (
        <div className="absolute -top-8 right-0 text-xs font-medium text-white bg-gray-800/90 backdrop-blur-sm px-3 py-1.5 rounded-full flex items-center gap-1.5 animate-pulse border border-gray-700 shadow-lg z-20">
          <div className="size-2 rounded-full bg-pink-500 animate-pulse"></div>
          {responseMode === 'diff' ? 'Applying changes to code...' : 'Generating new HTML...'}
        </div>
      )}
      
      <div
        className={classNames(
          "h-screen w-screen bg-black/20 fixed left-0 top-0 z-10",
          {
            "opacity-0 pointer-events-none": !open,
          }
        )}
        onClick={() => setOpen(false)}
      ></div>
      <div
        className={classNames(
          "absolute top-0 -translate-y-[calc(100%+8px)] right-0 z-10 w-80 bg-white border border-gray-200 rounded-lg shadow-lg transition-all duration-75 overflow-hidden",
          {
            "opacity-0 pointer-events-none": !open,
          }
        )}
      >
        <Login html={html}>
          <p className="text-gray-500 text-sm mb-3">
            You reached the limit of free AI usage. Please login to continue.
          </p>
        </Login>
      </div>
    </div>
  );
}

export default AskAI;