package com.lideeyunji.tool.framework.ai.websearch;


import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.lideeyunji.tool.framework.ai.chat.ChatBaseRequest;
import com.lideeyunji.tool.framework.ai.chat.ChatBaseResponse;
import com.lideeyunji.tool.framework.ai.service.Configuration;
import com.lideeyunji.tool.framework.ai.service.IChatService;
import com.lideeyunji.tool.framework.ai.websearch.searxng.SearXNGConfig;
import com.lideeyunji.tool.framework.ai.websearch.searxng.SearXNGRequest;
import com.lideeyunji.tool.framework.ai.websearch.searxng.SearXNGResponse;
import com.lideeyunji.tool.framework.exception.lideeYunJiException;
import com.lideeyunji.tool.framework.yunji.utils.FuncBase;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;


public class ChatWithWebSearchEnhance implements IChatService {
    private final IChatService chatService;
    private final SearXNGConfig searXNGConfig;
    private final OkHttpClient okHttpClient;
    public ChatWithWebSearchEnhance(IChatService chatService, Configuration configuration) {
        this.chatService = chatService;
        this.searXNGConfig = configuration.getSearXNGConfig();
        this.okHttpClient = configuration.getOkHttpClient();
    }

    @Override
    public ChatBaseResponse chatCompletion(String baseUrl, String apiKey, ChatBaseRequest chatBaseRequest) throws Exception {
        return chatService.chatCompletion(baseUrl, apiKey, addWebSearchResults(chatBaseRequest));
    }

    @Override
    public ChatBaseResponse chatCompletion(ChatBaseRequest chatBaseRequest) throws Exception {
        return chatService.chatCompletion(addWebSearchResults(chatBaseRequest));
    }



    private ChatBaseRequest addWebSearchResults(ChatBaseRequest chatBaseRequest) {
        int chatLen = chatBaseRequest.getMessages().size();
        String prompt = chatBaseRequest.getMessages().get(chatLen - 1).getContent();
        // 执行联网搜索并将结果附加到提示词中
        String searchResults = performWebSearch(prompt);
        chatBaseRequest.getMessages().get(chatLen - 1).setContent("我将提供一段来自互联网的资料信息, 请根据这段资料以及用户提出的问题来给出回答。请确保在回答中使用Markdown格式，并在回答末尾列出参考资料。如果资料中的信息不足以回答用户的问题，可以根据自身知识库进行补充，或者说明无法提供确切的答案。\n" +
                "网络资料:\n"
                + "============\n"
                + searchResults
                + "============\n"
                + "用户问题:\n"
                + "============\n"
                + prompt
                + "============\n");
        return chatBaseRequest;
    }

    private String performWebSearch(String query) {

        SearXNGRequest searXNGRequest = SearXNGRequest.builder()
                .q(query)
                .engines(searXNGConfig.getEngines())
                .build();


        if(StringUtils.isBlank(searXNGConfig.getUrl())){
            throw new lideeYunJiException("SearXNG url is not configured");
        }


        Request request = new Request.Builder()
                .url(searXNGConfig.getUrl()+ "?format=json&q=" + query + "&engines=" + searXNGConfig.getEngines())
                .get()
                .build();


        try(Response execute = okHttpClient.newCall(request).execute()) {
            if (execute.isSuccessful() && execute.body() != null){
                SearXNGResponse searXNGResponse = FuncBase.json2Bean(execute.body().string(), SearXNGResponse.class);

                if(searXNGResponse.getResults().size() > searXNGConfig.getNums()) {
                    return FuncBase.json2Str(searXNGResponse.getResults().subList(0, searXNGConfig.getNums()));
                }
                return FuncBase.json2Str(searXNGResponse.getResults());


            }else{
                throw new Exception("SearXNG request failed");
            }


        } catch (Exception e) {
            throw new lideeYunJiException("SearXNG request failed");
        }


    }
}
