package com.hanserwei.snailsai.controller;

import com.hanserwei.snailsai.model.AIResponse;
import com.hanserwei.snailsai.dto.ChatMessageDTO;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

@Slf4j
@RequestMapping("/dashscope")
@RestController
@CrossOrigin
public class DashScopeController {

    @Resource
    private ChatClient dashScopeChatClient;

    @PostMapping(value = "/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<AIResponse> generateStream(@RequestBody ChatMessageDTO chatMessageDTO) {
        // 构建提示词
        Prompt prompt = new Prompt(new UserMessage(chatMessageDTO.getMessage()));

        // 流式输出
        return dashScopeChatClient.prompt(prompt)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatMessageDTO.getConversionId()))
                .stream() // 流式输出
                .chatResponse()
                .mapNotNull(chatResponse -> {
                    Generation generation = chatResponse.getResult();
                    String text = generation.getOutput().getText();
                    return AIResponse.builder().v(text).build();
                });
    }
}