diff --git a/src/components/ai-optimized-editor.tsx b/src/components/ai-optimized-editor.tsx
index 87d1c2c..68b082d 100644
--- a/src/components/ai-optimized-editor.tsx
+++ b/src/components/ai-optimized-editor.tsx
@@ -1,250 +1,124 @@
 "use client";
 
-import { useState } from "react";
+import { useCallback, useState } from "react";
 import { DiffEditor } from "@monaco-editor/react";
 import { optimizeCode } from "@/app/actions/ai-improve";
-import { OptimizeCodeInput } from "@/types/ai-improve";
-import dynamic from "next/dynamic";
-import { highlighter } from "@/lib/shiki";
-import type { editor } from "monaco-editor";
-import { Loading } from "@/components/loading";
-import { shikiToMonaco } from "@shikijs/monaco";
-import { useProblem } from "@/hooks/use-problem";
-import type { Monaco } from "@monaco-editor/react";
-import { useCallback, useEffect, useRef } from "react";
-import { connectToLanguageServer } from "@/lib/language-server";
-import type { MonacoLanguageClient } from "monaco-languageclient";
-import { DefaultEditorOptionConfig } from "@/config/editor-option";
+import type { OptimizeCodeInput } from "@/types/ai-improve";
+import { CoreEditor } from "./core-editor"; // 引入你刚刚的组件
+// import { Loading } from "@/components/loading";
+import type { LanguageServerConfig } from "@/generated/client";
 
