diff --git a/src/app/actions/ai-improve.ts b/src/app/actions/ai-improve.ts
index d73a47a..ae78494 100644
--- a/src/app/actions/ai-improve.ts
+++ b/src/app/actions/ai-improve.ts
@@ -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);
diff --git a/src/components/ai-optimized-editor.tsx b/src/components/ai-optimized-editor.tsx
index adfd522..3e40f94 100644
--- a/src/components/ai-optimized-editor.tsx
+++ b/src/components/ai-optimized-editor.tsx
@@ -1,127 +1,80 @@
 "use client";
 
-import { useCallback, useState } from "react";
+import React, { useState, useEffect } from "react";
 import { DiffEditor } from "@monaco-editor/react";
+import { useMonacoTheme } from "@/hooks/use-monaco-theme";
 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 { AIOptimizeButton } from "@/features/problems/code/components/toolbar/actions/AIOptimizeButton";
 
 interface AIEditorWrapperProps {
-  language?: string;
-  value?: string;
-  path?: string;
-  problemId?: string;
-  languageServerConfigs?: LanguageServerConfig[];
-  onChange?: (value: string) => void;
-  className?: string;
+    language: string;
+    originalCode: string;
+    onReset: () => void;
 }
 
-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);
+export const AIEditorWrapper = ({ language, originalCode, onReset }: AIEditorWrapperProps) => {
+    const [optimizedCode, setOptimizedCode] = useState<string>("");
+    const [loading, setLoading] = useState(false);
+    const { theme } = useMonacoTheme();
 
-  const handleCodeChange = useCallback(
-    (val: string) => {
-      setCurrentCode(val);
-      onChange?.(val);
-    },
-    [onChange]
-  );
+    // 自动在组件首次挂载后执行 AI 优化
+    useEffect(() => {
+        if (!optimizedCode) {
+            handleOptimize();
+        }
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
 
-  const handleOptimize = useCallback(async () => {
-    if (!problemId || !currentCode) return;
-    setIsOptimizing(true);
-    setError(null);
+    const handleOptimize = async () => {
+        setLoading(true);
+        try {
+            const res = await optimizeCode({
+                code: originalCode,
+                error: "",
+                problemId: "",
+            });
+            setOptimizedCode(res.optimizedCode);
+        } catch (err) {
+            console.error("优化失败", err);
+            setOptimizedCode("// 优化失败，请稍后重试");
+        } finally {
+            setLoading(false);
+        }
+    };
 
-    try {
-      const input: OptimizeCodeInput = {
-        code: currentCode,
-        problemId,
-      };
-      const result = await optimizeCode(input);
-      setOptimizedCode(result.optimizedCode);
-      setShowDiff(true);
-    } catch (err) {
-      setError("AI 优化失败，请稍后重试");
-      console.error(err);
-    } finally {
-      setIsOptimizing(false);
-    }
-  }, [currentCode, problemId]);
+    const handleClick = () => {
+        if (optimizedCode) {
+            // 已有优化，点击返回
+            setOptimizedCode("");
+            onReset();
+        } else {
+            // 手动触发优化（如果需要）
+            handleOptimize();
+        }
+    };
 
-  const handleApplyOptimized = useCallback(() => {
-    setCurrentCode(optimizedCode);
-    onChange?.(optimizedCode);
-    setShowDiff(false);
-  }, [optimizedCode, onChange]);
+    return (
+        <div className="w-full h-[80vh] flex flex-col gap-4">
+            <div>
+                <AIOptimizeButton
+                    loading={loading}
+                    hasOptimized={!!optimizedCode}
+                    onClick={handleClick}
+                />
+            </div>
 
-  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 ? (
-          <DiffEditor
-            original={currentCode}
-            modified={optimizedCode}
-            language={language}
-            theme="vs-dark"
-            className="h-full w-full"
-            options={{ readOnly: true, minimap: { enabled: false } }}
-          />
-        ) : (
-          <CoreEditor
-            language={language}
-            value={currentCode}
-            path={path}
-            languageServerConfigs={languageServerConfigs}
-            onChange={handleCodeChange}
-            className="h-full w-full"
-          />
-        )}
-      </div>
-    </div>
-  );
+            {optimizedCode && (
+                <div className="flex-1">
+                    <DiffEditor
+                        language={language}
+                        original={originalCode}
+                        modified={optimizedCode}
+                        height="100%"
+                        theme={theme}
+                        options={{ readOnly: true, renderSideBySide: true, automaticLayout: true }}
+                    />
+                </div>
+            )}
+        </div>
+    );
 };
diff --git a/src/components/problem-editor.tsx b/src/components/problem-editor.tsx
index c9a8adc..a591203 100644
--- a/src/components/problem-editor.tsx
+++ b/src/components/problem-editor.tsx
@@ -1,46 +1,68 @@
 "use client";
 
-import { useEffect } from "react";
+import { useState, useEffect } from "react";
 import { CoreEditor } from "@/components/core-editor";
 import { useProblemEditorStore } from "@/stores/problem-editor";
 import type { LanguageServerConfig, Template } from "@/generated/client";
+import { AIEditorWrapper } from "@/components/ai-optimized-editor";
 
 interface ProblemEditorProps {
-  problemId: string;
-  templates: Template[];
-  languageServerConfigs?: LanguageServerConfig[];
+    problemId: string;
+    templates: Template[];
+    languageServerConfigs?: LanguageServerConfig[];
 }
 
 export const ProblemEditor = ({
-  problemId,
-  templates,
-  languageServerConfigs,
-}: ProblemEditorProps) => {
-  const {
-    language,
-    value,
-    path,
-    setProblem,
-    setValue,
-    setEditor,
-    setLspWebSocket,
-    setMarkers,
-  } = useProblemEditorStore();
+                                  problemId,
+                                  templates,
+                                  languageServerConfigs,
+                              }: ProblemEditorProps) => {
+    const {
+        language,
+        value,
+        path,
+        setProblem,
+        setValue,
+        setEditor,
+        setLspWebSocket,
+        setMarkers,
+    } = useProblemEditorStore();
 
-  useEffect(() => {
-    setProblem(problemId, templates);
-  }, [problemId, setProblem, templates]);
+    const [useAIEditor, setUseAIEditor] = useState(false);
 
-  return (
-    <CoreEditor
-      language={language}
-      value={value}
-      path={path}
-      languageServerConfigs={languageServerConfigs}
-      onEditorReady={setEditor}
-      onLspWebSocketReady={setLspWebSocket}
-      onMarkersReady={setMarkers}
-      onChange={setValue}
-    />
-  );
+    useEffect(() => {
+        setProblem(problemId, templates);
+    }, [problemId, setProblem, templates]);
+
+    return (
+        <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
+                    language={language}
+                    originalCode={value}
+                    onReset={() => setUseAIEditor(false)}
+                />
+            )}
+        </div>
+    );
 };