package tech.easyflow.ai.easyagentsflow.repository;

import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.NodeState;
import com.easyagents.flow.core.chain.repository.ChainStateField;
import com.easyagents.flow.core.chain.repository.NodeStateField;
import tech.easyflow.ai.easyagentsflow.support.WorkflowExecutionStepKey;
import tech.easyflow.common.cache.VersionedFields;

import java.io.Serializable;
import java.util.EnumSet;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 工作流状态对象与 Redis 字段之间的无反射映射。
 */
final class WorkflowStateFields {

    static final String FORMAT_FIELD = "_format";
    static final String FORMAT_VERSION = "2";
    private static final String NODE_ID_FIELD = "_nodeId";
    private static final String CHAIN_INSTANCE_ID_FIELD = "_chainInstanceId";

    private WorkflowStateFields() {
    }

    /**
     * 判断快照是否采用字段化格式。
     *
     * @param snapshot Redis 快照
     * @return 字段化格式时为 {@code true}
     */
    static boolean isFieldFormat(VersionedFields snapshot) {
        return snapshot != null
                && FORMAT_VERSION.equals(snapshot.getFields().get(FORMAT_FIELD));
    }

    /**
     * 将完整工作流状态编码为字段。
     *
     * @param state 工作流状态
     * @return 完整字段映射
     */
    static Map<String, Serializable> allChainFields(ChainState state) {
        EnumSet<ChainStateField> fields = EnumSet.allOf(ChainStateField.class);
        fields.remove(ChainStateField.VERSION);
        return chainFields(state, fields);
    }

    /**
     * 将变化的工作流状态字段编码为可独立提交的值。
     *
     * @param state 工作流状态
     * @param fields 变化字段
     * @return 字段映射
     */
    static Map<String, Serializable> chainFields(
            ChainState state, EnumSet<ChainStateField> fields) {
        Map<String, Serializable> values = new LinkedHashMap<>();
        values.put(FORMAT_FIELD, FORMAT_VERSION);
        for (ChainStateField field : fields) {
            Serializable value = switch (field) {
                case INSTANCE_ID -> state.getInstanceId();
                case STATUS -> state.getStatus();
                case MESSAGE -> state.getMessage();
                case ERROR -> state.getError();
                case MEMORY -> state.getMemory();
                case COMPUTE_COST -> state.getComputeCost();
                case SUSPEND_NODE_IDS -> serializable(state.getSuspendNodeIds());
                case SUSPEND_FOR_PARAMETERS -> serializable(state.getSuspendForParameters());
                case EXECUTE_RESULT -> serializable(state.getExecuteResult());
                case CHAIN_DEFINITION_ID -> state.getChainDefinitionId();
                case ENVIRONMENT -> serializable(state.getEnvironment());
                case PARENT_INSTANCE_ID -> state.getParentInstanceId();
                case AUDIT_INSTANCE_ID -> state.getAuditInstanceId();
                case TRIGGER_NODE_IDS -> serializable(state.getTriggerNodeIds());
                case TRIGGER_EDGE_IDS -> serializable(state.getTriggerEdgeIds());
                case UNCHECKED_EDGE_IDS -> serializable(state.getUncheckedEdgeIds());
                case UNCHECKED_NODE_IDS -> serializable(state.getUncheckedNodeIds());
                case STARTED_AT -> state.getStartedAt();
                case CHILD_EXECUTION_COUNT -> state.getChildExecutionCount();
                case VERSION, PAYLOAD, NODE_STATES, CHILD_STATE_IDS -> null;
            };
            if (field != ChainStateField.VERSION
                    && field != ChainStateField.PAYLOAD
                    && field != ChainStateField.NODE_STATES
                    && field != ChainStateField.CHILD_STATE_IDS) {
                values.put(field.name(), value);
            }
        }
        return values;
    }

