import { ConvexError, v } from "convex/values";

import type { Doc } from "./_generated/dataModel";
import { internalMutation, internalQuery } from "./_generated/server";

const flowRunStatus = v.union(
  v.literal("running"),
  v.literal("completed"),
  v.literal("failed"),
  v.literal("cancelled")
);

/**
 * Create a FlowRun keyed idempotently by its idempotencyKey
 * (`event:<eventId>:flow:<flowType>:<flowVersion>`). Retries reuse the existing
 * run rather than creating duplicate specialist executions.
 */
export const create = internalMutation({
  args: {
    agentId: v.optional(v.string()),
    flowType: v.string(),
    flowVersion: v.string(),
    idempotencyKey: v.string(),
    organizationId: v.id("organizations"),
    projectId: v.optional(v.id("projects")),
    sourceEventId: v.id("events"),
    threadId: v.optional(v.id("threads")),
    workId: v.optional(v.id("works")),
  },
  handler: async (ctx, args): Promise<Doc<"flowRuns">> => {
    // Idempotency: look up an existing run for the same source event + flow type.
    const existing = await ctx.db
      .query("flowRuns")
      .withIndex("by_organizationId_and_source_event", (q) =>
        q
          .eq("organizationId", args.organizationId)
          .eq("sourceEventId", args.sourceEventId)
      )
      .filter((q) => q.eq(q.field("flowType"), args.flowType))
      .unique();
    if (existing) {
      return existing;
    }
    const timestamp = Date.now();
    const runId = await ctx.db.insert("flowRuns", {
      agentId: args.agentId,
      attempt: 1,
      flowType: args.flowType,
      flowVersion: args.flowVersion,
      idempotencyKey: args.idempotencyKey,
      organizationId: args.organizationId,
      projectId: args.projectId,
      sourceEventId: args.sourceEventId,
      startedAt: timestamp,
      state: {},
      status: "running",
      threadId: args.threadId,
      updatedAt: timestamp,
      workId: args.workId,
    });
    const run = await ctx.db.get(runId);
    if (!run) {
      throw new ConvexError("Flow run could not be read after insert");
    }
    return run;
  },
});

/**
 * Patch a FlowRun's state/status. A terminal run cannot be revived by update.
 */
export const update = internalMutation({
  args: {
    finishedAt: v.optional(v.number()),
    flowRunId: v.id("flowRuns"),
    lastError: v.optional(v.string()),
    state: v.optional(v.any()),
    status: v.optional(flowRunStatus),
  },
  handler: async (ctx, args): Promise<Doc<"flowRuns">> => {
    const run = await ctx.db.get(args.flowRunId);
    if (!run) {
      throw new ConvexError("Flow run not found");
    }
    // Terminal runs are immutable.
    if (
      run.status === "completed" ||
      run.status === "cancelled" ||
      run.status === "failed"
    ) {
      return run;
    }
    const patch: Record<string, unknown> = { updatedAt: Date.now() };
    if (args.state !== undefined) {
      patch.state = args.state;
    }
    if (args.status !== undefined) {
      patch.status = args.status;
    }
    if (args.finishedAt !== undefined) {
      patch.finishedAt = args.finishedAt;
    }
    if (args.lastError !== undefined) {
      patch.lastError = args.lastError;
    }
    await ctx.db.patch(args.flowRunId, patch);
    const updated = await ctx.db.get(args.flowRunId);
    if (!updated) {
      throw new ConvexError("Flow run could not be read after update");
    }
    return updated;
  },
});

export const get = internalQuery({
  args: { flowRunId: v.id("flowRuns") },
  handler: async (ctx, args): Promise<Doc<"flowRuns"> | null> =>
    await ctx.db.get(args.flowRunId),
});

export { flowRunStatus };