package com.infoepoch.pms.agent.domain.care.orchestrator;

import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.alibaba.cloud.ai.graph.streaming.OutputType;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.infoepoch.pms.agent.common.utils.LogHelper;
import com.infoepoch.pms.agent.domain.care.log.CareTraceLogSupport;
import com.infoepoch.pms.agent.domain.care.model.CareQuery;
import com.infoepoch.pms.agent.domain.care.stream.CareStreamingResponseAssembler;
import com.infoepoch.pms.agent.observability.ai.AiComponentType;
import com.infoepoch.pms.agent.observability.ai.AiObservationContext;
import com.infoepoch.pms.agent.observability.ai.AiObservationSupport;
import org.springframework.ai.chat.messages.AbstractMessage;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * 通用对话执行器
 */
@Component
public class GeneralConversationExecutor {

    private final CareStreamingResponseAssembler responseAssembler;
    private final ReactAgent careAgent;
    private final AiObservationSupport observationSupport;

    public GeneralConversationExecutor(CareStreamingResponseAssembler responseAssembler,
                                       @Qualifier("CareAgent") ReactAgent careAgent,
                                       AiObservationSupport observationSupport) {
        this.responseAssembler = responseAssembler;
        this.careAgent = careAgent;
        this.observationSupport = observationSupport;
    }

    /**
     * 执行通用对话。
     */
    public Flux<String> execute(String traceId, String sessionId, CareQuery query) {
        String prompt = query.originalMessage();
        LogHelper.info(this, CareTraceLogSupport.format(
                traceId,
                sessionId,
                "通用对话开始",
                "用户输入=%s".formatted(CareTraceLogSupport.safeMessageSummary(query.originalMessage()))
        ));
        RunnableConfig runnableConfig = buildRunnableConfig(traceId, sessionId, "通用对话");
        try {
            AtomicInteger chunkIndex = new AtomicInteger(0);
            return careAgent.stream(prompt, runnableConfig)
                    .filter(StreamingOutput.class::isInstance)
                    .map(StreamingOutput.class::cast)
                    .filter(output -> output.getOutputType() == OutputType.AGENT_MODEL_STREAMING)
                    .map(StreamingOutput::message)
                    .filter(AssistantMessage.class::isInstance)
                    .map(AssistantMessage.class::cast)
                    .map(AbstractMessage::getText)
                    .filter(StringUtils::hasText)
                    .doOnNext(chunk -> chunkIndex.incrementAndGet())
                    .switchIfEmpty(Flux.defer(() -> {
                        LogHelper.info(this, CareTraceLogSupport.format(
                                traceId,
                                sessionId,
                                "通用对话兜底",
                                "模型未返回有效内容"
                        ));
                        return Flux.just(responseAssembler.buildGeneralFallbackReply());
                    }))
                    .doOnComplete(() -> {
                        LogHelper.info(this, CareTraceLogSupport.format(
                                traceId,
                                sessionId,
                                "通用对话完成",
                                "通用对话处理完成"
                        ));
                    })
                    .doOnError(error -> LogHelper.error(this, CareTraceLogSupport.format(
                            traceId,
                            sessionId,
                            "流式请求异常",
                            "通用对话失败，异常=" + CareTraceLogSupport.safeText(error.getMessage()) +
                                    "，已接收chunk数=%s".formatted(chunkIndex.get())
                    ), error));
        } catch (Exception e) {
            LogHelper.error(this, CareTraceLogSupport.format(
                    traceId,
                    sessionId,
                    "流式请求异常",
                    "通用对话立即失败，异常=" + CareTraceLogSupport.safeText(e.getMessage())
            ), e);
            return Flux.error(new IllegalArgumentException("调用模型失败"));
        }
    }

    private RunnableConfig buildRunnableConfig(String traceId, String sessionId, String scene) {
        AiObservationContext currentContext = observationSupport.currentContext();
        String requestId = StringUtils.hasText(currentContext.requestId()) ? currentContext.requestId() : traceId;
        String flowId = StringUtils.hasText(currentContext.flowId()) ? currentContext.flowId() : requestId;
        AiObservationContext context = currentContext.toBuilder()
                .traceId(traceId)
                .requestId(requestId)
                .sessionId(sessionId)
                .flowId(flowId)
                .scene(scene)
                .componentType(AiComponentType.AGENT)
                .componentName("careAgent")
                .agentName("careAgent")
                .build();
        return observationSupport.buildRunnableConfig(sessionId, context);
    }
}
