feat(ai): integrate AI code optimization with template support and editor toggling

src/app/actions/ai-improve.ts

@@ -5,9 +5,9 @@ import {
   OptimizeCodeOutput,
   OptimizeCodeOutputSchema,
 } from "@/types/ai-improve";
+import prisma from "@/lib/prisma";
 import { deepseek } from "@/lib/ai";
 import { CoreMessage, generateText } from "ai";
-import prisma from "@/lib/prisma";
 
 /**
  * 调用AI优化代码
@@ -17,13 +17,29 @@ import prisma from "@/lib/prisma";
 export const optimizeCode = async (
   input: OptimizeCodeInput
 ): Promise<OptimizeCodeOutput> => {
-  const model = deepseek("chat");
+  const model = deepseek("deepseek-chat");
 
   // 获取题目详情（如果提供了problemId）
   let problemDetails = "";
+  let templateDetails = "";
 
   if (input.problemId) {
     try {
+      const templates = await prisma.template.findMany({
+        where: { problemId: input.problemId },
+      });
+      if (templates && templates.length > 0) {
+        const tplStrings = templates
+          .map(
+            (t) =>
+              `Template (${t.language}):\n-------------------\n\`\`\`\n${t.content}\n\`\`\``
+          )
+          .join("\n");
+        templateDetails = `\nCode Templates:\n-------------------\n${tplStrings}`;
+      } else {
+        templateDetails = "\nNo code templates found for this problem.";
+      }
+
       // 尝试获取英文描述
       const problemLocalizationEn = await prisma.problemLocalization.findUnique(
         {
@@ -102,6 +118,12 @@ Error message (if any): ${input.error || "No error message provided"}
 
 ${problemDetails}
 
+The following is the code template section, do not modify the part that gives the same code as the code template
+
+${templateDetails}
+
+Write the code in conjunction with the topic and fill in the gaps in the code
+
 Respond ONLY with the JSON object containing the optimized code and explanations.
 Format:
 {
@@ -127,9 +149,13 @@ Format:
   }
 
   // 解析LLM响应
+  const fenceMatch = text.match(/```(?:json)?\s*([\s\S]*?)```/);
+  const jsonMatch = fenceMatch ? fenceMatch[1] : text.match(/\{[\s\S]*}/)?.[0];
+  const jsonString = jsonMatch ? jsonMatch.trim() : text.trim();
+
   let llmResponseJson;
   try {
-    const cleanedText = text.trim();
+    const cleanedText = jsonString.trim();
     llmResponseJson = JSON.parse(cleanedText);
   } catch (error) {
     console.error("Failed to parse LLM response as JSON:", error);

src/components/ai-optimized-editor.tsx

@@ -1,127 +1,68 @@
 "use client";
 
-import { useCallback, useState } from "react";
 import { DiffEditor } from "@monaco-editor/react";
 import { optimizeCode } from "@/app/actions/ai-improve";
-import type { OptimizeCodeInput } from "@/types/ai-improve";
-import { CoreEditor } from "./core-editor"; // 引入你刚刚的组件
-// import { Loading } from "@/components/loading";
-import type { LanguageServerConfig } from "@/generated/client";
+import { useMonacoTheme } from "@/hooks/use-monaco-theme";
+import React, { useState, useEffect, useCallback } from "react";
+import { useProblemEditorStore } from "@/stores/problem-editor";
 
-interface AIEditorWrapperProps {
-  language?: string;
-  value?: string;
-  path?: string;
-  problemId?: string;
-  languageServerConfigs?: LanguageServerConfig[];
-  onChange?: (value: string) => void;
-  className?: string;
-}
+export const AIEditorWrapper = () => {
+  const {
+    language,
+    value: originalCode,
+    setLoading,
+    AIgenerate,
+    LastOptimizedCode,
+    setLastOptimizedCode,
+  } = useProblemEditorStore();
 
-export const AIEditorWrapper = ({
-  language,
-  value,
-  path,
-  problemId,
-  languageServerConfigs,
-  onChange,
-}: // className,
-AIEditorWrapperProps) => {
-  const [currentCode, setCurrentCode] = useState(value ?? "");
-  const [optimizedCode, setOptimizedCode] = useState("");
-  const [isOptimizing, setIsOptimizing] = useState(false);
-  const [error, setError] = useState<string | null>(null);
-  const [showDiff, setShowDiff] = useState(false);
-
-  const handleCodeChange = useCallback(
-    (val: string) => {
-      setCurrentCode(val);
-      onChange?.(val);
-    },
-    [onChange]
-  );
+  const [optimizedCode, setOptimizedCode] = useState<string>("");
+  const { theme } = useMonacoTheme();
 
   const handleOptimize = useCallback(async () => {
-    if (!problemId || !currentCode) return;
-    setIsOptimizing(true);
-    setError(null);
-
+    setLoading(true);
     try {
-      const input: OptimizeCodeInput = {
-        code: currentCode,
-        problemId,
-      };
-      const result = await optimizeCode(input);
-      setOptimizedCode(result.optimizedCode);
-      setShowDiff(true);
+      const res = await optimizeCode({
+        code: originalCode,
+        error: "",
+        problemId: "",
+      });
+      setOptimizedCode(res.optimizedCode);
+      setLastOptimizedCode(res.optimizedCode);
     } catch (err) {
-      setError("AI 优化失败，请稍后重试");
-      console.error(err);
+      console.error("优化失败", err);
+      setOptimizedCode("// 优化失败，请稍后重试");
     } finally {
-      setIsOptimizing(false);
+      setLoading(false);
     }
-  }, [currentCode, problemId]);
+  }, [originalCode, setLoading, setLastOptimizedCode]);
 
-  const handleApplyOptimized = useCallback(() => {
-    setCurrentCode(optimizedCode);
-    onChange?.(optimizedCode);
-    setShowDiff(false);
-  }, [optimizedCode, onChange]);
+  useEffect(() => {
+    if (AIgenerate) {
+      handleOptimize();
+    } else if (LastOptimizedCode) {
+      setOptimizedCode(LastOptimizedCode);
+    }
+  }, [AIgenerate, LastOptimizedCode, handleOptimize]);
 
   return (
-    <div className="flex flex-col h-full w-full">
-      <div className="flex items-center justify-between p-4">
-        <button
-          onClick={handleOptimize}
-          disabled={isOptimizing}
-          className="px-4 py-2 bg-primary text-white rounded hover:bg-primary/90"
-        >
-          {isOptimizing ? "优化中..." : "AI优化代码"}
-        </button>
-
-        {showDiff && (
-          <div className="space-x-2">
-            <button
-              onClick={() => setShowDiff(false)}
-              className="px-4 py-2 bg-secondary text-white rounded"
-            >
-              隐藏对比
-            </button>
-            <button
-              onClick={handleApplyOptimized}
-              className="px-4 py-2 bg-green-500 text-white rounded"
-            >
-              应用优化结果
-            </button>
-          </div>
-        )}
-      </div>
-
-      {error && (
-        <div className="p-3 bg-red-100 text-red-600 rounded-md">{error}</div>
-      )}
-
-      <div className="flex-grow overflow-hidden">
-        {showDiff ? (
+    <div className="w-full h-[80vh] flex flex-col gap-4">
+      {optimizedCode && (
+        <div className="flex-1">
           <DiffEditor
-            original={currentCode}
+            language={language}
+            original={originalCode}
             modified={optimizedCode}
-            language={language}
-            theme="vs-dark"
-            className="h-full w-full"
-            options={{ readOnly: true, minimap: { enabled: false } }}
+            height="100%"
+            theme={theme}
+            options={{
+              readOnly: true,
+              renderSideBySide: true,
+              automaticLayout: true,
+            }}
           />
-        ) : (
-          <CoreEditor
-            language={language}
-            value={currentCode}
-            path={path}
-            languageServerConfigs={languageServerConfigs}
-            onChange={handleCodeChange}
-            className="h-full w-full"
-          />
-        )}
-      </div>
+        </div>
+      )}
     </div>
   );
 };

src/components/problem-editor.tsx

@@ -3,6 +3,7 @@
 import { useEffect } from "react";
 import { CoreEditor } from "@/components/core-editor";
 import { useProblemEditorStore } from "@/stores/problem-editor";
+import { AIEditorWrapper } from "@/components/ai-optimized-editor";
 import type { LanguageServerConfig, Template } from "@/generated/client";
 
 interface ProblemEditorProps {
@@ -25,6 +26,8 @@ export const ProblemEditor = ({
     setEditor,
     setLspWebSocket,
     setMarkers,
+    useAIEditor,
+    // setUseAIEditor
   } = useProblemEditorStore();
 
   useEffect(() => {
@@ -32,15 +35,30 @@ export const ProblemEditor = ({
   }, [problemId, setProblem, templates]);
 
   return (
-    <CoreEditor
-      language={language}
-      value={value}
-      path={path}
-      languageServerConfigs={languageServerConfigs}
-      onEditorReady={setEditor}
-      onLspWebSocketReady={setLspWebSocket}
-      onMarkersReady={setMarkers}
-      onChange={setValue}
-    />
+    <div className="w-full h-[85vh] relative">
+      {!useAIEditor ? (
+        <>
+          {/*<button*/}
+          {/*    className="absolute right-4 top-4 bg-blue-600 text-white px-3 py-1 rounded z-10"*/}
+          {/*    onClick={() => setUseAIEditor(true)}*/}
+          {/*>*/}
+          {/*    AI 优化代码*/}
+          {/*</button>*/}
+          <CoreEditor
+            language={language}
+            value={value}
+            path={path}
+            languageServerConfigs={languageServerConfigs}
+            onEditorReady={setEditor}
+            onLspWebSocketReady={setLspWebSocket}
+            onMarkersReady={setMarkers}
+            onChange={setValue}
+            className="h-[80vh] w-full"
+          />
+        </>
+      ) : (
+        <AIEditorWrapper />
+      )}
+    </div>
   );
 };

