package com.easyagents.llm.openai;

import com.easyagents.core.agent.react.ReActAgent;
import com.easyagents.core.agent.react.ReActAgentListener;
import com.easyagents.core.agent.react.ReActAgentState;
import com.easyagents.core.agent.react.ReActStep;
import com.easyagents.core.memory.ChatMemory;
import com.easyagents.core.message.ToolCall;
import com.easyagents.core.message.ToolMessage;
import com.easyagents.core.message.UserMessage;
import com.easyagents.core.model.chat.ChatModel;
import com.easyagents.core.model.chat.ChatOptions;
import com.easyagents.core.model.chat.StreamResponseListener;
import com.easyagents.core.model.chat.response.AiMessageResponse;
import com.easyagents.core.model.chat.tool.Tool;
import com.easyagents.core.model.chat.tool.ToolScanner;
import com.easyagents.core.model.client.StreamContext;
import com.easyagents.core.model.exception.ModelException;
import com.easyagents.core.prompt.SimplePrompt;
import org.junit.Test;

import java.util.List;
import java.util.concurrent.TimeUnit;

public class OpenAIChatModelTest {

    @Test(expected = ModelException.class)
    public void testChat() {

        String output = OpenAIChatConfig.builder()
            .endpoint("https://ai.gitee.com")
            .provider("GiteeAI")
            .model("Qwen3-32B")
            .apiKey("PXW1****D12")
            .buildModel()
            .chat("你叫什么名字");

        System.out.println(output);
    }

    @Test()
    public void testChatStream() {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setApiKey("PXW1GXE***");
        config.setEndpoint("https://ai.gitee.com");
        config.setModel("Qwen3-32B");
        config.setLogEnabled(true);

        ChatOptions options = ChatOptions.builder().thinkingEnabled(false).build();

        ChatModel chatModel = new OpenAIChatModel(config);

        ChatModelTestUtils.waitForStream(chatModel, "你叫什么名字", new StreamResponseListener() {
            @Override
            public void onMessage(StreamContext context, AiMessageResponse response) {
                System.out.println(response.getMessage().getContent());
            }

            @Override
            public void onFailure(StreamContext context, Throwable throwable) {
                System.out.println("onFailure>>>>" + throwable);
            }

            @Override
            public void onStop(StreamContext context) {
                System.out.println("stop!!!!");
            }
        }, options);

    }


    @Test()
    public void testChatStreamBailian() {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setApiKey("sk-32ab******57502");
        config.setEndpoint("https://dashscope.aliyuncs.com");
        config.setRequestPath("/compatible-mode/v1/chat/completions");
        config.setModel("qwen3-max");
        ChatModel chatModel = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("北京的天气如何？");
        prompt.addToolsFromClass(WeatherFunctions.class);

        ChatModelTestUtils.waitForStream(chatModel, prompt, new StreamResponseListener() {
            @Override
            public void onFailure(StreamContext context, Throwable throwable) {
                System.out.println("onFailure>>>>" + throwable);
            }

            @Override
            public void onMessage(StreamContext context, AiMessageResponse response) {
                if (response.getMessage().getContent() == null) {
                    System.out.println(response.getMessage());
                }

                if (response.getMessage().isFinalDelta()) {
                    List<ToolCall> toolCalls = response.getMessage().getToolCalls();
                    System.out.println(toolCalls);
                }

                System.out.println("onMessage >>>>>" + response.getMessage().getContent());
            }

            @Override
            public void onStop(StreamContext context) {
                System.out.println("stop!!!!");
            }
        });

    }

    @Test
    public void testChatOllama() {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setEndpoint("http://localhost:11434");
        config.setModel("llama3");
        config.setLogEnabled(true);

        ChatModel chatModel = new OpenAIChatModel(config);
        chatModel.chatStream("who are you", new StreamResponseListener() {
            @Override
            public void onMessage(StreamContext context, AiMessageResponse response) {
                System.out.println(response.getMessage().getContent());
            }

            @Override
            public void onStop(StreamContext context) {
                System.out.println("stop!!!!");
            }
        });

//        try {
//            Thread.sleep(2000);
//        } catch (InterruptedException e) {
//            throw new RuntimeException(e);
//        }
    }


    @Test()
    public void testChatWithImage() {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setApiKey("sk-5gqOcl*****");
        config.setModel("gpt-4-turbo");


        ChatModel chatModel = new OpenAIChatModel(config);
        SimplePrompt prompt = new SimplePrompt("What's in this image?");
        prompt.addImageUrl("https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg");


        AiMessageResponse response = chatModel.chat(prompt);
        System.out.println(response);
    }


    @Test()
    public void testFunctionCalling1() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setApiKey("sk-rts5NF6n*******");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("今天北京的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);
        AiMessageResponse response = llm.chat(prompt);

        System.out.println(response.executeToolCallsAndGetResults());
        // 阴转多云
    }

