import { vResultValidator } from "@convex-dev/workpool";
import { assert } from "convex-helpers";
import {
  paginationOptsValidator,
  type FunctionHandle,
  type PaginationResult,
} from "convex/server";
import { type Infer, type Value, v } from "convex/values";
import {
  internalQuery,
  mutation,
  type MutationCtx,
  query,
} from "./_generated/server.js";
import { type Logger, logLevel } from "./logging.js";
import { getWorkflow } from "./model.js";
import { getWorkpool } from "./pool.js";
import schema, {
  vOnComplete,
  type JournalEntry,
} from "./schema.js";
import { getDefaultLogger } from "./utils.js";
import {
  type WorkflowId,
  type OnCompleteArgs,
  type WorkflowStep,
  type EventId,
  vPaginationResult,
  vWorkflowStep,
  type SchedulerOptions,
  type PublicWorkflow,
  vPublicWorkflow,
} from "../types.js";
import { api, internal } from "./_generated/api.js";
import { formatErrorWithStack } from "../shared.js";
import type { Doc, Id } from "./_generated/dataModel.js";
import { paginator } from "convex-helpers/server/pagination";
import {
  deletePayload,
  deleteRunResult,
  loadPayload,
  loadRunResult,
  storePayload,
  storeRunResult,
} from "./payload.js";

const createArgs = v.object({
  workflowName: v.string(),
  workflowHandle: v.string(),
  workflowArgs: v.any(),
  maxParallelism: v.optional(v.number()),
  onComplete: v.optional(vOnComplete),
  startAsync: v.optional(v.boolean()),
  // TODO: ttl
});
export const create = mutation({
  args: createArgs,
  returns: v.id("workflows"),
  handler: createHandler,
});

export async function createHandler(
  ctx: MutationCtx,
  args: Infer<typeof createArgs>,
  schedulerOptions?: SchedulerOptions,
) {
  const console = await getDefaultLogger(ctx);
  await updateMaxParallelism(ctx, console, args.maxParallelism);
  const workflowArgs = await storePayload(ctx, args.workflowArgs as Value);
  const workflowId = await ctx.db.insert("workflows", {
    name: args.workflowName,
    workflowHandle: args.workflowHandle,
    args: workflowArgs,
    generationNumber: 0,
    onComplete: args.onComplete,
  });
  console.debug(
    `Created workflow ${workflowId}:`,
    args.workflowArgs,
    args.workflowHandle,
  );
  if (args.startAsync) {
    const workpool = await getWorkpool(ctx, args);
    await workpool.enqueueMutation(
      ctx,
      args.workflowHandle as FunctionHandle<"mutation">,
      { workflowId, generationNumber: 0 },
      {
        name: args.workflowName,
        onComplete: internal.pool.handlerOnComplete,
        context: { workflowId, generationNumber: 0 },
        ...schedulerOptions,
      },
    );
  } else {
    // If we can't start it, may as well not create it, eh? Fail fast...
    await ctx.runMutation(args.workflowHandle as FunctionHandle<"mutation">, {
      workflowId,
      generationNumber: 0,
    });
  }
  return workflowId;
}

export const getStatus = query({
  args: {
    workflowId: v.id("workflows"),
  },
  returns: v.object({
    workflow: v.any(),
    inProgress: v.any(),
    logLevel: logLevel,
  }),
  handler: async (ctx, args) => {
    const workflow = await ctx.db.get(args.workflowId);
    assert(workflow, `Workflow not found: ${args.workflowId}`);
    const console = await getDefaultLogger(ctx);

    const inProgress = await ctx.db
      .query("steps")
      .withIndex("inProgress", (q) =>
        q.eq("step.inProgress", true).eq("workflowId", args.workflowId),
      )
      .collect();
    console.debug(`${args.workflowId} blocked by`, inProgress);
    return {
      workflow: await hydrateWorkflow(ctx, workflow),
      inProgress: await Promise.all(
        inProgress.map((step) => hydrateJournalEntry(ctx, step)),
      ),
      logLevel: console.logLevel,
    };
  },
});

function publicWorkflowId(workflowId: Id<"workflows">): WorkflowId {
  return workflowId as any;
}