src/features/problems/code/components/toolbar/actions/AIDisplayButton.tsx

@@ -0,0 +1,44 @@
+"use client";
+
+import { TooltipButton } from "@/components/tooltip-button";
+import { useProblemEditorStore } from "@/stores/problem-editor";
+import { ArrowLeftRight, LoaderCircleIcon, Undo2Icon } from "lucide-react";
+
+export const AIDisplayButton = () => {
+  const { useAIEditor, setUseAIEditor, setAIgenerate, loading } =
+    useProblemEditorStore();
+
+  const handleClick = () => {
+    setAIgenerate(false);
+    if (!loading) {
+      setUseAIEditor(!useAIEditor);
+    }
+  };
+
+  const tooltipContent = loading
+    ? "AI 正在优化中…"
+    : useAIEditor
+    ? "返回原始编辑器"
+    : "查看 AI 优化代码";
+
+  return (
+    <TooltipButton
+      tooltipContent={tooltipContent}
+      onClick={handleClick}
+      disabled={loading}
+    >
+      {loading ? (
+        <LoaderCircleIcon
+          className="opacity-60 animate-spin"
+          size={16}
+          strokeWidth={2}
+          aria-hidden="true"
+        />
+      ) : useAIEditor ? (
+        <Undo2Icon size={16} strokeWidth={2} aria-hidden="true" />
+      ) : (
+        <ArrowLeftRight size={16} strokeWidth={2} aria-hidden="true" />
+      )}
+    </TooltipButton>
+  );
+};

