import {
  createInsertSchema,
  createSelectSchema,
  createUpdateSchema,
} from "drizzle-zod";
import { z } from "zod";
import { schema } from "@/database";

export type TaskHandler = (payload: Record<string, unknown>) => Promise<void>;

export const TaskStatusSchema = z.enum([
  "pending",
  "processing",
  "completed",
  "dead",
]);
export type TaskStatus = z.infer<typeof TaskStatusSchema>;

export const TaskTypeSchema = z.enum([
  "connector_sync",
  "batch_embedding",
  "process_uploaded_files",
  "check_due_connectors",
  "check_due_schedule_triggers",
  "schedule_trigger_run_execute",
]);
export type TaskType = z.infer<typeof TaskTypeSchema>;

export type ConnectorSyncPayload = {
  connectorId: string;
  continuationCount?: number;
};
export type BatchEmbeddingPayload = {
  documentIds: string[];
  connectorRunId: string;
};

export const SelectTaskSchema = createSelectSchema(schema.tasksTable, {
  taskType: TaskTypeSchema,
  status: TaskStatusSchema,
});
export const InsertTaskSchema = createInsertSchema(schema.tasksTable, {
  taskType: TaskTypeSchema,
  status: TaskStatusSchema.optional(),
}).omit({ id: true, createdAt: true });
export const UpdateTaskSchema = createUpdateSchema(schema.tasksTable, {
  status: TaskStatusSchema.optional(),
}).pick({
  status: true,
  startedAt: true,
  completedAt: true,
  lastError: true,
  scheduledFor: true,
});

export type Task = z.infer<typeof SelectTaskSchema>;
export type InsertTask = z.infer<typeof InsertTaskSchema>;
export type UpdateTask = z.infer<typeof UpdateTaskSchema>;
