package com.sonic.cow.client.impl;

import com.byteplus.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.byteplus.ark.runtime.model.completion.chat.ChatCompletionResult;
import com.byteplus.ark.runtime.model.completion.chat.ChatMessage;
import com.byteplus.ark.runtime.model.completion.chat.ChatMessageRole;
import com.byteplus.ark.runtime.service.ArkService;
import com.sonic.cow.client.DeepSeekClient;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.Collections;

@Slf4j
@Service
public class DeepSeekClientImpl implements DeepSeekClient {

    @Value("${volcengine.deepSeekModel}")
    private String contextCachingModel;

    @Autowired
    private ArkService volcengineArkService;

    @Override
    public String aiGen(String systemPrompt, Double temperature) {
        log.info("===>DeepSeekClientImpl aiGen systemPrompt : {}", systemPrompt);
        Long startTime = System.currentTimeMillis();
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(contextCachingModel)
                .messages(Collections.singletonList(ChatMessage.builder()
                        .role(ChatMessageRole.SYSTEM)
                        .content(systemPrompt).build()))
                .temperature(temperature == null ? 0.5 : temperature)
                .build();
        ChatCompletionResult chatCompletionResult = volcengineArkService.createChatCompletion(chatCompletionRequest);
        StringBuffer stringBuffer = new StringBuffer();
        chatCompletionResult.getChoices().forEach(choice -> stringBuffer.append(choice.getMessage().getContent()));
        Long endTime = System.currentTimeMillis();
        log.info("===>DeepSeekClientImpl aiGen time : {}", (endTime - startTime) / 1000);
        return stringBuffer.toString();
    }

}
