package tech.easyflow.ai.easyagentsflow.listener;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.easyagents.flow.core.chain.*;
import com.easyagents.flow.core.chain.event.*;
import com.easyagents.flow.core.chain.listener.ChainEventListener;
import com.easyagents.flow.core.chain.repository.NodeStateField;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.dao.DuplicateKeyException;
import org.springframework.stereotype.Component;
import tech.easyflow.ai.easyagentsflow.support.PublishedWorkflowDefinitionIds;
import tech.easyflow.ai.entity.Workflow;
import tech.easyflow.ai.entity.WorkflowExecResult;
import tech.easyflow.ai.entity.WorkflowExecStep;
import tech.easyflow.ai.service.WorkflowExecResultService;
import tech.easyflow.ai.service.WorkflowExecStepService;
import tech.easyflow.ai.service.WorkflowService;
import tech.easyflow.ai.utils.WorkFlowUtil;

import javax.annotation.Resource;
import java.util.Date;
import java.util.EnumSet;

@Component
public class ChainEventListenerForSave implements ChainEventListener {


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

    @Resource
    private WorkflowService workflowService;
    @Resource
    private WorkflowExecResultService workflowExecResultService;
    @Resource
    private WorkflowExecStepService workflowExecStepService;

    @Override
    public void onEvent(Event event, Chain chain) {
        if (event instanceof ChainStartEvent) {
            handleChainStartEvent((ChainStartEvent) event, chain);
        }
        if (event instanceof ChainEndEvent) {
            handleChainEndEvent((ChainEndEvent) event, chain);
        }
        if (event instanceof NodeStartEvent) {
            handleNodeStartEvent((NodeStartEvent) event, chain);
        }
        if (event instanceof NodeEndEvent) {
            handleNodeEndEvent((NodeEndEvent) event, chain);
        }
        if (event instanceof ChainStatusChangeEvent) {
            handleChainStatusChangeEvent((ChainStatusChangeEvent) event, chain);
        }
        if (event instanceof ChainResumeEvent) {
            handleChainResumeEvent((ChainResumeEvent) event, chain);
        }
    }

    private void handleChainStartEvent(ChainStartEvent event, Chain chain) {
        log.info("ChainStartEvent: {}", event);
        ChainDefinition definition = chain.getDefinition();
        ChainState state = chain.getState();
        Workflow workflow = resolveWorkflow(definition);
        if (workflow == null) {
            log.error("ChainStartEvent: workflow not found, definitionId={}", definition.getId());
            return;
        }
        String instanceId = state.getInstanceId();
        WorkflowExecResult existed = workflowExecResultService.getByExecKey(instanceId);
        if (existed != null) {
            return;
        }
        WorkflowExecResult record = new WorkflowExecResult();
        record.setExecKey(instanceId);
        record.setWorkflowId(workflow.getId());
        record.setTitle(workflow.getTitle());
        record.setDescription(workflow.getDescription());
        record.setInput(JSON.toJSONString(event.getVariables()));
        record.setWorkflowJson(workflow.getContent());
        record.setStartTime(new Date());
        record.setStatus(state.getStatus().getValue());
        record.setCreatedKey(WorkFlowUtil.getCreatedKey(chain));
        record.setCreatedBy(WorkFlowUtil.getOperator(chain).getId().toString());
        try {
            workflowExecResultService.save(record);
        } catch (DuplicateKeyException e) {
            // 多节点重试时可能并发写同一 exec_key，按幂等处理。
            log.debug("exec result already exists, execKey={}", instanceId, e);
        }
    }

    private void handleChainEndEvent(ChainEndEvent event, Chain chain) {
        log.info("ChainEndEvent: {}", event);
        ChainState state = chain.getState();
        String instanceId = state.getInstanceId();
        WorkflowExecResult record = workflowExecResultService.getByExecKey(instanceId);
        if (record == null) {
            log.error("ChainEndEvent: record not found: {}", instanceId);
        } else {
            record.setEndTime(new Date());
            record.setStatus(state.getStatus().getValue());
            record.setOutput(JSON.toJSONString(state.getExecuteResult()));
            ExceptionSummary error = state.getError();
            if (error != null) {
                record.setErrorInfo(error.getRootCauseClass() + " --> " + error.getRootCauseMessage());
            }
            workflowExecResultService.updateById(record);
        }
    }

