package tech.easyflow.ai.easyagentsflow.service;

import com.easyagents.flow.core.code.CodeRuntimeEngineManager;
import com.easyagents.flow.core.filestoreage.FileStorageManager;
import com.easyagents.flow.core.filestoreage.FileStorageProvider;
import com.easyagents.flow.core.knowledge.KnowledgeManager;
import com.easyagents.flow.core.knowledge.KnowledgeProvider;
import com.easyagents.flow.core.llm.LlmManager;
import com.easyagents.flow.core.llm.LlmProvider;
import com.easyagents.flow.core.parser.ChainParser;
import com.easyagents.flow.core.searchengine.SearchEngine;
import com.easyagents.flow.core.searchengine.SearchEngineManager;
import com.easyagents.flow.core.searchengine.SearchEngineProvider;
import com.easyagents.flow.core.searchengine.impl.BochaaiSearchEngineImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import tech.easyflow.ai.config.BochaaiProps;
import tech.easyflow.ai.config.CodeEngineProps;
import tech.easyflow.ai.easyagentsflow.code.PythonRuntimeEngine;
import tech.easyflow.ai.easyagentsflow.code.UnavailableCodeRuntimeEngine;
import tech.easyflow.ai.node.*;
import tech.easyflow.common.util.StringUtil;

import javax.annotation.Resource;
import java.util.Arrays;
import java.util.List;

@Component
public class TinyFlowConfigService {

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

    @Resource
    private LlmProvider llmProvider;
    @Resource
    private FileStorageProvider fileStorageProvider;
    @Resource
    private BochaaiProps bochaaiProps;
    @Resource
    private KnowledgeProvider knowledgeProvider;
    @Resource
    private CodeEngineProps codeEngineProps;
    @Resource
    private CodeEngineCapabilityService codeEngineCapabilityService;

    public void initProvidersAndNodeParsers(ChainParser chainParser) {
        setExtraNodeParser(chainParser);
        setLlmProvider();
        setFileStorage();
        setKnowledgeProvider();
        setSearchEngineProvider();
        setCodeRuntimeEngineProvider();
    }

    private void setFileStorage() {
        FileStorageManager.getInstance().registerProvider(fileStorageProvider);
    }

    public void setExtraNodeParser(ChainParser chainParser) {

        // 文档解析
        DocNodeParser docNodeParser = new DocNodeParser();
        // 文件生成
        MakeFileNodeParser makeFileNodeParser = new MakeFileNodeParser();
        // 插件
        PluginToolNodeParser pluginToolNodeParser = new PluginToolNodeParser();
        // 下载文件节点
        DownloadNodeParser downloadNodeParser = new DownloadNodeParser();
        // 保存数据节点
        SaveDatasetNodeParser saveDatasetNodeParser = new SaveDatasetNodeParser();
        // 查询数据节点
        SearchDatasetNodeParser searchDatasetNodeParser = new SearchDatasetNodeParser();
        // 工作流节点
        WorkflowNodeParser workflowNodeParser = new WorkflowNodeParser();
        // 条件判断节点
        ConditionNodeParser conditionNodeParser = new ConditionNodeParser();

        chainParser.addNodeParser(docNodeParser.getNodeName(), docNodeParser);
        chainParser.addNodeParser(makeFileNodeParser.getNodeName(), makeFileNodeParser);
        chainParser.addNodeParser(pluginToolNodeParser.getNodeName(), pluginToolNodeParser);
        chainParser.addNodeParser(downloadNodeParser.getNodeName(), downloadNodeParser);
        chainParser.addNodeParser(saveDatasetNodeParser.getNodeName(), saveDatasetNodeParser);
        chainParser.addNodeParser(searchDatasetNodeParser.getNodeName(), searchDatasetNodeParser);
        chainParser.addNodeParser(workflowNodeParser.getNodeName(), workflowNodeParser);
        chainParser.addNodeParser(conditionNodeParser.getNodeName(), conditionNodeParser);
    }

    public void setSearchEngineProvider() {
        BochaaiSearchEngineImpl engine = new BochaaiSearchEngineImpl();
        engine.setApiKey(bochaaiProps.getApiKey());
        SearchEngineManager.getInstance().registerProvider(new SearchEngineProvider() {
            @Override
            public SearchEngine getSearchEngine(Object id) {
                return id.equals("bocha-search") ? engine : null;
            }
        });
    }

    public void setLlmProvider() {
        LlmManager.getInstance().registerProvider(llmProvider);
    }

    public void setKnowledgeProvider() {
        KnowledgeManager.getInstance().registerProvider(knowledgeProvider);
    }

    public void setCodeRuntimeEngineProvider() {
        CodeRuntimeEngineManager manager = CodeRuntimeEngineManager.getInstance();

        PythonRuntimeEngine.ProbeResult probeResult;
        String pythonCommand = codeEngineProps.getCommand();
        if (!codeEngineProps.isEnabled()) {
            probeResult = new PythonRuntimeEngine.ProbeResult(false, "配置已关闭");
        } else {
            probeResult = PythonRuntimeEngine.probe(pythonCommand, 2000L);
        }

        boolean pythonAvailable = codeEngineProps.isEnabled() && probeResult.isAvailable();
        codeEngineCapabilityService.setPythonCapability(pythonAvailable, probeResult.getMessage());

        if (pythonAvailable) {
            PythonRuntimeEngine pythonRuntimeEngine = new PythonRuntimeEngine(
                    codeEngineProps.getCommand(),
                    codeEngineProps.getTimeoutMs(),
                    codeEngineProps.getMaxOutputBytes(),
                    codeEngineProps.getWorkingDir()
            );
            manager.registerProvider(engineId -> {
                if ("python".equals(engineId) || "py".equals(engineId)) {
                    return pythonRuntimeEngine;
                }
                return null;
            });
            log.info("已启用 Python 代码执行引擎: command={}, timeoutMs={}, maxOutputBytes={}, workingDir={}",
                    codeEngineProps.getCommand(), codeEngineProps.getTimeoutMs(), codeEngineProps.getMaxOutputBytes(), codeEngineProps.getWorkingDir());
        } else {
            final String reason = probeResult.getMessage();
            manager.registerProvider(engineId -> {
                if ("python".equals(engineId) || "py".equals(engineId)) {
                    return new UnavailableCodeRuntimeEngine(reason);
                }
                return null;
            });
            log.warn("Python 代码执行引擎不可用: enabled={}, command={}, reason={}",
                    codeEngineProps.isEnabled(), codeEngineProps.getCommand(), reason);
        }

        manager.registerProvider(engineId -> {
            if (!(engineId instanceof String)) {
                return new UnavailableCodeRuntimeEngine("未配置代码执行引擎");
            }
            String engine = ((String) engineId).trim();
            if (StringUtil.noText(engine)) {
                return new UnavailableCodeRuntimeEngine("未配置代码执行引擎");
            }
            return new UnavailableCodeRuntimeEngine("不支持的代码执行引擎: " + engine + "，当前支持: " + supportedEngineLabels());
        });
    }

    private String supportedEngineLabels() {
        List<String> engines = codeEngineCapabilityService.listSupportedCodeEngines()
                .stream()
                .filter(it -> Boolean.TRUE.equals(it.get("available")))
                .map(it -> (String) it.get("value"))
                .toList();
        if (engines.isEmpty()) {
            engines = Arrays.asList("js");
        }
        return String.join(", ", engines);
    }
}