async function publicWorkflow(
  ctx: { db: any },
  workflow: Doc<"workflows">,
): Promise<PublicWorkflow> {
  const hydrated = await hydrateWorkflow(ctx, workflow);
  return {
    workflowId: publicWorkflowId(hydrated._id),
    name: hydrated.name,
    args: hydrated.args,
    context: hydrated.onComplete?.context,
    runResult: hydrated.runResult,
  } satisfies PublicWorkflow;
}

async function publicStep(
  ctx: { db: any },
  step: JournalEntry,
): Promise<WorkflowStep> {
  const hydrated = await hydrateJournalEntry(ctx, step);
  return {
    workflowId: publicWorkflowId(hydrated.workflowId),
    name: hydrated.step.name,
    stepId: hydrated._id,
    stepNumber: hydrated.stepNumber,

    args: hydrated.step.args,
    runResult: hydrated.step.runResult,

    startedAt: hydrated.step.startedAt,
    completedAt: hydrated.step.completedAt,

    ...(hydrated.step.kind === "event"
      ? {
          kind: "event",
          eventId: hydrated.step.eventId as unknown as EventId,
        }
      : hydrated.step.kind === "workflow"
        ? {
            kind: "workflow",
            nestedWorkflowId: publicWorkflowId(hydrated.step.workflowId!),
          }
        : hydrated.step.kind === "function"
          ? {
              kind: "function",
              workId: hydrated.step.workId,
            }
          : {
              kind: "sleep",
              workId: hydrated.step.workId!,
            }),
  } satisfies WorkflowStep;
}

async function hydrateWorkflow(
  ctx: { db: any },
  workflow: Doc<"workflows">,
) {
  return {
    ...workflow,
    args: await loadPayload(ctx, workflow.args as Value),
    runResult: await loadRunResult(ctx, workflow.runResult),
  };
}

async function hydrateJournalEntry(
  ctx: { db: any },
  entry: JournalEntry,
) {
  return {
    ...entry,
    step: {
      ...entry.step,
      args: await loadPayload(ctx, entry.step.args as Value),
      runResult: await loadRunResult(ctx, entry.step.runResult),
    },
  };
}

export const list = query({
  args: {
    order: v.union(v.literal("asc"), v.literal("desc")),
    paginationOpts: paginationOptsValidator,
  },
  returns: vPaginationResult(vPublicWorkflow),
  handler: async (ctx, args) => {
    const result = await paginator(ctx.db, schema)
      .query("workflows")
      .order(args.order)
      .paginate(args.paginationOpts);
    return {
      ...result,
      page: await Promise.all(
        result.page.map((workflow) => publicWorkflow(ctx, workflow)),
      ),
    } as PaginationResult<Infer<typeof vPublicWorkflow>>;
  },
});

export const listByName = query({
  args: {
    name: v.string(),
    order: v.union(v.literal("asc"), v.literal("desc")),
    paginationOpts: paginationOptsValidator,
  },
  returns: vPaginationResult(vPublicWorkflow),
  handler: async (ctx, args) => {
    const result = await paginator(ctx.db, schema)
      .query("workflows")
      .withIndex("name", (q) => q.eq("name", args.name))
      .order(args.order)
      .paginate(args.paginationOpts);
    return {
      ...result,
      page: await Promise.all(
        result.page.map((workflow) => publicWorkflow(ctx, workflow)),
      ),
    } as PaginationResult<Infer<typeof vPublicWorkflow>>;
  },
});

export const listSteps = query({
  args: {
    workflowId: v.id("workflows"),
    order: v.union(v.literal("asc"), v.literal("desc")),
    paginationOpts: paginationOptsValidator,
  },
  returns: vPaginationResult(vWorkflowStep),
  handler: async (ctx, args) => {
    const result = await paginator(ctx.db, schema)
      .query("steps")
      .withIndex("workflow", (q) => q.eq("workflowId", args.workflowId))
      .order(args.order)
      .paginate(args.paginationOpts);
    return {
      ...result,
      page: await Promise.all(result.page.map((step) => publicStep(ctx, step))),
    } as PaginationResult<Infer<typeof vWorkflowStep>>;
  },
});

const restartArgs = v.object({
  workflowId: v.id("workflows"),
  from: v.optional(v.union(v.number(), v.string())),
  startAsync: v.optional(v.boolean()),
});

export const restart = mutation({
  args: restartArgs,
  returns: v.null(),
  handler: restartHandler,
});

