import z from 'zod';

const types = {
  GPT35Turbo: 'gpt-3.5-turbo',
  GPT35Turbo0613: 'gpt-3.5-turbo-0613',
  GPT35Turbo16K: 'gpt-3.5-turbo-16k',
  GPT4: 'gpt-4',
  GPT40613: 'gpt-4-0613',
  GPT41106: 'gpt-4-1106-preview',
  TextEmbeddingAda002: 'text-embedding-ada-002',
};

export const zGPTModels = z.nativeEnum(types);
export type GPTModels = z.infer<typeof zGPTModels>;

export const maxContextLength = (model: GPTModels): number => {
  switch (model) {
    case zGPTModels.enum.GPT35Turbo:
      return 4096;
    case zGPTModels.enum.GPT35Turbo0613:
      return 4096;
    case zGPTModels.enum.GPT35Turbo16K:
      return 16384;
    case zGPTModels.enum.GPT4:
      return 8192;
    case zGPTModels.enum.GPT40613:
      return 8192;
    case zGPTModels.enum.GPT41106:
      return 8192;
    case zGPTModels.enum.TextEmbeddingAda002:
      return 8192;
    default:
      return 0;
  }
};

export const queryContextLength = (model: GPTModels): number => {
  switch (model) {
    case zGPTModels.enum.GPT35Turbo:
      return 4000;
    case zGPTModels.enum.GPT35Turbo0613:
      return 4000;
    case zGPTModels.enum.GPT35Turbo16K:
      return 10000;
    case zGPTModels.enum.GPT4:
      return 8000;
    case zGPTModels.enum.GPT40613:
      return 8000;
    case zGPTModels.enum.GPT41106:
      return 8000;
    default:
      return 0;
  }
};

export const modelCostPerQuery = (model: GPTModels): number => {
  switch (model) {
    case zGPTModels.enum.GPT35Turbo:
      return 1;
    case zGPTModels.enum.GPT35Turbo0613:
      return 1;
    case zGPTModels.enum.GPT35Turbo16K:
      return 5;
    case zGPTModels.enum.GPT4:
      return 20;
    case zGPTModels.enum.GPT40613:
      return 20;
    case zGPTModels.enum.GPT41106:
      return 20;
    default:
      return 0;
  }
};