package com.sonic.cow.client.impl;

import com.google.common.collect.Lists;
import com.sonic.cow.client.ChatStructuredOutputsClient;
import com.byteplus.ark.runtime.model.completion.chat.ChatCompletionContentPart;
import com.byteplus.ark.runtime.model.completion.chat.ChatCompletionRequest;
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 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.ArrayList;
import java.util.List;

@Slf4j
@Service
public class ChatStructuredOutputsClientImpl implements ChatStructuredOutputsClient {

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

    @Autowired
    private ArkService volcengineArkService;

    @Override
    public String chatStructuredOutputs(String message, ChatCompletionRequest.ChatCompletionRequestResponseFormat responseFormat) {
        List<ChatMessage> messages = Lists.newArrayList();
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(message)
                .build());
        //在结构化输出中，如果关闭了推理的话是无法结构化输出内容的
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(contextCachingModel)
                .messages(messages)
                .responseFormat(responseFormat)
                .build();
        StringBuffer stringBuffer = new StringBuffer();
        Long startTime = System.currentTimeMillis();
        volcengineArkService.createChatCompletion(chatCompletionRequest).getChoices().forEach(
                choice -> {
                    stringBuffer.append(choice.getMessage().getContent());
                }
        );
        Long endTime = System.currentTimeMillis();
        log.info("chatStructuredOutputs 所需要的时间：{}" , (endTime - startTime) / 1000);
        return stringBuffer.toString();
    }

    @Override
    public String aiGenStructuredOutputs(String systemPrompt, Double temperature, ChatCompletionRequest.ChatCompletionRequestResponseFormat responseFormat, Boolean thinking) {
        log.info("aiGenStructuredOutputs systemPrompt : {}", systemPrompt);
        List<ChatMessage> messages = Lists.newArrayList();
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(systemPrompt)
                .build());
        //在结构化输出中，如果关闭了推理的话是无法结构化输出内容的
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(contextCachingModel)
                .messages(messages)
                .temperature(temperature)
                .responseFormat(responseFormat)
//                .thinking(thinking != null && thinking ? new ChatCompletionRequest.ChatCompletionRequestThinking("enabled") : new ChatCompletionRequest.ChatCompletionRequestThinking("disabled"))
                .build();
        StringBuffer stringBuffer = new StringBuffer();

        Long startTime = System.currentTimeMillis();
        volcengineArkService.createChatCompletion(chatCompletionRequest).getChoices().forEach(
                choice -> {
                    stringBuffer.append(choice.getMessage().getContent());
                    log.info("===> ReasoningContent : {}", choice.getMessage().getReasoningContent());
                }
        );
        Long endTime = System.currentTimeMillis();
        log.info("aiGenStructuredOutputs time : {}", (endTime - startTime) / 1000);
        return stringBuffer.toString();
    }

    @Override
    public String imageGenText(String systemPrompt, Double temperature, String imageUrl, ChatCompletionRequest.ChatCompletionRequestResponseFormat responseFormat) {
        log.info("图生文系统提示词：{}" , systemPrompt);
        final List<ChatMessage> messages = new ArrayList<>();
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content(systemPrompt).build());
        final List<ChatCompletionContentPart> multiParts = new ArrayList<>();
        multiParts.add(ChatCompletionContentPart.builder().type("image_url").imageUrl(
                new ChatCompletionContentPart.ChatCompletionContentPartImageURL(
                        imageUrl
                )
        ).build());
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER)
                .multiContent(multiParts).build();
        messages.add(userMessage);

        Long startTime = System.currentTimeMillis();
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(contextCachingModel)
                .messages(messages)
//                .thinking(new ChatCompletionRequest.ChatCompletionRequestThinking("disabled"))
                .temperature(temperature == null ? 0.5 : temperature)
                .responseFormat(responseFormat)
                .build();

        StringBuffer stringBuffer = new StringBuffer();
        volcengineArkService.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> stringBuffer.append(choice.getMessage().getContent()));

        Long endTime = System.currentTimeMillis();
        log.info("图生文所需要的时间：{}" , (endTime - startTime) / 1000);
        return stringBuffer.toString();
    }
}
