package tech.easyflow.ai.easyagentsflow.event;

import com.alibaba.fastjson2.JSON;
import com.easyagents.flow.core.chain.repository.LoopInputReference;
import com.easyagents.flow.core.chain.repository.LoopResultReference;
import com.easyagents.flow.core.chain.repository.LoopResultRepository;
import org.springframework.dao.DuplicateKeyException;
import org.springframework.stereotype.Component;
import tech.easyflow.ai.entity.WorkflowExecResult;
import tech.easyflow.ai.entity.WorkflowExecStep;
import tech.easyflow.ai.service.WorkflowExecResultService;
import tech.easyflow.ai.service.WorkflowExecStepService;
import tech.easyflow.common.mq.core.MQConsumerHandler;
import tech.easyflow.common.mq.core.MQMessage;
import tech.easyflow.common.mq.core.MQSubscription;
import tech.easyflow.common.mq.config.MQProperties;

import java.util.List;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * 工作流执行审计事件消费者。
 */
@Component
public class WorkflowExecutionAuditConsumer implements MQConsumerHandler {

    private final WorkflowExecResultService workflowExecResultService;
    private final WorkflowExecStepService workflowExecStepService;
    private final MQProperties mqProperties;
    private final LoopResultRepository loopResultRepository;

    /**
     * 创建工作流执行审计事件消费者。
     *
     * @param workflowExecResultService 工作流执行记录服务
     * @param workflowExecStepService 节点执行步骤服务
     * @param mqProperties MQ 配置
     * @param loopResultRepository 循环与大型查询结果仓储
     */
    public WorkflowExecutionAuditConsumer(WorkflowExecResultService workflowExecResultService,
                                          WorkflowExecStepService workflowExecStepService,
                                          MQProperties mqProperties,
                                          LoopResultRepository loopResultRepository) {
        this.workflowExecResultService = workflowExecResultService;
        this.workflowExecStepService = workflowExecStepService;
        this.mqProperties = mqProperties;
        this.loopResultRepository =
                loopResultRepository;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public MQSubscription subscription() {
        MQSubscription subscription = new MQSubscription();
        subscription.setTopic(WorkflowExecutionAuditMqConstants.TOPIC);
        subscription.setConsumerGroup(WorkflowExecutionAuditMqConstants.CONSUMER_GROUP);
        subscription.setShardCount(Math.max(
                1,
                mqProperties.getRedis().getChatPersistShardCount()));
        subscription.setBatchEnabled(true);
        return subscription;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void handle(List<MQMessage> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }
        for (MQMessage message : messages) {
            WorkflowExecutionAuditEvent event = JSON.parseObject(
                    message.getBody(), WorkflowExecutionAuditEvent.class);
            if (event == null || event.getType() == null) {
                throw new IllegalArgumentException("Invalid workflow execution audit event");
            }
            apply(event);
        }
    }

    /**
     * 按事件顺序幂等写入执行记录。
     *
     * @param event 审计事件
     */
    private void apply(WorkflowExecutionAuditEvent event) {
        switch (event.getType()) {
            case CHAIN_STARTED -> createExecution(event);
            case CHAIN_ENDED -> finishExecution(event);
            case NODE_STARTED -> createStep(event);
            case NODE_ENDED -> finishStep(event);
            default -> throw new IllegalArgumentException(
                    "Unsupported workflow execution audit event: " + event.getType());
        }
    }

    /**
     * 创建工作流执行记录。
     *
     * @param event 启动事件
     */
    private void createExecution(WorkflowExecutionAuditEvent event) {
        WorkflowExecResult incoming = requireResult(event);
        try {
            workflowExecResultService.save(incoming);
        } catch (DuplicateKeyException ignored) {
            // MQ 至少一次投递下的重复启动事件按 exec_key 幂等处理。
        }
    }

    /**
     * 完成工作流执行记录。
     *
     * @param event 结束事件
     */
    private void finishExecution(WorkflowExecutionAuditEvent event) {
        WorkflowExecResult incoming = requireResult(event);
        incoming.setOutput(resolveAuditOutput(
                incoming.getOutput()));
        if (workflowExecResultService.updateByExecKey(incoming) != 1) {
            throw new IllegalStateException(
                    "Unable to update workflow execution record: " + incoming.getExecKey());
        }
    }

    /**
     * 创建节点执行步骤。
     *
     * @param event 节点启动事件
     */
    private void createStep(WorkflowExecutionAuditEvent event) {
        WorkflowExecStep incoming = requireStep(event);
        WorkflowExecResult record = workflowExecResultService.getByExecKey(event.getInstanceId());
        if (record == null) {
            throw new IllegalStateException(
                    "Workflow execution record not found: " + event.getInstanceId());
        }
        incoming.setRecordId(record.getId());
        incoming.setInput(resolveAuditOutput(
                incoming.getInput()));
        try {
            workflowExecStepService.save(incoming);
        } catch (DuplicateKeyException ignored) {
            // 同一基础设施触发器恢复后重复投递时按稳定 exec_key 幂等处理。
        }
    }

    /**
     * 完成节点执行步骤。
     *
     * @param event 节点结束事件
     */
    private void finishStep(WorkflowExecutionAuditEvent event) {
        WorkflowExecStep incoming = requireStep(event);
        incoming.setOutput(resolveAuditOutput(
                incoming.getOutput()));
        if (workflowExecStepService.updateByExecKey(incoming) != 1) {
            throw new IllegalStateException(
                    "Unable to update workflow execution step: " + incoming.getExecKey());
        }
    }

    /**
     * 获取事件中的工作流执行记录。
     *
     * @param event 审计事件
     * @return 工作流执行记录
     */
    private WorkflowExecResult requireResult(WorkflowExecutionAuditEvent event) {
        if (event.getResult() == null || event.getResult().getExecKey() == null) {
            throw new IllegalArgumentException("Workflow execution audit result is required");
        }
        return event.getResult();
    }

    /**
     * 获取事件中的节点执行步骤。
     *
     * @param event 审计事件
     * @return 节点执行步骤
     */
    private WorkflowExecStep requireStep(WorkflowExecutionAuditEvent event) {
        if (event.getStep() == null || event.getStep().getExecKey() == null) {
            throw new IllegalArgumentException("Workflow execution audit step is required");
        }
        return event.getStep();
    }

    /**
     * 在审计消费线程还原轻量引用，保持执行记录既有完整 JSON 语义。
     *
     * @param output 可能包含内部引用的 JSON
     * @return 已还原的完整 JSON
     */
    private String resolveAuditOutput(
            String output) {
        if (output == null || output.isBlank()) {
            return output;
        }
        Object parsed = JSON.parse(output);
        Object references = restoreReferences(parsed);
        Object resolved =
                loopResultRepository.resolveReferences(
                        references);
        return JSON.toJSONString(resolved);
    }

    /**
     * 将异步消息中的稳定引用标记恢复为引擎引用对象。
     *
     * @param value JSON 值
     * @return 引擎可解析值
     */
    private Object restoreReferences(Object value) {
        if (value instanceof Map<?, ?> map) {
            if ("easyflow.loop-input.v1".equals(
                    map.get("referenceType"))
                    && map.get("resultId") != null
                    && map.get("itemCount")
                    instanceof Number) {
                return new LoopInputReference(
                        String.valueOf(
                                map.get("resultId")),
                        ((Number) map.get(
                                "itemCount")).intValue());
            }
            if ("easyflow.loop-result.v1".equals(
                    map.get("referenceType"))
                    && map.get("resultId") != null
                    && map.get("iterationCount")
                    instanceof Number
                    && map.get("outputName") != null) {
                return new LoopResultReference(
                        String.valueOf(
                                map.get("resultId")),
                        ((Number) map.get(
                                "iterationCount"))
                                .intValue(),
                        String.valueOf(
                                map.get("outputName")),
                        Boolean.TRUE.equals(
                                map.get("flattenAggregation")));
            }
            Map<Object, Object> restored =
                    new LinkedHashMap<>();
            map.forEach((key, item) ->
                    restored.put(
                            key,
                            restoreReferences(item)));
            return restored;
        }
        if (value instanceof List<?> list) {
            return list.stream()
                    .map(this::restoreReferences)
                    .toList();
        }
        return value;
    }
}
