import {
  vResultValidator,
  type RunResult,
} from "@convex-dev/workpool";
import { getConvexSize, type Infer, type Value, v } from "convex/values";
import { internalQuery } from "./_generated/server.js";

const PAYLOAD_REF_KIND = "@convex-dev/workflow/payloadRef";

export const vStoredPayloadRef = v.object({
  kind: v.literal(PAYLOAD_REF_KIND),
  payloadId: v.id("payloads"),
  size: v.number(),
});
export type StoredPayloadRef = Infer<typeof vStoredPayloadRef>;

const vOffloadedRunResult = v.union(
  v.object({
    kind: v.literal("success"),
    returnValue: vStoredPayloadRef,
  }),
  v.object({
    kind: v.literal("failed"),
    error: vStoredPayloadRef,
  }),
  v.object({
    kind: v.literal("canceled"),
  }),
);

export const vPersistedRunResult = v.union(vResultValidator, vOffloadedRunResult);
export type PersistedRunResult = Infer<typeof vPersistedRunResult>;

type PayloadReader = { db: any };
type PayloadWriter = { db: any };

export function isStoredPayloadRef(value: unknown): value is StoredPayloadRef {
  return (
    !!value &&
    typeof value === "object" &&
    "kind" in value &&
    (value as { kind?: unknown }).kind === PAYLOAD_REF_KIND &&
    "payloadId" in value &&
    typeof (value as { payloadId?: unknown }).payloadId === "string" &&
    "size" in value &&
    typeof (value as { size?: unknown }).size === "number"
  );
}

export async function storePayload(
  ctx: PayloadWriter,
  value: Value,
): Promise<StoredPayloadRef> {
  const size = getConvexSize(value);
  const payloadId = await ctx.db.insert("payloads", {
    value,
    size,
  });
  return {
    kind: PAYLOAD_REF_KIND,
    payloadId: payloadId as StoredPayloadRef["payloadId"],
    size,
  };
}

export async function loadPayload<T extends Value>(
  ctx: PayloadReader,
  valueOrRef: T | StoredPayloadRef,
): Promise<T> {
  if (!isStoredPayloadRef(valueOrRef)) {
    return valueOrRef;
  }
  const payload = await ctx.db.get(valueOrRef.payloadId);
  if (!payload) {
    throw new Error(`Payload not found: ${valueOrRef.payloadId}`);
  }
  return payload.value as T;
}

export async function deletePayload(
  ctx: PayloadWriter,
  valueOrRef: unknown,
): Promise<void> {
  if (!isStoredPayloadRef(valueOrRef)) {
    return;
  }
  await ctx.db.delete(valueOrRef.payloadId);
}

export async function storeRunResult(
  ctx: PayloadWriter,
  runResult: RunResult,
): Promise<PersistedRunResult> {
  switch (runResult.kind) {
    case "success":
      return {
        kind: "success",
        returnValue: await storePayload(ctx, runResult.returnValue as Value),
      };
    case "failed":
      return {
        kind: "failed",
        error: await storePayload(ctx, runResult.error),
      };
    case "canceled":
      return runResult;
  }
}

export async function loadRunResult(
  ctx: PayloadReader,
  runResult: PersistedRunResult | undefined,
): Promise<RunResult | undefined> {
  if (!runResult) {
    return undefined;
  }
  switch (runResult.kind) {
    case "success":
      return {
        kind: "success",
        returnValue: await loadPayload(ctx, runResult.returnValue as Value),
      };
    case "failed":
      if (!isStoredPayloadRef(runResult.error)) {
        return runResult as RunResult;
      }
      return {
        kind: "failed",
        error: (await loadPayload(ctx, runResult.error as Value)) as string,
      };
    case "canceled":
      return runResult;
  }
}

export async function deleteRunResult(
  ctx: PayloadWriter,
  runResult: PersistedRunResult | undefined,
): Promise<void> {
  if (!runResult) {
    return;
  }
  switch (runResult.kind) {
    case "success":
      await deletePayload(ctx, runResult.returnValue);
      break;
    case "failed":
      await deletePayload(ctx, runResult.error);
      break;
    case "canceled":
      break;
  }
}

export const load = internalQuery({
  args: {
    payloadId: v.id("payloads"),
  },
  returns: v.any(),
  handler: async (ctx, args) => {
    const payload = await ctx.db.get(args.payloadId);
    if (!payload) {
      throw new Error(`Payload not found: ${args.payloadId}`);
    }
    return payload.value;
  },
});