    @Test()
    public void testFunctionCalling2() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setApiKey("sk-rts5NF6n*******");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("今天北京的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);
        AiMessageResponse response = llm.chat(prompt);

        if (response.hasToolCalls()) {
            prompt.setToolMessages(response.executeToolCallsAndGetToolMessages());
            AiMessageResponse response1 = llm.chat(prompt);
            System.out.println(response1.getMessage().getContent());
        } else {
            System.out.println(response);
        }
    }

    @Test()
    public void testFunctionCalling3() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setLogEnabled(true);
        config.setEndpoint("https://ark.cn-beijing.volces.com");
        config.setRequestPath("/api/v3/chat/completions");
        config.setModel("deepseek-v3-250324");
        config.setApiKey("2d57a");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("今天北京的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);
        AiMessageResponse response = llm.chat(prompt);

        if (response.hasToolCalls()) {
            prompt.setToolMessages(response.executeToolCallsAndGetToolMessages());
            AiMessageResponse response1 = llm.chat(prompt);
            System.out.println(response1.getMessage().getContent());
        } else {
            System.out.println(response);
        }
    }

    @Test()
    public void testFunctionCalling4() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setLogEnabled(true);
        config.setEndpoint("https://ark.cn-beijing.volces.com");
        config.setRequestPath("/api/v3/chat/completions");
        config.setModel("deepseek-v3-250324");
        config.setApiKey("2d57aa75");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("今天北京的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);
        llm.chatStream(prompt, new StreamResponseListener() {
            @Override
            public void onMessage(StreamContext context, AiMessageResponse response) {
                System.out.println(" onMessage >>>>>" + response.hasToolCalls());
            }
        });

        TimeUnit.SECONDS.sleep(5);
    }

    @Test()
    public void testFunctionCalling44() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setLogEnabled(true);
        config.setEndpoint("https://ark.cn-beijing.volces.com");
        config.setRequestPath("/api/v3/chat/completions");
        config.setModel("deepseek-v3-250324");
        config.setApiKey("2d5");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("今天北京的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);
        llm.chatStream(prompt, new StreamResponseListener() {
            @Override
            public void onMessage(StreamContext context, AiMessageResponse response) {
                System.out.println(" onMessage >>>>>" + response.hasToolCalls());
            }
        });

        TimeUnit.SECONDS.sleep(5);
    }

    @Test()
    public void testFunctionCalling444() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setLogEnabled(true);
        config.setEndpoint("https://ai.gitee.com");
//        config.setRequestPath("/api/v3/chat/completions");
//        config.setModel("Qwen3-32B");
        config.setModel("DeepSeek-V3.2");
        config.setApiKey("PXW1G***L7D12");
//        config.setLogEnabled(false);

        OpenAIChatModel llm = new OpenAIChatModel(config);


        SimplePrompt prompt = new SimplePrompt("北京和上海的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);
        llm.chatStream(prompt, new StreamResponseListener() {
            @Override
            public void onMessage(StreamContext context, AiMessageResponse response) {

//                System.out.println("onMessage11 >>>>>" + response);
                if (response.getMessage().isFinalDelta() && response.hasToolCalls()) {
                    System.out.println(":::::::: start....");
                    List<ToolMessage> toolMessages = response.executeToolCallsAndGetToolMessages();
                    prompt.setAiMessage(response.getMessage());
                    prompt.setToolMessages(toolMessages);
                    llm.chatStream(prompt, new StreamResponseListener() {
                        @Override
                        public void onMessage(StreamContext context, AiMessageResponse response) {
                            String msg = response.getMessage().getContent() != null ? response.getMessage().getContent() : response.getMessage().getReasoningContent();
                            System.out.println(":::22" + msg);
                        }

                        @Override
                        public void onStop(StreamContext context) {
                            System.out.println("onStop >>>>>");
                        }
                    });
                } else {
                    String msg = response.getMessage().getContent() != null ? response.getMessage().getContent() : response.getMessage().getReasoningContent();
                    System.out.println(">>>" + msg);
                }
            }
        });

        TimeUnit.SECONDS.sleep(25);

    }


    @Test()
    public void testFunctionCalling5() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setLogEnabled(true);
        config.setEndpoint("https://ai.gitee.com");
        config.setModel("Qwen3-32B");
        config.setApiKey("PXW1G*********D12");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("北京和上海的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);
        llm.chatStream(prompt, new StreamResponseListener() {
            @Override
            public void onMessage(StreamContext context, AiMessageResponse response) {
//                System.out.println("onMessage >>>>>" + response);
            }
        });

        TimeUnit.SECONDS.sleep(25);
    }


    @Test()
    public void testFunctionCalling55() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setLogEnabled(true);
        config.setEndpoint("https://ai.gitee.com");
        config.setModel("Qwen3-32B");
//        config.setModel("DeepSeek-V3");
//        config.setSupportToolMessage(false);
        config.setApiKey("PXW1");


