package tech.easyflow.ai.node;

import com.easyagents.core.util.StringUtil;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.repository.LoopInputReference;
import com.easyagents.flow.core.node.BaseNode;
import com.easyagents.flow.core.util.IoBulkhead;
import com.mybatisflex.core.row.Row;
import com.mybatisflex.core.tenant.TenantManager;
import tech.easyflow.common.util.SpringContextUtil;
import tech.easyflow.common.web.exceptions.BusinessException;
import tech.easyflow.datacenter.execution.model.DatacenterSqlQueryRequest;
import tech.easyflow.datacenter.execution.model.DatasetRef;
import tech.easyflow.datacenter.execution.service.DatacenterDatasetQueryService;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

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


    private static final Pattern PARAM_PATTERN = Pattern.compile("\\{\\{(.+?)\\}\\}");
    private static final Pattern SQL_CODE_BLOCK_PATTERN = Pattern.compile(
            "\\A```(?:sql)?\\s*([\\s\\S]*?)\\s*```\\z",
            Pattern.CASE_INSENSITIVE);
    private static final int QUERY_PAGE_SIZE = Math.max(
            1,
            Integer.getInteger(
                    "easyflow.workflow.dataset.page-size",
                    1_000));

    private DatasetRef datasetRef;
    private String querySql;

    public SearchDatasetNode() {
    }

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

    public SearchDatasetNode(DatasetRef datasetRef, String querySql) {
        this.datasetRef = datasetRef;
        this.querySql = querySql;
    }

    @Override
    public Map<String, Object> execute(Chain chain) {
        Map<String, Object> params =
                chain.getExecutionState().resolveParameters(this);
        DatacenterDatasetQueryService queryService = SpringContextUtil.getBean(DatacenterDatasetQueryService.class);
        DatacenterSqlQueryRequest request = buildRuntimeRequest(params);
        Map<String, Object> result = new HashMap<>();
        try {
            TenantManager.ignoreTenantCondition();
            try (IoBulkhead.Permit ignored = IoBulkhead.dataset().acquire(resolveIoTarget())) {
                String resultId = chain.getStateInstanceId()
                        + ":dataset:"
                        + UUID.randomUUID();
                int rowCount =
                        chain.storeProducedLoopInputOutsideLock(
                        resultId,
                        sink -> queryService.consumeBySql(
                                request,
                                QUERY_PAGE_SIZE,
                                sink::accept),
                        0L,
                        chain.currentFencingClaimId(),
                        chain.currentClaimGeneration());
                result.put(
                        resolveOutputKey("data"),
                        new LoopInputReference(
                                resultId, rowCount));
                return result;
            }
        } finally {
            TenantManager.restoreTenantCondition();
        }
    }

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

    private DatacenterSqlQueryRequest buildRuntimeRequest(Map<String, Object> params) {
        DatacenterSqlQueryRequest request = new DatacenterSqlQueryRequest();
        request.setDatasetRef(copyDatasetRef());
        request.setSql(resolveQuerySql(params));
        return request;
    }

    private String resolveQuerySql(Map<String, Object> params) {
        String sql = normalizeSqlCodeBlock(resolveTemplateString(querySql, params));
        if (!StringUtil.hasText(sql)) {
            throw new BusinessException("查询数据节点未设置 SQL");
        }
        return sql.trim();
    }

    /**
     * 去除完整单个 SQL Markdown 代码块的边界标记。
     *
     * @param sql 原始 SQL 文本
     * @return 可交给 SQL 解析器处理的文本
     */
    private String normalizeSqlCodeBlock(String sql) {
        if (!StringUtil.hasText(sql)) {
            return sql;
        }
        String trimmed = sql.trim();
        Matcher matcher = SQL_CODE_BLOCK_PATTERN.matcher(trimmed);
        return matcher.matches() ? matcher.group(1).trim() : trimmed;
    }

    private DatasetRef copyDatasetRef() {
        DatasetRef copy = new DatasetRef();
        copy.setSourceId(datasetRef == null ? null : datasetRef.getSourceId());
        copy.setCatalogId(datasetRef == null ? null : datasetRef.getCatalogId());
        copy.setCatalogName(datasetRef == null ? null : datasetRef.getCatalogName());
        copy.setTableId(null);
        copy.setTableName(null);
        copy.setVersionId(null);
        return copy;
    }

    private String resolveTemplateString(String text, Map<String, Object> params) {
        if (!StringUtil.hasText(text) || params == null || params.isEmpty()) {
            return text;
        }
        Matcher matcher = PARAM_PATTERN.matcher(text);
        StringBuffer buffer = new StringBuffer();
        while (matcher.find()) {
            Object replacement = params.get(matcher.group(1));
            matcher.appendReplacement(buffer, replacement == null ? "" : Matcher.quoteReplacement(String.valueOf(replacement)));
        }
        matcher.appendTail(buffer);
        return buffer.toString();
    }

    private String resolveOutputKey(String defaultName) {
        List<Parameter> outputDefs = getOutputDefs();
        if (outputDefs == null || outputDefs.isEmpty()) {
            return defaultName;
        }
        String name = outputDefs.get(0).getName();
        return StringUtil.hasText(name) ? name : defaultName;
    }

    public DatasetRef getDatasetRef() {
        return datasetRef;
    }

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

    public String getQuerySql() {
        return querySql;
    }

    public void setQuerySql(String querySql) {
        this.querySql = querySql;
    }
}
