package com.hanserwei.chat.controller;

import com.hanserwei.chat.model.dto.ChatMessageDTO;
import com.hanserwei.chat.model.vo.AIResponse;
import com.hanserwei.chat.reader.MyPdfReader;
import com.hanserwei.chat.utils.ConversationContext;
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.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.List;

@Slf4j
@RestController
@RequestMapping("/ai")
public class AiChatController {

    @Resource
    private ChatClient dashScopeChatClient;
    @Resource
    private MyPdfReader myPdfReader;
    @Resource
    private VectorStore vectorStore;

    @PostMapping(path = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<AIResponse> chatWithAi(@RequestBody ChatMessageDTO chatMessageDTO) {
        log.info("会话ID：{}", chatMessageDTO.getConversionId());
        ConversationContext.setConversationId(chatMessageDTO.getConversionId());

        return dashScopeChatClient.prompt()
                .user(chatMessageDTO.getMessage())
                .advisors(p -> p.param(ChatMemory.CONVERSATION_ID, chatMessageDTO.getConversionId()))
                .stream()
                .chatResponse()
                .mapNotNull(chatResponse -> AIResponse.builder()
                        .v(chatResponse.getResult().getOutput().getText())
                        .build())
                .contextWrite(ctx -> ConversationContext.withConversationId(chatMessageDTO.getConversionId()))
                .doFinally(signalType -> ConversationContext.clear());
    }

    @GetMapping("/readpdf")
    public String readPdf() {
        List<Document> docsFromPdf = myPdfReader.getDocsFromPdf();
        vectorStore.add(docsFromPdf);
        return "ok!";
    }
}