package com.axis.innovators.box.browser;

import org.cef.browser.CefBrowser;
import org.cef.browser.CefFrame;
import org.cef.browser.CefMessageRouter;
import org.cef.callback.CefQueryCallback;
import org.cef.handler.CefMessageRouterHandlerAdapter;
import org.tzd.lm.LM;

import javax.swing.*;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 这是一个简单的示例程序，用于展示如何使用JCEF来创建一个简单的浏览器窗口。
 */
public class MainApplication {
    private static final ExecutorService executor = Executors.newCachedThreadPool();
    private static long modelHandle;
    private static long ctxHandle;
    private static boolean isSystem = true;
    public static void main(String[] args) {
        LM.loadLibrary(LM.CUDA);
        modelHandle = LM.llamaLoadModelFromFile(LM.DEEP_SEEK);
        ctxHandle = LM.createContext(modelHandle);

        AtomicReference<BrowserWindow> window = new AtomicReference<>();
        SwingUtilities.invokeLater(() -> {
            WindowRegistry.getInstance().createNewWindow("main", builder ->
                    window.set(builder.title("Axis Innovators Box")
                            .size(1280, 720)
                            .htmlPath("C:\\Users\\Administrator\\MCreatorWorkspaces\\AxisInnovatorsBox\\src\\main\\java\\com\\axis\\innovators\\box\\browser\\DeepSeek - 探索未至之境.html")
                            .operationHandler(createOperationHandler())
                            .build())
            );

            CefMessageRouter msgRouter = window.get().getMsgRouter();
            if (msgRouter != null) {
                msgRouter.addHandler(new CefMessageRouterHandlerAdapter() {
                    @Override
                    public boolean onQuery(CefBrowser browser, CefFrame frame, long queryId,
                                           String request, boolean persistent, CefQueryCallback callback) {
                        // 处理浏览器请求
                        handleBrowserQuery(browser, request, callback);
                        return true;
                    }

                    @Override
                    public void onQueryCanceled(CefBrowser browser, CefFrame frame, long queryId) {
                        // 处理请求取消
                    }
                }, true);
            }
        });

        // 关闭钩子
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            LM.llamaFreeContext(ctxHandle);
            LM.llamaFreeModel(modelHandle);
            executor.shutdown();
        }));
    }

    private static void handleBrowserQuery(CefBrowser browser, String request, CefQueryCallback callback) {
        try {
            String[] parts = request.split(":", 3);
            if (parts.length < 3) {
                callback.failure(400, "请求格式错误");
                return;
            }

            String operation = parts[0];
            String requestId = parts[1];
            String prompt = parts[2];

            if ("ai-inference".equals(operation)) {
                executor.execute(() -> {
                    //String system;
                    if (isSystem) {
                        //system = ;
                        isSystem = false;
                    } //else {
                        //system = null;
                    //}
                    List<String> messageList = new java.util.ArrayList<>(List.of());
                    // 修改后的推理回调处理
                    String jsCode = String.format(
                            "if (typeof updateResponse === 'function') {" +
                                    "  updateResponse('%s', '%s');" +
                                    "}",
                            requestId, "<details open style=background: rgba(0,0,0,0.03);border: 1px solid #eee;border-radius: 4px;padding: 8px;margin: 8px 0;>" +
                                    "\\n" +
                                    "<summary>推理内容</summary><font face=\\\"黑体\\\" color=grey size=3>"
                    );
                    browser.executeJavaScript(jsCode, null, 0);
                    LM.inference(modelHandle, ctxHandle, 0.6f, prompt + "<think>\n",
                            """
                                # 角色设定
                                你是一个严格遵循规则的AI助手,
                                当遇到简单的问题时你可以直接回答（回答先添加</think>）,在遇到复杂问题时请你思考后再给出答案，并且需满足以下要求：

                                ## 核心指令
                                1. **强制推理标记**：无论是否推理，回答**必须**以`<think>`开头，推理结束时闭合`</think>`。
                                2. **推理流程**：
                                   - 使用固定句式（如“好的，用户现在需要...”）启动分析
                                3. **回答规范**：
                                   - 使用Markdown标题、列表、加粗突出重点
                                   - LaTeX公式严格包裹在`$$...$$`中（示例：`$$E=mc^2$$`）
                                   - 不同可能性答案用列表呈现

                                ## 违规惩罚
                                - 若未包含`<think>`标签，需在回答开头添加 </think>
                                """,
                            new LM.MessageCallback() {
                                private boolean thinkingClosed = false;

                                @Override
                                public void onMessage(String message) {
                                    messageList.add(message);
                                    SwingUtilities.invokeLater(() -> {
                                        // 统一转义处理
                                        String escaped = message
                                                .replace("\\", "\\\\")
                                                .replace("'", "\\'")
                                                .replace("\"", "\\\"")
                                                .replace("\n", "\\n")
                                                .replace("\r", "\\r");

                                        if (messageList.contains("</think>") && !thinkingClosed) {
                                            String endJs = String.format(
                                                    "if (typeof updateResponse === 'function') {" +
                                                            "  updateResponse('%s', '%s');" +
                                                            "}",
                                                    requestId, "</font></details>"
                                            );
                                            browser.executeJavaScript(endJs, null, 0);
                                            thinkingClosed = true;
                                        }

                                        // 实时更新内容
                                        String jsCode = String.format(
                                                "if (typeof updateResponse === 'function') {" +
                                                        "  updateResponse('%s', '%s');" +
                                                        "}",
                                                requestId, escaped
                                        );
                                        browser.executeJavaScript(jsCode, null, 0);
                                    });
                                }
                            },isSystem);
                    messageList.clear();
                    callback.success("COMPLETED:" + requestId);
                });
            }
        } catch (Exception e) {
            callback.failure(500, "服务器错误: " + e.getMessage());
        }
    }

    private static WindowOperationHandler createOperationHandler() {
        return new WindowOperationHandler.Builder()
                .withDefaultOperations()
                .build();
    }
}