src/features/problems/code/components/toolbar/actions/AIOptimizeButton.tsx

@@ -0,0 +1,42 @@
+"use client";
+
+import { Wand2Icon } from "lucide-react";
+import { TooltipButton } from "@/components/tooltip-button";
+//import { LoaderCircleIcon, Undo2Icon } from "lucide-react";
+import { useProblemEditorStore } from "@/stores/problem-editor";
+
+export const AIOptimizeButton = () => {
+  const { useAIEditor, setUseAIEditor, setAIgenerate, loading } =
+    useProblemEditorStore();
+
+  const handleClick = () => {
+    setAIgenerate(true);
+    if (!loading) {
+      setUseAIEditor(!useAIEditor);
+    }
+  };
+
+  //     ? "AI 正在优化中…"
+  //     : useAIEditor
+  //         ? "返回原始编辑器"
+  //         : "使用 AI 优化代码";
+
+  const tooltipContent = "使用 AI 优化代码"; // 仅保留默认提示内容
+
+  return (
+    <TooltipButton
+      tooltipContent={tooltipContent}
+      onClick={handleClick}
+      disabled={loading || useAIEditor}
+    >
+      {loading ? //     className="opacity-60 animate-spin" // <LoaderCircleIcon
+      //     size={16}
+      //     strokeWidth={2}
+      //     aria-hidden="true"
+      // />
+      null : useAIEditor ? null : ( // <Undo2Icon size={16} strokeWidth={2} aria-hidden="true" />
+        <Wand2Icon size={16} strokeWidth={2} aria-hidden="true" />
+      )}
+    </TooltipButton>
+  );
+};

