import * as restate from "@restatedev/restate-sdk";

import type { RunnerFuncs, RunnerOptions } from "@karakeep/shared/queueing";
import { QueueRetryAfterError } from "@karakeep/shared/queueing";
import { tryCatch } from "@karakeep/shared/tryCatch";

import type { RunnerJobData, RunnerResult, SerializedError } from "./types";

function serializeError(error: Error): SerializedError {
  return {
    name: error.name,
    message: error.message,
    stack: error.stack,
  };
}

export function runnerServiceName(queueName: string): string {
  return `${queueName}-runner`;
}

export function buildRunnerService<T, R>(
  queueName: string,
  funcs: RunnerFuncs<T, R>,
  opts: RunnerOptions<T>,
) {
  return restate.service({
    name: runnerServiceName(queueName),
    options: {
      ingressPrivate: true,
      inactivityTimeout: {
        seconds: opts.timeoutSecs * 2,
      },
      // No retries at runner level - dispatcher handles retry logic
      retryPolicy: {
        maxAttempts: 1,
      },
      journalRetention: {
        days: 3,
      },
    },
    handlers: {
      run: async (
        ctx: restate.Context,
        jobData: RunnerJobData<T>,
      ): Promise<RunnerResult<R>> => {
        // Validate payload if validator provided
        let payload = jobData.data;
        if (opts.validator) {
          const res = opts.validator.safeParse(jobData.data);
          if (!res.success) {
            return {
              type: "error",
              error: {
                name: "ValidationError",
                message: res.error.message,
              },
            };
          }
          payload = res.data;
        }

        const res = await tryCatch(
          ctx
            .run(
              "main logic",
              async () => {
                const result = await tryCatch(
                  funcs.run({
                    id: jobData.id,
                    data: payload,
                    priority: jobData.priority,
                    runNumber: jobData.runNumber,
                    abortSignal: AbortSignal.timeout(
                      jobData.timeoutSecs * 1000,
                    ),
                  }),
                );
                if (result.error) {
                  if (result.error instanceof QueueRetryAfterError) {
                    return {
                      type: "rate_limit" as const,
                      delayMs: result.error.delayMs,
                    };
                  }
                  throw result.error;
                }
                return { type: "success" as const, value: result.data };
              },
              {
                maxRetryAttempts: 1,
              },
            )
            .orTimeout({
              seconds: jobData.timeoutSecs * 1.1,
            }),
        );

        if (res.error) {
          return {
            type: "error",
            error: serializeError(res.error),
          };
        }

        return res.data as RunnerResult<R>;
      },

      onCompleted: async (
        ctx: restate.Context,
        data: { job: RunnerJobData<T>; result: R },
      ): Promise<void> => {
        await ctx.run(
          "onComplete",
          async () => {
            await funcs.onComplete?.(
              {
                id: data.job.id,
                data: data.job.data,
                priority: data.job.priority,
                runNumber: data.job.runNumber,
                abortSignal: AbortSignal.timeout(data.job.timeoutSecs * 1000),
              },
              data.result,
            );
          },
          {
            maxRetryAttempts: 1,
          },
        );
      },

      onError: async (
        ctx: restate.Context,
        data: { job: RunnerJobData<T>; error: SerializedError },
      ): Promise<void> => {
        // Reconstruct the error
        const reconstructedError = Object.assign(
          new Error(data.error.message),
          {
            name: data.error.name,
            stack: data.error.stack,
          },
        );

        await ctx.run(
          "onError",
          async () => {
            await funcs.onError?.({
              id: data.job.id,
              data: data.job.data,
              priority: data.job.priority,
              runNumber: data.job.runNumber,
              numRetriesLeft: data.job.numRetriesLeft,
              error: reconstructedError,
            });
          },
          {
            maxRetryAttempts: 1,
          },
        );
      },
    },
  });
}