package cn.com.poc.thirdparty.resource.demand.ai.function.web_seach;

import cn.com.poc.agent_application.entity.KnowledgeContentResult;
import cn.com.poc.agent_application.entity.Variable;
import cn.com.poc.common.utils.DocumentLoad;
import cn.com.poc.common.utils.JsonUtils;
import cn.com.poc.common.utils.StringUtils;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dbchain.DBChainResult;
import cn.com.poc.thirdparty.resource.demand.ai.function.AbstractFunctionResult;
import cn.com.poc.thirdparty.resource.demand.ai.function.AbstractLargeModelFunction;
import cn.com.poc.thirdparty.resource.demand.ai.function.entity.FunctionLLMConfig;
import cn.com.poc.thirdparty.resource.demand.ai.function.entity.Parameters;
import cn.com.poc.thirdparty.resource.demand.ai.function.entity.Properties;
import cn.com.yict.framemax.core.config.Config;
import cn.com.yict.framemax.core.exception.BusinessException;
import cn.hutool.core.collection.ListUtil;
import com.alibaba.fastjson.JSONObject;
import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.http.apache.v2.ApacheHttpTransport;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.gson.GsonFactory;
import com.google.api.services.customsearch.v1.CustomSearchAPI;
import com.google.api.services.customsearch.v1.CustomSearchAPIRequestInitializer;
import com.google.api.services.customsearch.v1.model.Result;
import com.google.api.services.customsearch.v1.model.Search;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.concurrent.*;

/**
 * @author alex.yao
 * @date 2025/1/14
 */
@Component
public class WebSearchFunction extends AbstractLargeModelFunction {

    private Logger logger = LoggerFactory.getLogger(WebSearchFunction.class);

    private CustomSearchAPI customSearchAPI;
    private final HttpTransport transport = new ApacheHttpTransport();

    private final JsonFactory jsonFactory = new GsonFactory();

    private final String KEY = Config.get("google.custom.search.key");
    private final String CX = Config.get("google.custom.search.cx");

    private final String DESC = "该方法是使用Google搜索，对用户给出的关键词进行网页搜索检索并且获取网页内容";

    private final FunctionLLMConfig functionLLMConfig = new FunctionLLMConfig.FunctionLLMConfigBuilder()
            .name("web_search")
            .description(DESC)
            .parameters(new Parameters("object")
                    .addProperties("query", new Properties("string", "搜索关键词")))
            .build();

    private CustomSearchAPI getCustomSearchAPI() {
        try {
            if (customSearchAPI == null) {
                customSearchAPI = new CustomSearchAPI.Builder(transport, jsonFactory, GoogleNetHttpTransport.newTrustedTransport().createRequestFactory().getInitializer())
                        .setRootUrl("https://google-api.gsstcloud.com")
                        .setCustomSearchAPIRequestInitializer(new CustomSearchAPIRequestInitializer("AIzaSyCV8PTQ10rG5wo4E004dR3mcGD1RM_PrBw"))
                        .build();
            }
            return customSearchAPI;
        } catch (Exception e) {
            throw new BusinessException("初始化google搜索服务失败");
        }
    }

    @Override
    public AbstractFunctionResult<List<WebSearchFunctionResult>> doFunction(String content, String identifier, List<DBChainResult> dbChainResults, List<KnowledgeContentResult> knowledgeContentResults) {
        AbstractFunctionResult<List<WebSearchFunctionResult>> result = new AbstractFunctionResult<>();
        if (StringUtils.isBlank(content)) {
            result.setPromptContent("FAIL");
            return result;
        }
        JSONObject jsonObject = JSONObject.parseObject(content);
        String query = jsonObject.getString("query");
        query = query.replaceAll(StringUtils.SPACE, StringUtils.EMPTY);
        List<WebSearchFunctionResult> results = new CopyOnWriteArrayList<>();
        try {
            Search execute = getCustomSearchAPI()
                    .cse().list().setCx(CX)
                    .setKey(KEY)
                    .setFileType("pdf,html")
                    .setQ(query)
                    .setStart(1L)
                    .setNum(3)
                    .execute();
            loadHTML(execute, results);
            logger.info("web search result：{}", results);
            StringBuilder stringBuilder = new StringBuilder();
            int i = 0;
            for (WebSearchFunctionResult webSearchFunctionResult : results) {
                stringBuilder.append(i).append(":").append(webSearchFunctionResult.getHtmlContent()).append("\n");
            }
            result.setPromptContent(stringBuilder.toString());
            result.setFunctionResult(results);
            return result;
        } catch (Exception e) {
            logger.error("web search error:" + e.getMessage());
            result.setPromptContent("FAIL");
            return result;
        }
    }

    private void loadHTML(Search execute, List<WebSearchFunctionResult> results) {
        List<Result> items = execute.getItems();
        try {
            ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(3, 3, 10, TimeUnit.SECONDS, new LinkedBlockingQueue<>(3));
            CountDownLatch countDownLatch = new CountDownLatch(items.size());
            for (Result item : items) {
                threadPoolExecutor.submit(() -> {
                    String link = item.getLink();
                    String htmlContent = StringUtils.EMPTY;
                    if (link.endsWith("pdf")) {
                        try {
                            File file = DocumentLoad.downloadURLDocument(link);
                            htmlContent = DocumentLoad.loadPDF(file);
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    } else {
                        htmlContent = DocumentLoad.htmlToMarkdown(link);
                    }
                    if (StringUtils.isNotBlank(htmlContent)) {
                        htmlContent = htmlContent.replaceAll(StringUtils.SPACE, StringUtils.EMPTY);
                        if (htmlContent.length() > 1500) {
                            int index = htmlContent.length() / 2;
                            htmlContent = htmlContent.substring(index - 725, index + 725);
                        }
                        WebSearchFunctionResult webSearchResult = new WebSearchFunctionResult();
                        webSearchResult.setTitle(item.getTitle());
                        webSearchResult.setUrl(item.getLink());
                        webSearchResult.setContent(item.getSnippet());
                        webSearchResult.setHtmlContent(htmlContent);
                        results.add(webSearchResult);
                    }
                    countDownLatch.countDown();
                });
            }
            countDownLatch.await(30, TimeUnit.SECONDS);
            threadPoolExecutor.shutdown();
        } catch (Exception e) {
            logger.error("web search error:" + e.getMessage());
        }
    }

    @Override
    public String getDesc() {
        return DESC;
    }

    @Override
    public List<String> getLLMConfig() {
        return ListUtil.toList(JsonUtils.serialize(functionLLMConfig));
    }

    @Override
    public List<String> getLLMConfig(List<Variable> variableStructure) {
        return this.getLLMConfig();
    }
}
