import { vWorkflowId } from '@convex-dev/workflow';
import { v } from 'convex/values';

import { internalMutation } from '../_generated/server';
import { jsonValueValidator } from '../lib/validators/json';
import { workflowManagers } from './engine';
import * as EngineHelpers from './helpers/engine';
import { safeShardIndex } from './helpers/engine/shard';
import { recoverStuckExecutions } from './helpers/recovery';

export const onWorkflowComplete = internalMutation({
  args: {
    workflowId: vWorkflowId,
    context: jsonValueValidator,
    result: jsonValueValidator,
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    await EngineHelpers.handleWorkflowComplete(ctx, args);
    return null;
  },
});

export const cleanupComponentWorkflow = internalMutation({
  args: {
    workflowId: vWorkflowId,
    shardIndex: v.optional(v.number()),
  },
  returns: v.null(),
  handler: async (ctx, args) => {
    const manager = workflowManagers[safeShardIndex(args.shardIndex)];
    await EngineHelpers.cleanupComponentWorkflow(manager, ctx, args.workflowId);
    return null;
  },
});

export const recoverStuck = internalMutation({
  args: {},
  returns: v.object({ recovered: v.number() }),
  handler: async (ctx) => {
    return await recoverStuckExecutions(ctx, workflowManagers);
  },
});