export async function restartHandler(
  ctx: MutationCtx,
  args: Infer<typeof restartArgs>,
) {
  const workflow = await ctx.db.get(args.workflowId);
  assert(workflow, `Workflow not found: ${args.workflowId}`);
  const console = await getDefaultLogger(ctx);

  if (!workflow.runResult) {
    throw new Error(`Workflow is still running: ${args.workflowId}`);
  }

  // Delete steps from the specified point
  if (args.from !== undefined) {
    if (typeof args.from === "number") {
      if (args.from < 0) {
        throw new Error(`Step number cannot be negative: ${args.from}`);
      }
      const stepsToDelete = await ctx.db
        .query("steps")
        .withIndex("workflow", (q) =>
          q
            .eq("workflowId", args.workflowId)
            .gte("stepNumber", args.from as number),
        )
        .collect();
      if (stepsToDelete.length === 0) {
        console.warn(
          `Step number ${args.from} not found in workflow ${args.workflowId}`,
        );
      }
      await deleteSteps(ctx, stepsToDelete);
    } else {
      // Walk backwards to find step by name, collecting steps to delete
      const stepsDesc = ctx.db
        .query("steps")
        .withIndex("workflow", (q) => q.eq("workflowId", args.workflowId))
        .order("desc");
      let found = false;
      const toDelete: Doc<"steps">[] = [];
      for await (const step of stepsDesc) {
        toDelete.push(step);
        if (step.step.name === args.from) {
          found = true;
          break;
        }
      }
      if (!found) {
        throw new Error(
          `Step "${args.from}" not found in workflow ${args.workflowId}`,
        );
      }
      await deleteSteps(ctx, toDelete);
    }
  }

  // Increment generation number and clear result
  const generationNumber = workflow.generationNumber + 1;
  await deleteRunResult(ctx, workflow.runResult);
  await ctx.db.patch(args.workflowId, {
    generationNumber,
    runResult: undefined,
  });

  console.event("retry", {
    workflowId: args.workflowId,
    name: workflow.name,
    from: args.from,
  });

  if (args.startAsync) {
    const workpool = await getWorkpool(ctx, {});
    await workpool.enqueueMutation(
      ctx,
      workflow.workflowHandle as FunctionHandle<"mutation">,
      { workflowId: args.workflowId, generationNumber },
      {
        name: workflow.name,
        onComplete: internal.pool.handlerOnComplete,
        context: { workflowId: args.workflowId, generationNumber },
      },
    );
  } else {
    await ctx.runMutation(
      workflow.workflowHandle as FunctionHandle<"mutation">,
      {
        workflowId: args.workflowId,
        generationNumber,
      },
    );
  }
}

export const cancel = mutation({
  args: {
    workflowId: v.id("workflows"),
  },
  returns: v.null(),
  handler: async (ctx, { workflowId }) => {
    const workflow = await ctx.db.get(workflowId);
    assert(workflow, `Workflow not found: ${workflowId}`);
    await completeHandler(ctx, {
      workflowId,
      generationNumber: workflow.generationNumber,
      runResult: { kind: "canceled" },
    });
  },
});

const completeArgs = v.object({
  workflowId: v.id("workflows"),
  generationNumber: v.number(),
  runResult: vResultValidator,
});

export const complete = mutation({
  args: completeArgs,
  returns: v.null(),
  handler: completeHandler,
});

