package tech.easyflow.ai.node;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.runtime.RetryableTriggerException;
import com.easyagents.flow.core.node.BaseNode;
import com.easyagents.flow.core.util.IoBulkhead;
import com.mybatisflex.core.tenant.TenantManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tech.easyflow.ai.utils.WorkFlowUtil;
import tech.easyflow.ai.easyagentsflow.config.WorkflowRuntimeProperties;
import tech.easyflow.common.entity.LoginAccount;
import tech.easyflow.common.cache.RedisIdempotencyExecutor.IdempotentOperationInProgressException;
import tech.easyflow.common.util.SpringContextUtil;
import tech.easyflow.datacenter.execution.model.DatasetRef;
import tech.easyflow.datacenter.execution.service.DatacenterDatasetQueryService;
import tech.easyflow.datacenter.execution.service.DatacenterDatasetWriteService;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class SaveDatasetNode extends BaseNode {
    private static final long serialVersionUID = 1L;

    private static final Logger log = LoggerFactory.getLogger(SaveDatasetNode.class);

    private DatasetRef datasetRef;

    public SaveDatasetNode() {
    }

    public SaveDatasetNode(DatasetRef datasetRef) {
        this.datasetRef = datasetRef;
    }

    @Override
    public Map<String, Object> execute(Chain chain) {
        Map<String, Object> state =
                chain.getExecutionState().resolveParameters(this);
        JSONObject payload = new JSONObject(state);
        JSONArray saveList = payload.getJSONArray("saveList");
        if (saveList == null || saveList.isEmpty()) {
            throw new RuntimeException("saveList 不能为空");
        }
        LoginAccount account = WorkFlowUtil.getOperator(chain);
        DatacenterDatasetWriteService writeService = SpringContextUtil.getBean(DatacenterDatasetWriteService.class);
        DatacenterDatasetQueryService queryService = SpringContextUtil.getBean(DatacenterDatasetQueryService.class);
        WorkflowRuntimeProperties runtimeProperties = SpringContextUtil.getBean(WorkflowRuntimeProperties.class);
        List<JSONObject> rows = new ArrayList<>(saveList.size());
        for (Object item : saveList) {
            rows.add(item instanceof JSONObject json ? json : JSONObject.from(item));
        }
        try {
            TenantManager.ignoreTenantCondition();
            try (IoBulkhead.Permit ignored = IoBulkhead.dataset().acquire(resolveIoTarget())) {
                writeService.saveRowsIdempotently(
                        datasetRef,
                        rows,
                        account,
                        runtimeProperties.getDataWriteBatchSize(),
                        chain.currentExecutionIdempotencyKey(getId()));
                var schema = queryService.getLocation(datasetRef);
                Map<String, Object> result = new HashMap<>();
                result.put("successRows", rows.size());
                result.put("source", schema.getSource());
                result.put("catalog", schema.getCatalog());
                result.put("table", schema.getTable());
                result.put("version", datasetRef.getVersionId());
                return result;
            }
        } catch (IdempotentOperationInProgressException conflict) {
            throw new RetryableTriggerException("数据集写入幂等操作仍在处理中", conflict);
        } catch (Exception ex) {
            log.error("工作流保存数据到统一数据集失败，datasetRef={}", datasetRef, ex);
            throw ex;
        } finally {
            TenantManager.restoreTenantCondition();
        }
    }

    /**
     * 获取数据源级 I/O 隔离目标。
     *
     * @return 数据源目标键
     */
    private String resolveIoTarget() {
        return "dataset:"
                + (datasetRef == null || datasetRef.getSourceId() == null
                ? "unknown"
                : datasetRef.getSourceId());
    }

    public DatasetRef getDatasetRef() {
        return datasetRef;
    }

    public void setDatasetRef(DatasetRef datasetRef) {
        this.datasetRef = datasetRef;
    }
}
