package cn.com.poc.thirdparty.resource.demand.ai.aggregate.impl;

import cn.com.poc.thirdparty.resource.demand.ai.aggregate.AIDialogueService;
import cn.com.poc.thirdparty.resource.demand.ai.entity.csvchain.CSVChainResponse;
import cn.com.poc.thirdparty.resource.demand.ai.entity.csvchain.CSVChainResult;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dbchain.DBChainResponse;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dbchain.DBChainResult;
import cn.com.poc.thirdparty.resource.demand.ai.entity.function.FunctionCallResponse;
import cn.com.poc.thirdparty.resource.demand.ai.entity.function.FunctionCallResult;
import cn.com.poc.thirdparty.resource.demand.ai.entity.largemodel.LargeModelDemandResponse;
import cn.com.poc.thirdparty.resource.demand.ai.entity.largemodel.LargeModelDemandResult;
import cn.com.poc.thirdparty.resource.demand.ai.entity.largemodel.LargeModelResponse;
import cn.com.poc.thirdparty.resource.demand.ai.route.DgtoolsApiRoute;
import cn.com.poc.thirdparty.resource.demand.dgTools.DgtoolsAbstractHttpClient;
import cn.com.poc.thirdparty.resource.demand.member.service.DemandAuthService;
import org.apache.http.Header;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.RequestBuilder;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.message.BasicHeader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

/**
 * @author alex.yao
 * @date 2023/5/24
 **/

@Component
public class AIDialogueServiceImpl implements AIDialogueService {
    private final String POST = "POST";

    private final Logger logger = LoggerFactory.getLogger(AIDialogueService.class);

    @Value("${dgtools.domain.url}")
    private String DOMAIN_URL;

    @Value(value = "${large-model.apikey}")
    private String API_KEY;

    @Resource
    private DemandAuthService demandAuthService;

    @Resource
    private DgtoolsAbstractHttpClient dgToolsAbstractHttpClient;

    @Override
    public BufferedReader polyStream(LargeModelResponse largeResponse) throws Exception {
        //获取one-api key
        largeResponse.setStream(true);
        LargeModelDemandResponse response = new LargeModelDemandResponse();
        response.setApiKey(API_KEY);
        response.setLargeResponse(largeResponse);
        return largeModelStream(response);
    }

    @Override
    public LargeModelDemandResult poly(LargeModelResponse largeResponse) {
        largeResponse.setStream(false);
        LargeModelDemandResponse response = new LargeModelDemandResponse();
        response.setApiKey(API_KEY);
        response.setLargeResponse(largeResponse);
        return largeModelRequest(response);
    }

    @Override
    public FunctionCallResult functionCall(FunctionCallResponse response) {
        String url = DgtoolsApiRoute.DgtoolsAI.FUNCTION_CALL;
        response.setApiKey(API_KEY);
        List<Header> headers = new ArrayList<Header>() {{
            add(DgtoolsApiRoute.JSON_HEADER);
            add(DgtoolsApiRoute.AI_HEADER);
            add(new BasicHeader(DgtoolsApiRoute.HEADER_X_PLATFORM_AUTHORIZATION, demandAuthService.getToken()));
        }};
        return dgToolsAbstractHttpClient.doRequest(url, response, headers);
    }

    @Override
    public DBChainResult dbChain(DBChainResponse response) {
        String url = DgtoolsApiRoute.DgtoolsAI.DATABASE_CHAIN;
        response.setApiKey(API_KEY);
        List<Header> headers = new ArrayList<Header>() {{
            add(DgtoolsApiRoute.JSON_HEADER);
            add(DgtoolsApiRoute.AI_HEADER);
            add(new BasicHeader(DgtoolsApiRoute.HEADER_X_PLATFORM_AUTHORIZATION, demandAuthService.getToken()));
        }};
        return dgToolsAbstractHttpClient.doRequest(url, response, headers);
    }

    @Override
    public CSVChainResult csvChain(CSVChainResponse response) {
        String url = DgtoolsApiRoute.DgtoolsAI.CSV_CHAIN;
        response.setApiKey(API_KEY);
        List<Header> headers = new ArrayList<Header>() {{
            add(DgtoolsApiRoute.JSON_HEADER);
            add(DgtoolsApiRoute.AI_HEADER);
            add(new BasicHeader(DgtoolsApiRoute.HEADER_X_PLATFORM_AUTHORIZATION, demandAuthService.getToken()));
        }};
        return dgToolsAbstractHttpClient.doRequest(url, response, headers);
    }


    private BufferedReader largeModelStream(LargeModelDemandResponse request) throws IOException {
        String jsonBody = dgToolsAbstractHttpClient.buildJson(request);
        CloseableHttpClient httpClient = HttpClients.createDefault();
        CloseableHttpResponse httpResponse = httpClient.execute(RequestBuilder.create(POST)
                .setUri(DOMAIN_URL + DgtoolsApiRoute.BASE_URL + DgtoolsApiRoute.DgtoolsAI.LARGE_MODEL)
                .addHeader(DgtoolsApiRoute.JSON_HEADER)
                .addHeader(DgtoolsApiRoute.AI_HEADER)
                .addHeader(new BasicHeader(DgtoolsApiRoute.HEADER_X_PLATFORM_AUTHORIZATION, demandAuthService.getToken()))
                .setEntity(new StringEntity(jsonBody, StandardCharsets.UTF_8))
                .build()
        );
        InputStream inputStream = httpResponse.getEntity().getContent();
        return new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
    }

    private LargeModelDemandResult largeModelRequest(LargeModelDemandResponse request) {
        String url = DgtoolsApiRoute.DgtoolsAI.LARGE_MODEL;
        List<Header> headers = new ArrayList<Header>() {{
            add(DgtoolsApiRoute.JSON_HEADER);
            add(DgtoolsApiRoute.AI_HEADER);
            add(new BasicHeader(DgtoolsApiRoute.HEADER_X_PLATFORM_AUTHORIZATION, demandAuthService.getToken()));
        }};
        return dgToolsAbstractHttpClient.doRequest(url, request, headers);
    }


}