    /**
     * 从字段快照还原工作流状态。
     *
     * @param snapshot Redis 字段快照
     * @return 工作流状态
     */
    @SuppressWarnings("unchecked")
    static ChainState decodeChain(VersionedFields snapshot) {
        Map<String, Object> fields = snapshot.getFields();
        ChainState state = new ChainState();
        state.setInstanceId((String) fields.get(ChainStateField.INSTANCE_ID.name()));
        state.setStatus((com.easyagents.flow.core.chain.ChainStatus)
                fields.get(ChainStateField.STATUS.name()));
        state.setMessage((String) fields.get(ChainStateField.MESSAGE.name()));
        state.setError((com.easyagents.flow.core.chain.ExceptionSummary)
                fields.get(ChainStateField.ERROR.name()));
        Object memory = fields.get(ChainStateField.MEMORY.name());
        state.setMemory(memory == null
                ? new ConcurrentHashMap<>()
                : new ConcurrentHashMap<>((Map<String, Object>) memory));
        state.setComputeCost(number(fields.get(ChainStateField.COMPUTE_COST.name())));
        state.setSuspendNodeIds((java.util.Set<String>)
                fields.get(ChainStateField.SUSPEND_NODE_IDS.name()));
        state.setSuspendForParameters((java.util.List<com.easyagents.flow.core.chain.Parameter>)
                fields.get(ChainStateField.SUSPEND_FOR_PARAMETERS.name()));
        state.setExecuteResult((Map<String, Object>)
                fields.get(ChainStateField.EXECUTE_RESULT.name()));
        state.setChainDefinitionId((String)
                fields.get(ChainStateField.CHAIN_DEFINITION_ID.name()));
        state.setEnvironment((Map<String, Object>)
                fields.get(ChainStateField.ENVIRONMENT.name()));
        state.setParentInstanceId((String)
                fields.get(ChainStateField.PARENT_INSTANCE_ID.name()));
        state.setAuditInstanceId((String)
                fields.get(ChainStateField.AUDIT_INSTANCE_ID.name()));
        state.setTriggerNodeIds((java.util.List<String>)
                fields.get(ChainStateField.TRIGGER_NODE_IDS.name()));
        state.setTriggerEdgeIds((java.util.List<String>)
                fields.get(ChainStateField.TRIGGER_EDGE_IDS.name()));
        state.setUncheckedEdgeIds((java.util.List<String>)
                fields.get(ChainStateField.UNCHECKED_EDGE_IDS.name()));
        state.setUncheckedNodeIds((java.util.List<String>)
                fields.get(ChainStateField.UNCHECKED_NODE_IDS.name()));
        state.setStartedAt(number(fields.get(ChainStateField.STARTED_AT.name())));
        state.setChildExecutionCount(number(
                fields.get(ChainStateField.CHILD_EXECUTION_COUNT.name())));
        state.setVersion(snapshot.getVersion());
        return state;
    }

    /**
     * 将完整节点状态编码为字段。
     *
     * @param state 节点状态
     * @return 完整字段映射
     */
    static Map<String, Serializable> allNodeFields(NodeState state) {
        EnumSet<NodeStateField> fields = EnumSet.allOf(NodeStateField.class);
        fields.remove(NodeStateField.VERSION);
        return nodeFields(state, fields);
    }

    /**
     * 将变化的节点状态字段编码为可独立提交的值。
     *
     * @param state 节点状态
     * @param fields 变化字段
     * @return 字段映射
     */
    static Map<String, Serializable> nodeFields(
            NodeState state, EnumSet<NodeStateField> fields) {
        Map<String, Serializable> values = new LinkedHashMap<>();
        values.put(FORMAT_FIELD, FORMAT_VERSION);
        values.put(NODE_ID_FIELD, state.getNodeId());
        values.put(CHAIN_INSTANCE_ID_FIELD, state.getChainInstanceId());
        for (NodeStateField field : fields) {
            Serializable value = switch (field) {
                case STATUS -> state.getStatus();
                case ERROR -> state.getError();
                case MEMORY -> state.getMemory();
                case RETRY_COUNT -> state.getRetryCount();
                case EXECUTE_COUNT -> state.getExecuteCount();
                case EXECUTE_EDGE_IDS -> serializable(state.getExecuteEdgeIds());
                case EXECUTION_ATTEMPT_KEY -> state.getExecutionAttemptKey();
                case LOOP_COUNT -> state.getLoopCount();
                case TRIGGER_COUNT -> state.getTriggerCount();
                case TRIGGER_EDGE_IDS -> serializable(state.getTriggerEdgeIds());
                case INSTANCE_ID, MESSAGE, PAYLOAD, NODE_STATES, COMPUTE_COST,
                     SUSPEND_NODE_IDS, SUSPEND_FOR_PARAMETERS, EXECUTE_RESULT,
                     ENVIRONMENT, VERSION -> null;
            };
            if (field == NodeStateField.STATUS
                    || field == NodeStateField.ERROR
                    || field == NodeStateField.MEMORY
                    || field == NodeStateField.RETRY_COUNT
                    || field == NodeStateField.EXECUTE_COUNT
                    || field == NodeStateField.EXECUTE_EDGE_IDS
                    || field == NodeStateField.EXECUTION_ATTEMPT_KEY
                    || field == NodeStateField.LOOP_COUNT
                    || field == NodeStateField.TRIGGER_COUNT
                    || field == NodeStateField.TRIGGER_EDGE_IDS) {
                values.put(field.name(), value);
            }
        }
        return values;
    }

