import type { Instance, SuspendInfo } from "../types/instance";
import type { Node } from "../types/node";
import type { MachineMessage } from "../types/messages";
import type {
  CommandContext,
  CommandResult,
  CommandExecutionResult,
} from "../types/commands";
import { isCommandValueResult, isResumeResult } from "../types/commands";
import type { SpawnTarget, SpawnOptions } from "../types/transitions";
import {
  isTransitionToResult,
  isSpawnResult,
  isCedeResult,
  isSuspendResult,
} from "../types/transitions";
import { cede, spawn, suspend } from "../helpers/cede-spawn";
import { shallowMerge } from "../types/state";
import { createInstance, createSuspendInfo, clearSuspension } from "../types/instance";
import { userMessage, instanceMessage } from "../types/messages";

/**
 * Execute a command on an instance.
 * Returns the result and the updated instance.
 */
export async function executeCommand(
  instance: Instance,
  commandName: string,
  input: unknown,
  instanceId: string,
  history: MachineMessage<unknown>[],
  enqueue: (msgs: MachineMessage<unknown>[]) => void,
): Promise<{
  result: CommandExecutionResult;
  instance: Instance;
  transitionResult?: CommandResult;
  suspendInfo?: SuspendInfo;
  messages?: string | MachineMessage<unknown>[];
}> {
  const command = instance.node.commands?.[commandName];
  if (!command) {
    return {
      result: { success: false, error: `Command not found: ${commandName}` },
      instance,
    };
  }

  // Validate input
  const parsed = command.inputSchema.safeParse(input);
  if (!parsed.success) {
    return {
      result: { success: false, error: `Invalid input: ${parsed.error.message}` },
      instance,
    };
  }

  // Track state updates (with validation matching the tool code path)
  let currentState = instance.state;
  const updateState = (patch: Partial<unknown>) => {
    const merged = shallowMerge(
      currentState as Record<string, unknown>,
      patch as Record<string, unknown>,
    );
    const result = instance.node.validator.safeParse(merged);
    if (!result.success) {
      throw new Error(`Command state update validation failed: ${result.error.message}`);
    }
    currentState = result.data;
    // Enqueue state update message
    enqueue([instanceMessage(
      { kind: "state", instanceId, patch: patch as Record<string, unknown> },
    )]);
  };

  // Create getInstanceMessages function that filters by source.instanceId
  const getInstanceMessages = (): MachineMessage[] => {
    return history.filter(
      (msg) => msg.metadata?.source?.instanceId === instanceId
    );
  };

  // Create context with helpers
  const ctx: CommandContext<unknown> = {
    state: currentState,
    updateState,
    instanceId,
    getInstanceMessages,
    cede,
    spawn,
    suspend,
  };

  try {
    // Execute the command
    const cmdResult = await command.execute(parsed.data, ctx);

    // Handle command result - returns optional messages to enqueue and optional payload
    if (isCommandValueResult(cmdResult)) {
      const updatedInstance: Instance = { ...instance, state: currentState };
      return {
        result: { success: true, value: cmdResult.payload },
        instance: updatedInstance,
        messages: cmdResult.messages,
      };
    }

    // Handle transition result
    if (isTransitionToResult(cmdResult)) {
      const newNode = cmdResult.node;
      const newState = cmdResult.state ?? newNode.initialState;
      const newInstance = createInstance(
        newNode,
        newState,
        undefined,
        instance.packStates,
      );
      return {
        result: { success: true },
        instance: newInstance,
        transitionResult: cmdResult,
      };
    }

    // Handle spawn result
    if (isSpawnResult(cmdResult)) {
      const newChildren = cmdResult.children.map((target) =>
        createInstance(
          target.node,
          target.state ?? target.node.initialState,
        ),
      );
      const updatedInstance: Instance = {
        ...instance,
        state: currentState,
        children: newChildren.length === 0 ? undefined : newChildren,
      };
      return {
        result: { success: true },
        instance: updatedInstance,
        transitionResult: cmdResult,
      };
    }

    // Handle cede result
    if (isCedeResult(cmdResult)) {
      // Return the cede result - caller must handle removing this instance
      return {
        result: { success: true, value: cmdResult.content },
        instance: { ...instance, state: currentState },
        transitionResult: cmdResult,
      };
    }

    // Handle suspend result
    if (isSuspendResult(cmdResult)) {
      const suspendInfo = createSuspendInfo(cmdResult);
      const updatedInstance: Instance = {
        ...instance,
        state: currentState,
        suspended: suspendInfo,
      };
      return {
        result: { success: true },
        instance: updatedInstance,
        transitionResult: cmdResult,
        suspendInfo,
      };
    }

    // Handle resume result
    if (isResumeResult(cmdResult)) {
      const updatedInstance: Instance = {
        ...clearSuspension(instance),
        state: currentState,
      };
      return {
        result: { success: true },
        instance: updatedInstance,
        transitionResult: cmdResult,
      };
    }

    return {
      result: { success: false, error: "Unknown command result type" },
      instance,
    };
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    return {
      result: { success: false, error: message },
      instance,
    };
  }
}