-// 动态导入Monaco Editor
-const Editor = dynamic(
-    async () => {
-        await import("vscode");
-        const monaco = await import("monaco-editor");
+interface AIEditorWrapperProps {
+  language?: string;
+  value?: string;
+  path?: string;
+  problemId?: string;
+  languageServerConfigs?: LanguageServerConfig[];
+  onChange?: (value: string) => void;
+  className?: string;
+}
 
-        self.MonacoEnvironment = {
-            getWorker(_, label) {
-                if (label === "json") {
-                    return new Worker(
-                        new URL("monaco-editor/esm/vs/language/json/json.worker.js", import.meta.url)
-                    );
-                }
-                if (label === "css" || label === "scss" || label === "less") {
-                    return new Worker(
-                        new URL("monaco-editor/esm/vs/language/css/css.worker.js", import.meta.url)
-                    );
-                }
-                if (label === "html" || label === "handlebars" || label === "razor") {
-                    return new Worker(
-                        new URL("monaco-editor/esm/vs/language/html/html.worker.js", import.meta.url)
-                    );
-                }
-                if (label === "typescript" || label === "javascript") {
-                    return new Worker(
-                        new URL("monaco-editor/esm/vs/language/typescript/ts.worker.js", import.meta.url)
-                    );
-                }
-                return new Worker(
-                    new URL("monaco-editor/esm/vs/editor/editor.worker.js", import.meta.url)
-                );
-            },
-        };
-        const { loader } = await import("@monaco-editor/react");
-        loader.config({ monaco });
-        return (await import("@monaco-editor/react")).Editor;
-    },
-    {
-        ssr: false,
-        loading: () => <Loading />,
+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 handleOptimize = useCallback(async () => {
+    if (!problemId || !currentCode) return;
+    setIsOptimizing(true);
+    setError(null);
+
+    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]);
 
-export function AIProblemEditor({
-                                    initialCode = "",
-                                    problemId = "",
-                                    onCodeChange
-                                }: {
-    initialCode?: string;
-    problemId?: string;
-    onCodeChange?: (code: string) => void;
-}) {
-    const {
-        editor,
-        setEditor,
-        setMarkers,
-        setWebSocket,
-        currentLang,
-        currentPath,
-        currentTheme,
-        currentValue,
-        changeValue,
-        currentEditorLanguageConfig,
-        currentLanguageServerConfig,
-    } = useProblem();
+  const handleApplyOptimized = useCallback(() => {
+    setCurrentCode(optimizedCode);
+    onChange?.(optimizedCode);
+    setShowDiff(false);
+  }, [optimizedCode, onChange]);
 
-    const monacoLanguageClientRef = useRef<MonacoLanguageClient | null>(null);
+  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>
 
-    // 保持原有AI优化的状态
-    const [showDiff, setShowDiff] = useState(false);
-    const [optimizedCode, setOptimizedCode] = useState("");
-    const [isOptimizing, setIsOptimizing] = useState(false);
-    const [error, setError] = useState<string | null>(null);
-
-    // 重用useProblem的状态管理
-    const currentCode = currentValue || initialCode;
-
-    const handleCodeChange = useCallback((value: string | undefined) => {
-        if (value !== undefined) {
-            changeValue(value);
-            if (onCodeChange) {
-                onCodeChange(value);
-            }
-        }
-    }, [onCodeChange, changeValue]);
-
-    // 保持原有LSP连接逻辑
-    const connectLSP = useCallback(async () => {
-        if (!(currentLang && editor)) return;
-
-        if (monacoLanguageClientRef.current) {
-            monacoLanguageClientRef.current.stop();
-            monacoLanguageClientRef.current = null;
-            setWebSocket(null);
-        }
-
-        if (!currentEditorLanguageConfig || !currentLanguageServerConfig) return;
-
-        try {
-            const { client: monacoLanguageClient, webSocket } = await connectToLanguageServer(
-                currentEditorLanguageConfig,
-                currentLanguageServerConfig
-            );
-            monacoLanguageClientRef.current = monacoLanguageClient;
-            setWebSocket(webSocket);
-        } catch (error) {
-            console.error("Failed to connect to LSP:", error);
-        }
-    }, [
-        currentEditorLanguageConfig,
-        currentLang,
-        currentLanguageServerConfig,
-        editor,
-        setWebSocket,
-    ]);
-
-    useEffect(() => {
-        connectLSP();
-    }, [connectLSP]);
-
-    useEffect(() => {
-        return () => {
-            if (monacoLanguageClientRef.current) {
-                monacoLanguageClientRef.current.stop();
-                monacoLanguageClientRef.current = null;
-                setWebSocket(null);
-            }
-        };
-    }, [setWebSocket]);
-
-    const handleEditorWillMount = useCallback((monaco: Monaco) => {
-        shikiToMonaco(highlighter, monaco);
-    }, []);
-
-    const handleOnMount = useCallback(
-        async (editor: editor.IStandaloneCodeEditor) => {
-            setEditor(editor);
-            await connectLSP();
-        },
-        [setEditor, connectLSP]
-    );
-
-    const handleEditorValidation = useCallback(
-        (markers: editor.IMarker[]) => {
-            setMarkers(markers);
-        },
-        [setMarkers]
-    );
-
-    const handleOptimizeCode = useCallback(async () => {
-        if (!currentCode || !problemId) return;
-
-        setIsOptimizing(true);
-        setError(null);
-
-        try {
-            const input: OptimizeCodeInput = {
-                code: currentCode,
-                problemId
-            };
-
-            const result = await optimizeCode(input);
-            setOptimizedCode(result.optimizedCode);
-            setShowDiff(true);
-        } catch (err) {
-            setError("代码优化失败，请重试");
-            console.error(err);
-        } finally {
-            setIsOptimizing(false);
-        }
-    }, [currentCode, problemId]);
-
-    return (
-        <div className="flex flex-col h-full w-full">
-            {/* 保持原有AI优化按钮 */}
-            <div className="flex justify-between items-center p-4">
+          {showDiff && (
+              <div className="space-x-2">
                 <button
-                    onClick={handleOptimizeCode}
-                    disabled={isOptimizing || !currentCode}
-                    className="px-4 py-2 bg-primary text-primary-foreground rounded-md hover:bg-primary/90"
+                    onClick={() => setShowDiff(false)}
+                    className="px-4 py-2 bg-secondary text-white rounded"
                 >
-                    {isOptimizing ? "优化中..." : "AI优化代码"}
+                  隐藏对比
                 </button>
-
-                {showDiff && (
-                    <button
-                        onClick={() => setShowDiff(!showDiff)}
-                        className="px-4 py-2 bg-secondary text-secondary-foreground rounded-md"
-                    >
-                        {"隐藏对比"}
-                    </button>
-                )}
-            </div>
-
-            {error && (
-                <div className="p-3 bg-destructive/10 text-destructive rounded-md">
-                    {error}
-                </div>
-            )}
-
-            <div className="flex-grow overflow-hidden">
-                {showDiff ? (
-                    <DiffEditor
-                        original={currentCode}
-                        modified={optimizedCode}
-                        language={currentLang}
-                        theme={currentTheme}
-                        className="h-full w-full"
-                        options={{
-                            readOnly: true,
-                            minimap: { enabled: false }
-                        }}
-                    />
-                ) : (
-                    <Editor
-                        language={currentLang}
-                        theme={currentTheme}
-                        path={currentPath}
-                        value={currentCode}
-                        beforeMount={handleEditorWillMount}
-                        onMount={handleOnMount}
-                        onChange={handleCodeChange}
-                        onValidate={handleEditorValidation}
-                        options={DefaultEditorOptionConfig}
-                        loading={<Loading />}
-                        className="h-full w-full"
-                    />
-                )}
-            </div>
+                <button
+                    onClick={handleApplyOptimized}
+                    className="px-4 py-2 bg-green-500 text-white rounded"
+                >
+                  应用优化结果
+                </button>
+              </div>
+          )}
         </div>
-    );
-}
\ No newline at end of file
+
+        {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>
+  );
+};