    /**
     * 从字段快照还原节点状态。
     *
     * @param snapshot Redis 字段快照
     * @return 节点状态
     */
    @SuppressWarnings("unchecked")
    static NodeState decodeNode(VersionedFields snapshot) {
        Map<String, Object> fields = snapshot.getFields();
        NodeState state = new NodeState();
        state.setNodeId((String) fields.get(NODE_ID_FIELD));
        state.setChainInstanceId((String) fields.get(CHAIN_INSTANCE_ID_FIELD));
        state.setStatus((com.easyagents.flow.core.chain.NodeStatus)
                fields.get(NodeStateField.STATUS.name()));
        state.setError((com.easyagents.flow.core.chain.ExceptionSummary)
                fields.get(NodeStateField.ERROR.name()));
        Object memory = fields.get(NodeStateField.MEMORY.name());
        state.setMemory(memory == null
                ? new ConcurrentHashMap<>()
                : new ConcurrentHashMap<>((Map<String, Object>) memory));
        state.setRetryCount(integer(fields.get(NodeStateField.RETRY_COUNT.name())));
        state.setExecuteCount(atomic(fields.get(NodeStateField.EXECUTE_COUNT.name())));
        state.setExecuteEdgeIds(defaultList(
                (java.util.List<String>) fields.get(NodeStateField.EXECUTE_EDGE_IDS.name())));
        state.setExecutionAttemptKey(
                (String) fields.get(
                        NodeStateField
                                .EXECUTION_ATTEMPT_KEY
                                .name()));
        state.setLoopCount(integer(fields.get(NodeStateField.LOOP_COUNT.name())));
        state.setTriggerCount(atomic(fields.get(NodeStateField.TRIGGER_COUNT.name())));
        state.setTriggerEdgeIds(defaultList(
                (java.util.List<String>) fields.get(NodeStateField.TRIGGER_EDGE_IDS.name())));
        state.setVersion(snapshot.getVersion());
        return normalizeNode(state);
    }

    /**
     * 补齐升级前节点快照缺失的执行尝试键。
     *
     * @param state 字段化或旧对象节点状态
     * @return 原节点状态
     */
    static NodeState normalizeNode(NodeState state) {
        if (state == null
                || (state.getExecutionAttemptKey() != null
                && !state.getExecutionAttemptKey().isBlank())
                || state.getMemory() == null) {
            return state;
        }
        Object legacyExecKey =
                state.getMemory().get(
                        "executeId");
        if (legacyExecKey instanceof String) {
            state.setExecutionAttemptKey(
                    WorkflowExecutionStepKey
                            .encodeLegacy(
                                    (String) legacyExecKey));
        }
        return state;
    }

    private static Serializable serializable(Object value) {
        return value == null ? null : (Serializable) value;
    }

    private static long number(Object value) {
        return value instanceof Number ? ((Number) value).longValue() : 0L;
    }

    private static int integer(Object value) {
        return value instanceof Number ? ((Number) value).intValue() : 0;
    }

    private static AtomicInteger atomic(Object value) {
        if (value instanceof AtomicInteger) {
            return new AtomicInteger(((AtomicInteger) value).get());
        }
        return new AtomicInteger(integer(value));
    }

    private static java.util.List<String> defaultList(java.util.List<String> value) {
        return value == null ? new java.util.ArrayList<>() : new java.util.ArrayList<>(value);
    }
}
