package tech.easyflow.ai.easyagents.tool;

import com.easyagents.core.model.chat.tool.BaseTool;
import com.easyagents.core.model.chat.tool.Parameter;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.DataType;
import com.easyagents.flow.core.chain.runtime.ChainExecutor;
import tech.easyflow.ai.entity.Workflow;
import tech.easyflow.common.util.SpringContextUtil;

import java.math.BigInteger;
import java.util.*;

public class WorkflowTool extends BaseTool {

    private BigInteger workflowId;
    private String definitionId;

    public WorkflowTool() {
    }

    public WorkflowTool(Workflow workflow, boolean needEnglishName) {
        this(workflow, needEnglishName, workflow.getId() == null ? null : workflow.getId().toString());
    }

    public WorkflowTool(Workflow workflow, boolean needEnglishName, String definitionId) {
        this.workflowId = workflow.getId();
        this.definitionId = definitionId;
        this.name = ChatToolNameHelper.resolveToolName(
                needEnglishName,
                workflow.getEnglishName(),
                workflow.getTitle(),
                "workflow",
                workflow.getId()
        );
        this.description = workflow.getDescription();
        this.parameters = toParameters(workflow, definitionId);
    }


    static Parameter[] toParameters(Workflow workflow, String definitionId) {
        ChainExecutor executor = SpringContextUtil.getBean(ChainExecutor.class);
        ChainDefinition definition = executor.getDefinitionRepository().getChainDefinitionById(definitionId);
        List<com.easyagents.flow.core.chain.Parameter> parameterDefs = definition.getStartParameters();
        if (parameterDefs == null || parameterDefs.isEmpty()) {
            return new Parameter[0];
        }

        Parameter[] parameters = new Parameter[parameterDefs.size()];
        for (int i = 0; i < parameterDefs.size(); i++) {
            com.easyagents.flow.core.chain.Parameter parameterDef = parameterDefs.get(i);
            parameters[i] = toToolParameter(parameterDef);
        }
        return parameters;
    }

    private static Parameter toToolParameter(com.easyagents.flow.core.chain.Parameter parameterDef) {
        Parameter parameter = new Parameter();
        parameter.setName(parameterDef.getName());
        parameter.setDescription(parameterDef.getDescription());
        parameter.setRequired(parameterDef.isRequired());

        DataType dataType = parameterDef.getDataType();
        if (dataType == null) {
            dataType = DataType.String;
        }

        switch (dataType) {
            case Object:
                parameter.setType("object");
                parameter.setChildren(toToolChildren(parameterDef.getChildren()));
                break;
            case Array:
                parameter.setType("array");
                parameter.setChildren(toToolChildren(parameterDef.getChildren()));
                break;
            case Array_Object:
                parameter.setType("array");
                parameter.setChildren(Collections.singletonList(buildArrayItem("object", parameterDef.getChildren())));
                break;
            case Array_String:
                parameter.setType("array");
                parameter.setChildren(Collections.singletonList(buildArrayItem("string", null)));
                break;
            case Array_Number:
                parameter.setType("array");
                parameter.setChildren(Collections.singletonList(buildArrayItem("number", null)));
                break;
            case Array_Boolean:
                parameter.setType("array");
                parameter.setChildren(Collections.singletonList(buildArrayItem("boolean", null)));
                break;
            case Array_File:
                parameter.setType("array");
                parameter.setChildren(Collections.singletonList(buildArrayItem("string", null)));
                break;
            case Number:
                parameter.setType("number");
                break;
            case Boolean:
                parameter.setType("boolean");
                break;
            case File:
                parameter.setType("string");
                break;
            case String:
            default:
                parameter.setType("string");
                break;
        }

        return parameter;
    }

    private static List<Parameter> toToolChildren(List<com.easyagents.flow.core.chain.Parameter> parameterDefs) {
        if (parameterDefs == null || parameterDefs.isEmpty()) {
            return null;
        }
        List<Parameter> children = new ArrayList<>(parameterDefs.size());
        for (com.easyagents.flow.core.chain.Parameter childDef : parameterDefs) {
            children.add(toToolParameter(childDef));
        }
        return children;
    }

    private static Parameter buildArrayItem(String itemType, List<com.easyagents.flow.core.chain.Parameter> childrenDefs) {
        Parameter item = new Parameter();
        item.setName(null);
        item.setType(itemType);
        item.setChildren(toToolChildren(childrenDefs));
        return item;
    }

    public BigInteger getWorkflowId() {
        return workflowId;
    }

    public void setWorkflowId(BigInteger workflowId) {
        this.workflowId = workflowId;
    }

    public String getDefinitionId() {
        return definitionId;
    }

    public void setDefinitionId(String definitionId) {
        this.definitionId = definitionId;
    }

    @Override
    public Object invoke(Map<String, Object> argsMap) {
        ChainExecutor executor = SpringContextUtil.getBean(ChainExecutor.class);
        return executor.execute(definitionId, argsMap);
    }

    @Override
    public String toString() {
        return "AiWorkflowFunction{" +
                "workflowId=" + workflowId +
                ", definitionId='" + definitionId + '\'' +
                ", name='" + name + '\'' +
                ", description='" + description + '\'' +
                ", parameters=" + Arrays.toString(parameters) +
                '}';
    }
}
