import { z } from 'zod';

import { get } from '../../../operators/get.js';
import { getById } from '../../../operators/getById.js';
import { Pagination } from '../../../shared.js';
import type { ThreekitAuthProps } from '../../../ThreekitAuthProps.js';
import { File } from '../../files/Files.js';
import { Route } from '../../Route.js';

export enum JobRunResultCode {
  Success = 'Success',
  Failed = 'Failed'
}

export const JobRun = z.object({
  id: z.string().uuid(),
  orgId: z.string().uuid(),
  jobId: z.string().uuid(),
  taskId: z.string().uuid(),

  createdBy: z.string().uuid().optional(), // I think this is a UUID?

  createdAt: z.string(),
  updatedAt: z.string(),

  startedAt: z.string().optional(),
  stoppedAt: z.string().optional(),
  wallClockTime: z.number().optional(),

  results: z.object({
    files: z.array(File).optional() // two common names are 'stdout' and 'stderr'
  }),
  resultCode: z.nativeEnum(JobRunResultCode).optional(),
  resultMessage: z.string().optional(),

  error: z.string().optional(),

  reportProgress: z.boolean(),
  progress: z.number(),

  // pod info
  zone: z.string().optional(),
  podId: z.string().optional(),
  instanceID: z.string().optional(),

  // machine type
  cpuPlatform: z.string().optional(),
  machineType: z.string().optional(),
  gpuType: z.string().optional(),
  preemptible: z.boolean().optional(),
  machine: z.object({}).optional(), // TODO: should be: machine: { [key: string]: any };

  // GPU resources
  GPUTime: z.number().optional(),

  // CPU resources
  requestCPUs: z.number().optional(),
  limitCPUs: z.number().optional(),
  nodeCPUs: z.number().optional(),
  CPUTime: z.number().optional(),

  // GPU resources
  requestRAM: z.number().optional(),
  limitRAM: z.number().optional(),
  nodeRAM: z.number().optional(),
  RAMUsage: z.number().optional()
});

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

export const JobRunListing = Pagination.merge(
  z.object({
    runs: z.array(JobRun)
  })
);
export type JobRunListing = z.infer<typeof JobRunListing>;

// the following can either be { jobId: string } or { taskId: string }, but not both.
// https://threekit.slack.com/archives/CDHMWSL80/p1712070371605319
export const QueryJobRunProps = z.union([
  z.object({ jobId: z.string().uuid() }),
  z.object({ taskId: z.string().uuid() })
]);
export type QueryJobRunProps = z.infer<typeof QueryJobRunProps>;

const API_ROUTE = '/api/jobs/runs';

export class JobRuns extends Route {
  constructor(auth: ThreekitAuthProps) {
    super(auth, API_ROUTE);
  }

  get(queryProps: QueryJobRunProps, pagination: Pagination = {}) {
    return get<QueryJobRunProps | Pagination, JobRunListing>(this.context, {
      ...queryProps,
      ...pagination
    });
  }

  getById(id: string) {
    const assetId = z.string().uuid().parse(id);
    return getById<JobRun>(this.context, assetId);
  }
}