// When the overall workflow completes (successfully or not).
export async function completeHandler(
  ctx: MutationCtx,
  args: Infer<typeof completeArgs>,
) {
  const workflow = await getWorkflow(
    ctx,
    args.workflowId,
    args.generationNumber,
  );
  const console = await getDefaultLogger(ctx);
  if (workflow.runResult) {
    throw new Error(`Workflow not running: ${workflow}`);
  }
  workflow.runResult = await storeRunResult(ctx, args.runResult);
  console.event("completed", {
    workflowId: workflow._id,
    name: workflow.name,
    status: args.runResult.kind,
    overallDurationMs: Date.now() - workflow._creationTime,
  });
  if (args.runResult.kind === "canceled") {
    // We bump it so no in-flight steps succeed / we don't race to complete.
    workflow.generationNumber += 1;
    // TODO: can we cancel these asynchronously if there's more than one?
    const inProgress = await ctx.db
      .query("steps")
      .withIndex("inProgress", (q) =>
        q.eq("step.inProgress", true).eq("workflowId", args.workflowId),
      )
      .collect();
    if (inProgress.length > 0) {
      const workpool = await getWorkpool(ctx, {});
      for (const { step } of inProgress) {
        if (!step.kind || step.kind === "function" || step.kind === "sleep") {
          if (step.workId) {
            await workpool.cancel(ctx, step.workId);
          }
        } else if (step.kind === "workflow") {
          if (step.workflowId) {
            await ctx.runMutation(api.workflow.cancel, {
              workflowId: step.workflowId,
            });
          }
        }
      }
    }
    console.debug(`Canceled workflow:`, workflow);
  }
  // Write the workflow so the onComplete can observe the updated status.
  await ctx.db.replace(workflow._id, workflow);
  if (workflow.onComplete) {
    try {
      await ctx.runMutation(
        workflow.onComplete.fnHandle as FunctionHandle<
          "mutation",
          OnCompleteArgs
        >,
        {
          workflowId: workflow._id as unknown as WorkflowId,
          result: args.runResult,
          context: workflow.onComplete.context,
        },
      );
    } catch (error) {
      const message = formatErrorWithStack(error);
      console.error("Error calling onComplete", message);
      await ctx.db.insert("onCompleteFailures", {
        ...args,
        runResult: await storeRunResult(ctx, args.runResult),
        error: message,
      });
    }
  }
  // TODO: delete everything unless ttl is set
  console.debug(`Completed workflow ${workflow._id}:`, workflow);
}

export const cleanup = mutation({
  args: {
    workflowId: v.string(),
    force: v.optional(v.boolean()),
  },
  returns: v.boolean(),
  handler: async (ctx, args) => {
    const workflowId = ctx.db.normalizeId("workflows", args.workflowId);
    if (!workflowId) {
      throw new Error(`Invalid workflow ID: ${args.workflowId}`);
    }
    const workflow = await ctx.db.get(workflowId);
    if (!workflow) {
      return false;
    }
    const logger = await getDefaultLogger(ctx);
    // TODO: allow cleaning up a workflow from inside it / in the onComplete hook
    if (!workflow.runResult) {
      if (!args.force) {
        logger.debug(
          `Can't clean up workflow ${workflowId} since it hasn't completed.`,
        );
        return false;
      }
      logger.debug(`Workflow ${workflowId} is not completed, forcing anyways`);
    }
    logger.debug(`Cleaning up workflow ${workflowId}`, workflow);
    await deletePayload(ctx, workflow.args);
    await deleteRunResult(ctx, workflow.runResult);
    await ctx.db.delete(workflowId);
    const journalEntries = await ctx.db
      .query("steps")
      .withIndex("workflow", (q) => q.eq("workflowId", workflowId))
      .collect();
    await deleteSteps(ctx, journalEntries);
    return true;
  },
});

async function updateMaxParallelism(
  ctx: MutationCtx,
  console: Logger,
  maxParallelism: number | undefined,
) {
  const config = await ctx.db.query("config").first();
  if (config) {
    if (maxParallelism && maxParallelism !== config.maxParallelism) {
      console.warn("Updating max parallelism to", maxParallelism);
      await ctx.db.patch(config._id, { maxParallelism });
    }
  } else {
    await ctx.db.insert("config", { maxParallelism });
  }
}

async function deleteSteps(ctx: MutationCtx, steps: Doc<"steps">[]) {
  for (const entry of steps) {
    await deletePayload(ctx, entry.step.args);
    await deleteRunResult(ctx, entry.step.runResult);
    await ctx.db.delete(entry._id);
    if (entry.step.kind === "event" && entry.step.eventId) {
      const event = await ctx.db.get(entry.step.eventId);
      if (event?.state.kind === "sent") {
        await deleteRunResult(ctx, event.state.result);
      }
      await ctx.db.delete(entry.step.eventId);
    } else if (entry.step.kind === "workflow" && entry.step.workflowId) {
      const workpool = await getWorkpool(ctx, {});
      await workpool.enqueueMutation(ctx, api.workflow.cleanup, {
        workflowId: entry.step.workflowId,
        force: true,
      });
    }
  }
}

export const sleep = internalQuery({
  args: {},
  returns: v.null(),
  handler: async () => null,
});

// eslint-disable-next-line @typescript-eslint/no-unused-vars
const console = "THIS IS A REMINDER TO USE getDefaultLogger";