    private void handleNodeStartEvent(NodeStartEvent event, Chain chain) {
        log.info("NodeStartEvent: {}", event);
        Node node = event.getNode();
        ChainState ancestorState = findAncestorState(chain.getState(), chain);

        String instanceId = ancestorState.getInstanceId();
        NodeState nodeState = chain.getNodeState(node.getId());

        String execKey = IdUtil.fastSimpleUUID();
        chain.updateNodeStateSafely(node.getId(), state -> {
            state.getMemory().put("executeId", execKey);
            return EnumSet.of(NodeStateField.MEMORY);
        });

        WorkflowExecResult record = workflowExecResultService.getByExecKey(instanceId);
        if (record == null) {
            log.error("NodeStartEvent: record not found: {}", instanceId);
        } else {
            WorkflowExecStep step = new WorkflowExecStep();
            step.setRecordId(record.getId());
            step.setExecKey(execKey);
            step.setNodeId(node.getId());
            step.setNodeName(node.getName());
            step.setInput(JSON.toJSONString(ancestorState.resolveParameters(node)));
            step.setNodeData(JSON.toJSONString(node));
            step.setStartTime(new Date());
            step.setStatus(nodeState.getStatus().getValue());
            workflowExecStepService.save(step);
        }
    }

    private void handleNodeEndEvent(NodeEndEvent event, Chain chain) {
        log.info("NodeEndEvent: {}", event);
        Node node = event.getNode();
        NodeState nodeState = chain.getNodeState(node.getId());
        String execKey = nodeState.getMemory().get("executeId").toString();
        WorkflowExecStep step = workflowExecStepService.getByExecKey(execKey);
        if (step == null) {
            log.error("NodeEndEvent: step not found: {}", execKey);
        } else {
            step.setOutput(JSON.toJSONString(event.getResult()));
            step.setEndTime(new Date());
            step.setStatus(nodeState.getStatus().getValue());
            ExceptionSummary error = nodeState.getError();
            if (error != null) {
                step.setErrorInfo(error.getRootCauseClass() + " --> " + error.getRootCauseMessage());
            }
            workflowExecStepService.updateById(step);
        }
    }

    private void handleChainStatusChangeEvent(ChainStatusChangeEvent event, Chain chain) {
        log.info("ChainStatusChangeEvent: {}", event);
    }

    private void handleChainResumeEvent(ChainResumeEvent event, Chain chain) {
        log.info("ChainResumeEvent: {}", event);
    }

    /**
     * 递归查找顶级状态
     */
    private ChainState findAncestorState(ChainState state, Chain chain) {
        String parentInstanceId = state.getParentInstanceId();
        if (StrUtil.isEmpty(parentInstanceId)) {
            return state;
        }
        ChainState chainState = chain.getChainStateRepository().load(parentInstanceId);
        return findAncestorState(chainState, chain);
    }

    /**
     * 根据定义 ID 解析当前执行所对应的工作流。
     * 已发布快照执行会使用 published 前缀，需要先还原为真实工作流 ID。
     */
    private Workflow resolveWorkflow(ChainDefinition definition) {
        if (definition == null || StrUtil.isBlank(definition.getId())) {
            return null;
        }
        String definitionId = definition.getId();
        String workflowId = PublishedWorkflowDefinitionIds.unwrap(definitionId);
        try {
            java.math.BigInteger id = new java.math.BigInteger(workflowId);
            if (PublishedWorkflowDefinitionIds.isPublished(definitionId)) {
                return workflowService.getPublishedById(id);
            }
            return workflowService.getById(id);
        } catch (NumberFormatException ex) {
            log.error("Unsupported workflow definition id: {}", definitionId, ex);
            return null;
        }
    }
}
