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

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

const dispatchStatus = v.union(
  v.literal("pending"),
  v.literal("sending"),
  v.literal("accepted"),
  v.literal("completed"),
  v.literal("failed")
);

/**
 * The handler-version stamp every dispatch carries. Bumped only when the
 * handler contract changes; encoded into the idempotency key so a new version
 * is a fresh dispatch rather than a duplicate of the old one.
 */
export const HANDLER_VERSION = "v1";

/**
 * Create an AgentDispatch row keyed idempotently by
 * `event:<sourceEventId>:agent:<agentId>:handler:<version>`. Retries reuse the
 * existing pending/sending dispatch rather than duplicating deliveries.
 * Returns `{ dispatchId, created }` so callers can decide whether to send.
 */
export const create = internalMutation({
  args: {
    agentId: 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<{ created: boolean; dispatch: Doc<"agentDispatches"> }> => {
    const existing = await ctx.db
      .query("agentDispatches")
      .withIndex("by_organizationId_and_source_event_handler", (q) =>
        q
          .eq("organizationId", args.organizationId)
          .eq("sourceEventId", args.sourceEventId)
          .eq("agentId", args.agentId)
          .eq("handlerVersion", HANDLER_VERSION)
      )
      .unique();
    if (existing) {
      return { created: false, dispatch: existing };
    }
    const timestamp = Date.now();
    const dispatchId = await ctx.db.insert("agentDispatches", {
      agentId: args.agentId,
      attempt: 1,
      createdAt: timestamp,
      handlerVersion: HANDLER_VERSION,
      idempotencyKey: args.idempotencyKey,
      organizationId: args.organizationId,
      projectId: args.projectId,
      sourceEventId: args.sourceEventId,
      status: "pending",
      threadId: args.threadId,
      updatedAt: timestamp,
      workId: args.workId,
    });
    const dispatch = await ctx.db.get(dispatchId);
    if (!dispatch) {
      throw new ConvexError("Agent dispatch could not be read after insert");
    }
    return { created: true, dispatch };
  },
});

/**
 * Transition a dispatch toward sending. Idempotent: a dispatch already accepted
 * or completed is not rewound back to sending on retry.
 */
export const markSending = internalMutation({
  args: { dispatchId: v.id("agentDispatches") },
  handler: async (ctx, args): Promise<Doc<"agentDispatches">> => {
    const dispatch = await ctx.db.get(args.dispatchId);
    if (!dispatch) {
      throw new ConvexError("Agent dispatch not found");
    }
    if (dispatch.status === "accepted" || dispatch.status === "completed") {
      return dispatch;
    }
    await ctx.db.patch(args.dispatchId, {
      status: "sending",
      updatedAt: Date.now(),
    });
    const updated = await ctx.db.get(args.dispatchId);
    if (!updated) {
      throw new ConvexError("Agent dispatch could not be read after update");
    }
    return updated;
  },
});

export const markAccepted = internalMutation({
  args: { dispatchId: v.id("agentDispatches") },
  handler: async (ctx, args): Promise<Doc<"agentDispatches">> => {
    const dispatch = await ctx.db.get(args.dispatchId);
    if (!dispatch) {
      throw new ConvexError("Agent dispatch not found");
    }
    // Completed is terminal; do not rewind to accepted.
    if (dispatch.status === "completed") {
      return dispatch;
    }
    await ctx.db.patch(args.dispatchId, {
      status: "accepted",
      updatedAt: Date.now(),
    });
    const updated = await ctx.db.get(args.dispatchId);
    if (!updated) {
      throw new ConvexError("Agent dispatch could not be read after update");
    }
    return updated;
  },
});

export const markCompleted = internalMutation({
  args: { dispatchId: v.id("agentDispatches") },
  handler: async (ctx, args): Promise<Doc<"agentDispatches">> => {
    const dispatch = await ctx.db.get(args.dispatchId);
    if (!dispatch) {
      throw new ConvexError("Agent dispatch not found");
    }
    if (dispatch.status === "completed") {
      return dispatch;
    }
    await ctx.db.patch(args.dispatchId, {
      status: "completed",
      updatedAt: Date.now(),
    });
    const updated = await ctx.db.get(args.dispatchId);
    if (!updated) {
      throw new ConvexError("Agent dispatch could not be read after update");
    }
    return updated;
  },
});

export const markFailed = internalMutation({
  args: {
    dispatchId: v.id("agentDispatches"),
    error: v.string(),
    retry: v.boolean(),
  },
  handler: async (ctx, args): Promise<Doc<"agentDispatches">> => {
    const dispatch = await ctx.db.get(args.dispatchId);
    if (!dispatch) {
      throw new ConvexError("Agent dispatch not found");
    }
    // Completed is terminal; a late failure after completion is ignored.
    if (dispatch.status === "completed") {
      return dispatch;
    }
    const nextAttempt = args.retry ? dispatch.attempt + 1 : dispatch.attempt;
    await ctx.db.patch(args.dispatchId, {
      attempt: nextAttempt,
      lastError: args.error,
      status: args.retry ? "pending" : "failed",
      updatedAt: Date.now(),
    });
    const updated = await ctx.db.get(args.dispatchId);
    if (!updated) {
      throw new ConvexError("Agent dispatch could not be read after update");
    }
    return updated;
  },
});

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

export { dispatchStatus };