src/features/problems/code/components/toolbar/code-toolbar.tsx

@@ -9,6 +9,8 @@ import {
 } from "@/features/problems/code/components/toolbar";
 import { AnalyzeButton } from "./actions/analyze-button";
 import { LspConnectionIndicator } from "./controls/lsp-connection-indicator";
+import { AIDisplayButton } from "@/features/problems/code/components/toolbar/actions/AIDisplayButton";
+import { AIOptimizeButton } from "@/features/problems/code/components/toolbar/actions/AIOptimizeButton";
 
 interface CodeToolbarProps {
   className?: string;
@@ -25,6 +27,8 @@ export const CodeToolbar = async ({ className }: CodeToolbarProps) => {
           <LspConnectionIndicator />
         </div>
         <div className="flex items-center gap-2">
+          <AIOptimizeButton />
+          <AIDisplayButton />
           <AnalyzeButton />
           <ResetButton />
           <UndoButton />

src/stores/problem-editor.ts

@@ -16,6 +16,10 @@ type ProblemEditorState = {
   editor: editor.IStandaloneCodeEditor | null;
   lspWebSocket: WebSocket | null;
   markers: editor.IMarker[];
+  useAIEditor: boolean;
+  loading: boolean;
+  AIgenerate: boolean;
+  LastOptimizedCode: string;
 };
 
 type ProblemEditorAction = {
@@ -26,6 +30,10 @@ type ProblemEditorAction = {
   setEditor: (editor: editor.IStandaloneCodeEditor) => void;
   setLspWebSocket: (lspWebSocket: WebSocket) => void;
   setMarkers: (markers: editor.IMarker[]) => void;
+  setUseAIEditor: (flag: boolean) => void;
+  setLoading: (flag: boolean) => void;
+  setAIgenerate: (flag: boolean) => void;
+  setLastOptimizedCode: (code: string) => void;
 };
 
 type ProblemEditorStore = ProblemEditorState & ProblemEditorAction;
@@ -38,6 +46,14 @@ export const useProblemEditorStore = create<ProblemEditorStore>((set, get) => ({
   editor: null,
   lspWebSocket: null,
   markers: [],
+  useAIEditor: false,
+  loading: false,
+  AIgenerate: false,
+  LastOptimizedCode: "",
+  setLastOptimizedCode: (code) => set({ LastOptimizedCode: code }),
+  setAIgenerate: (flag) => set({ AIgenerate: flag }),
+  setLoading: (loading) => set({ loading }),
+  setUseAIEditor: (loading) => set({ useAIEditor: loading }),
   setProblem: (problemId, templates) => {
     const language = getLanguage(problemId);
     const value = getValue(problemId, language, templates);