diff --git a/src/components/ai-optimized-editor.tsx b/src/components/ai-optimized-editor.tsx
index 68b082d..87d1c2c 100644
--- a/src/components/ai-optimized-editor.tsx
+++ b/src/components/ai-optimized-editor.tsx
@@ -1,124 +1,250 @@
 "use client";
 
-import { useCallback, useState } from "react";
+import { 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 { 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";
 
-interface AIEditorWrapperProps {
-  language?: string;
-  value?: string;
-  path?: string;
-  problemId?: string;
-  languageServerConfigs?: LanguageServerConfig[];
-  onChange?: (value: string) => void;
-  className?: string;
-}
+// 动态导入Monaco Editor
+const Editor = dynamic(
+    async () => {
+        await import("vscode");
+        const monaco = await import("monaco-editor");
 
-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);
+        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 />,
     }
-  }, [currentCode, problemId]);
+);
 
-  const handleApplyOptimized = useCallback(() => {
-    setCurrentCode(optimizedCode);
-    onChange?.(optimizedCode);
-    setShowDiff(false);
-  }, [optimizedCode, onChange]);
+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();
 
-  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>
+    const monacoLanguageClientRef = useRef<MonacoLanguageClient | null>(null);
 
-          {showDiff && (
-              <div className="space-x-2">
+    // 保持原有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">
                 <button
-                    onClick={() => setShowDiff(false)}
-                    className="px-4 py-2 bg-secondary text-white rounded"
+                    onClick={handleOptimizeCode}
+                    disabled={isOptimizing || !currentCode}
+                    className="px-4 py-2 bg-primary text-primary-foreground rounded-md hover:bg-primary/90"
                 >
-                  隐藏对比
+                    {isOptimizing ? "优化中..." : "AI优化代码"}
                 </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>
-        )}
+                {showDiff && (
+                    <button
+                        onClick={() => setShowDiff(!showDiff)}
+                        className="px-4 py-2 bg-secondary text-secondary-foreground rounded-md"
+                    >
+                        {"隐藏对比"}
+                    </button>
+                )}
+            </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"
-              />
-          )}
+            {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>
         </div>
-      </div>
-  );
-};
+    );
+}
\ No newline at end of file