package space.anyi.springAiAlibabaLearn.controller;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.NodeOutput;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.function.Consumer;

@RestController
@RequestMapping("/graph")
public class GraphController {
    private static final Logger log = LoggerFactory.getLogger(GraphController.class);
    private final CompiledGraph compiledGraph;
    private final CompiledGraph createStanceGraph;
    private final CompiledGraph conditionGraph;
    private final CompiledGraph loopGraph;
    private final CompiledGraph saveGraph;

    public GraphController(@Autowired @Qualifier("compiledGraph")CompiledGraph compiledGraph,
                           @Autowired @Qualifier("createStanceGraph")CompiledGraph createStanceGraph,
                           @Autowired @Qualifier("conditionGraph")CompiledGraph conditionGraph,
                           @Autowired @Qualifier("loopGraph")CompiledGraph loopGraph,
                           @Autowired @Qualifier("saveGraph")CompiledGraph saveGraph) {
        this.compiledGraph = compiledGraph;
        this.createStanceGraph = createStanceGraph;
        this.conditionGraph = conditionGraph;
        this.loopGraph = loopGraph;
        this.saveGraph = saveGraph;
    }

    @GetMapping("/quickStart")
    public Flux<NodeOutput> quickStart(){
        return compiledGraph.stream().log().doOnNext(new Consumer<NodeOutput>() {
            @Override
            public void accept(NodeOutput nodeOutput) {
                log.debug("nodeOutput:{}",nodeOutput);
                //获取工作流中存储的数据
                Map<String, Object> data = nodeOutput.state().data();
                log.debug("data:{}",data);
            }
        });
    }
    @GetMapping("/createStance")
    public Map<String, Object> createStance(@RequestParam("word")String word){
        return createStanceGraph.stream(Map.of("word", word)).log().last().block().state().data();
    }
    @GetMapping("/condition")
    public Map<String, Object> condition(@RequestParam("topic")String topic){
        return conditionGraph.stream(Map.of("topic", topic)).log().last().block().state().data();
    }
    @GetMapping("/loop")
    public Map<String, Object> loop(@RequestParam("topic")String topic){
        return loopGraph.stream(Map.of("topic", topic)).log().last().block().state().data();
    }
    @GetMapping("/saveTest")
    public Map<String, Object> saveTest(@RequestParam("msg")String msg,@RequestParam("id") String id){
        //graph工作流的运行配置
        RunnableConfig runnableConfig = RunnableConfig.builder()
                .threadId(id)
                .build();
        return saveGraph.stream(Map.of("msg", msg),runnableConfig)
                .log()
                .last()
                .block()
                .state()
                .data();
    }
}