import type { Edge, Node } from '@vue-flow/core';

import type { ExecutionLogEntry, WorkflowRun } from '#/api/ai-platform/ai-platform';

function mapLogStatus(
  status: string,
): 'failed' | 'pending' | 'running' | 'success' {
  if (status === 'completed') return 'success';
  if (status === 'failed') return 'failed';
  if (status === 'waiting') return 'running';
  return 'success';
}

function mapEdgeStatus(
  status: string,
): 'completed' | 'failed' | 'running' | undefined {
  if (status === 'completed') return 'completed';
  if (status === 'failed') return 'failed';
  if (status === 'waiting') return 'running';
  return 'completed';
}

function resolveNodeOutput(log: ExecutionLogEntry, run: WorkflowRun) {
  const stored = run.outputs?.[`_node_${log.node_id}`];
  const rawOutput = log.output ?? stored?.output ?? stored;

  if (rawOutput === undefined) {
    return undefined;
  }

  if (
    rawOutput &&
    typeof rawOutput === 'object' &&
    !Array.isArray(rawOutput) &&
    ('output' in rawOutput || 'output_variables' in rawOutput)
  ) {
    return rawOutput;
  }

  const outputVariables =
    stored && typeof stored === 'object'
      ? Object.fromEntries(
          Object.entries(stored).filter(([key]) => key !== 'output'),
        )
      : {};

  return {
    output: rawOutput,
    output_variables: outputVariables,
  };
}

function buildReplayInputs(
  logIndex: number,
  logs: ExecutionLogEntry[],
  run: WorkflowRun,
) {
  const log = logs[logIndex];
  if (log?.inputs) {
    return log.inputs;
  }

  const variables: Record<string, any> = { ...(run.inputs || {}) };

  for (let i = 0; i < logIndex; i += 1) {
    const prevLog = logs[i];
    if (!prevLog?.node_id) continue;
    const nodeKey = `_node_${prevLog.node_id}`;
    const stored = run.outputs?.[nodeKey];
    if (stored !== undefined) {
      variables[nodeKey] = stored;
    }
  }

  let previous_output: any = null;
  if (logIndex > 0) {
    const prevLog = logs[logIndex - 1];
    if (prevLog) {
      previous_output =
        prevLog.output ??
        run.outputs?.[`_node_${prevLog.node_id}`]?.output ??
        run.outputs?.[`_node_${prevLog.node_id}`] ??
        null;
    }
  }

  return {
    variables,
    previous_output,
  };
}

export function applyRunReplay(
  nodes: Node[],
  edges: Edge[],
  run: WorkflowRun,
): { edges: Edge[]; nodes: Node[] } {
  const nodeMap = new Map(nodes.map((node) => [node.id, { ...node, data: { ...node.data } }]));
  const edgeList = edges.map((edge) => ({
    ...edge,
    data: { ...(edge.data || {}), status: undefined as string | undefined },
  }));

  const logs = (run.execution_log || []) as ExecutionLogEntry[];
  const executedNodeIds: string[] = [];

  logs.forEach((log, logIndex) => {
    if (!log?.node_id) return;
    executedNodeIds.push(log.node_id);

    const node = nodeMap.get(log.node_id);
    if (!node) return;

    const isWaiting = log.status === 'waiting';
    const isSuccess = log.status === 'completed';
    const nodeOutput = resolveNodeOutput(log, run);
    const nodeInputs = buildReplayInputs(logIndex, logs, run);

    node.data = {
      ...node.data,
      executionStatus: {
        status: isWaiting ? 'running' : isSuccess ? 'success' : 'failed',
        elapsed_time: log.elapsed_time,
        tokens_used: log.tokens_used,
        error: log.error,
      },
      executionResult: {
        status: mapLogStatus(log.status || 'completed'),
        elapsed_time: log.elapsed_time,
        tokens_used: log.tokens_used,
        error_message: log.error,
        output: nodeOutput,
        inputs: nodeInputs,
      },
    };
    nodeMap.set(log.node_id, node);
  });

  if (run.current_node_id && run.status === 'waiting') {
    const waitingNode = nodeMap.get(run.current_node_id);
    if (waitingNode) {
      waitingNode.data = {
        ...waitingNode.data,
        executionStatus: {
          status: 'running',
        },
        executionResult: {
          status: 'running',
        },
      };
      nodeMap.set(run.current_node_id, waitingNode);
    }
  }

  for (let i = 1; i < executedNodeIds.length; i += 1) {
    const sourceId = executedNodeIds[i - 1];
    const targetId = executedNodeIds[i];
    const log = logs[i];
    const edgeStatus = mapEdgeStatus(log?.status || 'completed');

    edgeList.forEach((edge, index) => {
      if (edge.source === sourceId && edge.target === targetId) {
        edgeList[index] = {
          ...edge,
          data: {
            ...(edge.data || {}),
            status: edgeStatus,
          },
        };
      }
    });
  }

  return {
    nodes: Array.from(nodeMap.values()),
    edges: edgeList,
  };
}
