import { z } from 'zod';

import { AssetBase } from './SessionEvents.js';

export const JobEventBase = z
  .object({
    orgId: z.string().uuid(),
    userId: z.string().uuid().optional(),
    sessionId: z.string().uuid().optional(), // to tie into front end sessions
    clientTime: z.string().optional(),
    metadata: z.any().optional(),
    platformHost: z.string().optional(),
    parameters: z.any().optional()
  })
  .merge(AssetBase);

// machine specifications
export const WorkerParameters = z.object({
  machineType: z.string().optional(),
  cpuAllocation: z.number().optional(),
  memoryAllocation: z.number().optional(),
  gpuType: z.string().optional()
});

export enum JobEventType {
  RealtimeJob = 'realtime_task',
  BulkJobAction = 'bulk_job_action',
  BulkJobTask = 'bulk_job_task'
}

export const RealtimeJobEvent = JobEventBase.merge(
  z.object({
    eventType: z.literal(JobEventType.RealtimeJob),
    jobType: z.string(), // This is purposely not an enum, so as we add new task types, we do not have to update this SDK
    queueDuration: z.number().optional(),
    taskDuration: z.number().optional(),
    errorMessage: z.string().optional(),
    parameters: z.any().optional(),
    results: z.any().optional(),
    url: z.string().optional()
  })
).merge(WorkerParameters);
export type RealtimeJobEvent = z.infer<typeof RealtimeJobEvent>;

export enum BulkJobAction {
  Start = 'start',
  Cancel = 'cancel',
  Stop = 'stop'
}

export const BulkJobActionEvent = JobEventBase.merge(
  z.object({
    eventType: z.literal(JobEventType.BulkJobAction),
    jobId: z.string().uuid(),
    jobName: z.string(),
    jobType: z.string(), // This is purposely not an enum, so as we add new task types, we do not have to update this SDK
    jobAction: z.nativeEnum(BulkJobAction),
    taskCount: z.number(),
    jobPriority: z.number().optional(),
    jobPool: z.string().optional(),
    jobScheduler: z.string().optional(),
    parameters: z.any().optional()
  })
);
export type BulkJobActionEvent = z.infer<typeof BulkJobActionEvent>;

export enum TaskRunStatus {
  Success = 'success',
  Failed = 'failed',
  Cancelled = 'cancelled'
}

export const BulkJobTaskEvent = JobEventBase.merge(
  z.object({
    eventType: z.literal(JobEventType.BulkJobTask),
    jobType: z.string(), // This is purposely not an enum, so as we add new task types, we do not have to update this SDK
    jobId: z.string().uuid(),
    taskId: z.string().uuid(),
    runId: z.string().uuid(),
    runStatus: z.nativeEnum(TaskRunStatus),
    runDuration: z.number(),
    runStartTime: z.string().optional(),
    parameters: z.any().optional(),
    results: z.any().optional(),
    errorMessage: z.string().optional()
  })
).merge(WorkerParameters);
export type BulkJobTaskEvent = z.infer<typeof BulkJobTaskEvent>;

export const jobEventTypeToEventZod = {
  [JobEventType.RealtimeJob]: RealtimeJobEvent,
  [JobEventType.BulkJobAction]: BulkJobActionEvent,
  [JobEventType.BulkJobTask]: BulkJobTaskEvent
};

const JobEvent = z.union([
  RealtimeJobEvent,
  BulkJobActionEvent,
  BulkJobTaskEvent
]);

export type JobEvent = z.infer<typeof JobEvent>;