        OpenAIChatModel llm = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("/no_think 北京和上海的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);


        llm.chatStream(prompt, new StreamResponseListener() {
            @Override
            public void onMessage(StreamContext context, AiMessageResponse response) {
                if (response.getMessage().isFinalDelta() && response.hasToolCalls()) {
                    System.out.println(":::::::: start....");
                    prompt.setAiMessage(response.getMessage());
                    prompt.setToolMessages(response.executeToolCallsAndGetToolMessages());
                    llm.chatStream(prompt, new StreamResponseListener() {
                        @Override
                        public void onMessage(StreamContext context, AiMessageResponse response) {
                            String msg = response.getMessage().getContent() != null ? response.getMessage().getContent() : response.getMessage().getReasoningContent();
                            System.out.println(":::" + msg);
                        }
                    });
                } else {
                    String msg = response.getMessage().getContent() != null ? response.getMessage().getContent() : response.getMessage().getReasoningContent();
                    System.out.println(">>>" + msg);
                }
            }
        });

        TimeUnit.SECONDS.sleep(25);
    }

    @Test()
    public void testFunctionCalling6() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
        config.setLogEnabled(true);
        config.setEndpoint("https://ai.gitee.com");
        config.setModel("Qwen3-32B");
        config.setApiKey("PXW1");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        SimplePrompt prompt = new SimplePrompt("/no_think 北京和上海的天气怎么样");
        prompt.addToolsFromClass(WeatherFunctions.class);
        AiMessageResponse response = llm.chat(prompt);

        prompt.setToolMessages(response.executeToolCallsAndGetToolMessages());

        System.out.println(llm.chat(prompt));

    }


    @Test()
    public void testReAct1() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
//        config.setDebug(true);
        config.setEndpoint("https://ai.gitee.com");
        config.setModel("Qwen3-32B");
        config.setApiKey("****");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        List<Tool> tools = ToolScanner.scan(WeatherFunctions.class);
//        ReActAgent reActAgent = new ReActAgent(llm, functions, "北京和上海的天气怎么样？");
        ReActAgent reActAgent = new ReActAgent(llm, tools, "介绍一下北京");
        reActAgent.addListener(new ReActAgentListener() {

            @Override
            public void onActionStart(ReActStep step) {
                System.out.println(">>>>>>" + step.getThought());
                System.out.println("正在调用工具 >>>>> " + step.getAction() + ":" + step.getActionInput());
            }

            @Override
            public void onActionEnd(ReActStep step, Object result) {
                System.out.println("工具调用结束 >>>>> " + step.getAction() + ":" + step.getActionInput() + ">>>>结果：" + result);
            }

            @Override
            public void onFinalAnswer(String finalAnswer) {
                System.out.println("onFinalAnswer >>>>>" + finalAnswer);
            }

            @Override
            public void onNonActionResponse(AiMessageResponse response) {
                System.out.println("onNonActionResponse >>>>>" + response.getMessage().getContent());
            }
        });

        reActAgent.execute();
    }


    @Test()
    public void testReAct2() throws InterruptedException {
        OpenAIChatConfig config = new OpenAIChatConfig();
//        config.setDebug(true);
        config.setEndpoint("https://ai.gitee.com");
        config.setModel("Qwen2-72B-Instruct");
        config.setApiKey("*****");

        OpenAIChatModel llm = new OpenAIChatModel(config);

        List<Tool> tools = ToolScanner.scan(WeatherFunctions.class);
        ReActAgent reActAgent = new ReActAgent(llm, tools, "今天的天气怎么样？");
//        reActAgent.setStreamable(true);
        reActAgent.addListener(new ReActAgentListener() {

            @Override
            public void onChatResponseStream(StreamContext context, AiMessageResponse response) {
//                System.out.print(response.getMessage().getContent());
            }

            @Override
            public void onRequestUserInput(String question) {
                System.out.println("onRequestUserInput>>>" + question);

                ReActAgentState state = reActAgent.getState();
                state.addMessage(new UserMessage("我在北京市"));
                ReActAgent newAgent = new ReActAgent(llm, tools, state);
                newAgent.addListener(this);
                newAgent.execute();
            }

            @Override
            public void onActionStart(ReActStep step) {
                System.out.println(">>>>>>" + step.getThought());
                System.out.println("正在调用工具 >>>>> " + step.getAction() + ":" + step.getActionInput());
            }

            @Override
            public void onActionEnd(ReActStep step, Object result) {
                System.out.println("工具调用结束 >>>>> " + step.getAction() + ":" + step.getActionInput() + ">>>>结果：" + result);
            }

            @Override
            public void onFinalAnswer(String finalAnswer) {
                System.out.println("onFinalAnswer >>>>>" + finalAnswer);
                ChatMemory memory = reActAgent.getMemoryPrompt().getMemory();
                System.out.println(memory);
            }

            @Override
            public void onNonActionResponseStream(StreamContext context) {
                System.out.println("onNonActionResponseStream >>>>>" + context);
            }
        });

        reActAgent.execute();

        TimeUnit.SECONDS.sleep(20);
    }


}
