import { ChannelCredentials } from 'nice-grpc';
import { z } from 'zod/v4';
import type { Context } from '@hatchet/v1/client/worker/context';
import { Logger, LogLevel } from '@util/logger';

const ClientTLSConfigSchema = z.object({
  tls_strategy: z.enum(['tls', 'mtls', 'none']).optional(),
  cert_file: z.string().optional(),
  ca_file: z.string().optional(),
  key_file: z.string().optional(),
  server_name: z.string().optional(),
});

const HealthcheckConfigSchema = z.object({
  enabled: z.boolean().optional().default(false),
  port: z.number().optional().default(8001),
});

export const OpenTelemetryConfigSchema = z.object({
  /**
   * List of attribute keys to exclude from spans.
   * Useful for filtering sensitive or verbose data like payloads.
   */
  excludedAttributes: z.array(z.string()).optional().default([]),

  /**
   * If true, includes the task name in the span name for task run spans.
   * e.g., "hatchet.start_step_run.my_task" instead of "hatchet.start_step_run"
   */
  includeTaskNameInSpanName: z.boolean().optional().default(false),
});

export type OpenTelemetryConfig = z.infer<typeof OpenTelemetryConfigSchema>;

const TaskMiddlewareSchema = z
  .object({
    before: z.any().optional(),
    after: z.any().optional(),
  })
  .optional();

const DurationMsSchema = z.number().int().nonnegative().finite();

export const ClientConfigSchema = z.object({
  token: z.string(),
  tls_config: ClientTLSConfigSchema,
  healthcheck: HealthcheckConfigSchema.optional(),
  host_port: z.string(),
  api_url: z.string(),
  log_level: z.enum(['OFF', 'DEBUG', 'INFO', 'WARN', 'ERROR']).optional(),
  tenant_id: z.string(),
  namespace: z.string().optional(),
  otel: OpenTelemetryConfigSchema.optional(),
  middleware: TaskMiddlewareSchema,
  cancellation_grace_period: DurationMsSchema.optional().default(1000),
  cancellation_warning_threshold: DurationMsSchema.optional().default(300),
  grpc_max_recv_message_length: z
    .number()
    .int()
    .positive()
    .optional()
    .default(4 * 1024 * 1024),
  grpc_max_send_message_length: z
    .number()
    .int()
    .positive()
    .optional()
    .default(4 * 1024 * 1024),
});

export type LogConstructor = (context: string, logLevel?: LogLevel) => Logger;

/**
 * A middleware function that runs before every task invocation.
 * Returns extra fields to replace the task input, or void to skip.
 * @template T - The expected input type for the hook.
 * @param input - The current task input.
 * @param ctx - The task execution context.
 * @returns The new input value, or void to pass through unchanged.
 */
export type BeforeHookFn<T = any> = (
  input: T,
  ctx: Context<any>
) => Record<string, any> | void | Promise<Record<string, any> | void>;

/**
 * A middleware function that runs after every task invocation.
 * Returns extra fields to replace the task output, or void to skip.
 * @param output - The task output.
 * @param ctx - The task execution context.
 * @param input - The task input (after before-hooks have run).
 * @returns The new output value, or void to pass through unchanged.
 */
export type AfterHookFn<TOutput = any, TInput = any> = (
  output: TOutput,
  ctx: Context<any>,
  input: TInput
) => Record<string, any> | void | Promise<Record<string, any> | void>;

/**
 * Middleware hooks that run before/after every task invocation.
 *
 * Each hook can be a single function or an array of functions.
 * When an array is provided the functions run in order and each
 * result replaces the value (input for `before`, output for `after`).
 *
 * Return `void` (or `undefined`) from a hook to pass through unchanged.
 */
export type TaskMiddleware<TInput = any, TOutput = any> = {
  before?: BeforeHookFn<TInput> | readonly BeforeHookFn<TInput>[];
  after?: AfterHookFn<TOutput, TInput> | readonly AfterHookFn<TOutput, TInput>[];
};

type NonVoidReturn<F> = F extends (...args: any[]) => infer R
  ? Exclude<Awaited<R>, void | undefined>
  : {};

type MergeReturns<T> = T extends readonly [infer F, ...infer Rest]
  ? NonVoidReturn<F> & MergeReturns<Rest>
  : {};

export type InferMiddlewareBefore<M> = M extends { before: infer P }
  ? P extends (...args: any[]) => any
    ? NonVoidReturn<P>
    : P extends readonly any[]
      ? MergeReturns<P>
      : {}
  : {};

export type InferMiddlewareAfter<M> = M extends { after: infer P }
  ? P extends (...args: any[]) => any
    ? NonVoidReturn<P>
    : P extends readonly any[]
      ? MergeReturns<P>
      : {}
  : {};

type ClientConfigInferred = z.infer<typeof ClientConfigSchema>;

export type ClientConfig = Omit<
  ClientConfigInferred,
  | 'cancellation_grace_period'
  | 'cancellation_warning_threshold'
  | 'grpc_max_recv_message_length'
  | 'grpc_max_send_message_length'
> & {
  cancellation_grace_period?: number;
  cancellation_warning_threshold?: number;
  grpc_max_recv_message_length?: number;
  grpc_max_send_message_length?: number;
} & {
  credentials?: ChannelCredentials;
} & {
  logger: LogConstructor;
  middleware?: TaskMiddleware;
};
export type ClientTLSConfig = z.infer<typeof ClientTLSConfigSchema>;
