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.node.BaseNode;
import com.mybatisflex.core.tenant.TenantManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tech.easyflow.ai.utils.WorkFlowUtil;
import tech.easyflow.common.entity.LoginAccount;
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.HashMap;
import java.util.Map;

public class SaveDatasetNode extends BaseNode {

    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.getState().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);
        int successRows = 0;
        try {
            TenantManager.ignoreTenantCondition();
            for (Object item : saveList) {
                JSONObject row = item instanceof JSONObject json ? json : JSONObject.from(item);
                writeService.saveRow(datasetRef, row, account);
                successRows++;
            }
            var schema = queryService.getSchema(datasetRef);
            Map<String, Object> result = new HashMap<>();
            result.put("successRows", successRows);
            result.put("source", schema.getSource());
            result.put("catalog", schema.getCatalog());
            result.put("table", schema.getTable());
            result.put("version", datasetRef.getVersionId());
            return result;
        } catch (Exception ex) {
            log.error("工作流保存数据到统一数据集失败，datasetRef={}", datasetRef, ex);
            throw ex;
        } finally {
            TenantManager.restoreTenantCondition();
        }
    }

    public DatasetRef getDatasetRef() {
        return datasetRef;
    }

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