import {z} from 'zod'

export type YQLParams = {query: string};

export type MappingRecordText = {type?: 'text' | undefined, value?: string | undefined};

export type MappingRecordColumn = {type?: 'column' | undefined, name?: string | undefined};

export type MappingRecordNumber = {type?: 'number' | undefined, value?: number | undefined};

export type MappingRecordConstant = {type?: 'constant' | undefined, value?: string | undefined};

export type BatchInferenceMessagesMappingResponse = {type?: string | undefined, messages: MappingRecordText | MappingRecordColumn, max_tokens?: (MappingRecordNumber | MappingRecordColumn | MappingRecordText | MappingRecordConstant | null) | undefined, custom_id?: (MappingRecordColumn | null) | undefined};

export type BatchInferenceInputMapping = {type?: string | undefined, input: MappingRecordColumn, encoding_format?: (MappingRecordText | MappingRecordColumn | null) | undefined, dimensions?: (MappingRecordNumber | MappingRecordColumn | null) | undefined, custom_id?: (MappingRecordColumn | null) | undefined};

export type OperationSrcDataset = {id: string, version: string | null, mapping?: (BatchInferenceMessagesMappingResponse | BatchInferenceInputMapping | null) | undefined};

export type OperationDstDataset = {id: string, version: string | null};

export type CreateYQLOperationRequest = {type?: 'yql' | undefined, params: YQLParams, src: Array<OperationSrcDataset>, dst?: Array<OperationDstDataset> | undefined, ai_project_id?: (string | null) | undefined};

export type MaterializeHistoryParams = {start_datetime: string, end_datetime: string, dataset_name: string, model_name?: (string | null) | undefined};

export type CreateMaterializeHistoryOperationRequest = {type?: 'materialize_history' | undefined, params: MaterializeHistoryParams, ai_project_id?: (string | null) | undefined};

export type BatchInferenceParams = {model: string, completion_window: string};

export type BatchInferenceMessagesMapping = {type?: string | undefined, messages: MappingRecordText | MappingRecordColumn, max_tokens?: (MappingRecordNumber | MappingRecordColumn | MappingRecordText | null) | undefined, custom_id?: (MappingRecordColumn | null) | undefined};

export type CreateOperationSrcDataset = {id: string, version: string | null, mapping?: (BatchInferenceMessagesMapping | BatchInferenceInputMapping | null) | undefined};

export type CreateBatchInferenceOperationRequest = {type?: 'batch_inference' | undefined, params: BatchInferenceParams, src: Array<CreateOperationSrcDataset>, dst?: Array<OperationDstDataset> | undefined, ai_project_id?: (string | null) | undefined};

export type SupervisedHParameters = {batch_size?: number | undefined, learning_rate?: number | undefined, n_epochs?: number | undefined, warmup_ratio?: number | undefined, weight_decay?: number | undefined, packing?: boolean | undefined, max_grad_norm?: number | undefined, context_length?: number | undefined, lora?: boolean | undefined, lora_r?: number | undefined, lora_alpha?: number | undefined, lora_dropout?: number | undefined} | {[key: string]: unknown};

export type SupervisedConfig = {hyperparameters: SupervisedHParameters};

export type SupervisedMethodConfig = {type?: 'supervised' | undefined, supervised: SupervisedConfig};

export type SpecDraftLossParams = {type?: ('kl' | 'lk_alpha' | 'lk_hybrid') | undefined};

export type SpecDraftHParameters = {batch_size?: number | undefined, learning_rate?: number | undefined, n_epochs?: number | undefined, warmup_ratio?: number | undefined, weight_decay?: number | undefined, packing?: boolean | undefined, max_grad_norm?: number | undefined, context_length?: number | undefined, architecture?: string | undefined, num_decoding_heads?: number | undefined, loss?: SpecDraftLossParams | undefined} | {[key: string]: unknown};

export type SpecDraftConfig = {hyperparameters: SpecDraftHParameters};

export type SpecDraftMethodConfig = {type?: 'spec-draft' | undefined, spec_draft: SpecDraftConfig};

export type HFCheckpointParametersRequestInput = {type?: (string | 'hf') | undefined, repo: string, revision?: (string | null) | undefined, token?: (string | null) | undefined};

export type FTCheckpointParameters = {type?: (string | 'ft') | undefined};

export type HFCheckpointParametersResponse = {type?: (string | 'hf') | undefined, repo: string, revision?: (string | null) | undefined};

export type WandbConfigRequest = {project: string, name?: (string | null) | undefined, entity?: (string | null) | undefined, tags?: Array<string> | undefined, api_key: string};

export type WandbIntegrationRequest = {type?: (string | 'wandb') | undefined, wandb: WandbConfigRequest};

export type MlflowConfigRequest = {tracking_uri: string, experiment_name: string | null, run_name: string | null, tags?: {[key: string]: string} | undefined, certificate_file: string, certificate_file_id?: (string | null) | undefined, username: string, password: string};

export type MlflowIntegrationRequest = {type?: 'mlflow' | undefined, mlflow: MlflowConfigRequest};

export type HfExportConfigRequest = {output_repo_name: string, api_token: string};

export type HfExportIntegrationRequest = {type?: (string | 'hf') | undefined, hf: HfExportConfigRequest};

export type WandbConfigResponse = {project: string, name?: (string | null) | undefined, entity?: (string | null) | undefined, tags?: Array<string> | undefined};

export type WandbIntegrationResponse = {type?: (string | 'wandb') | undefined, wandb: WandbConfigResponse};

export type MlflowConfigResponse = {tracking_uri: string, experiment_name: string | null, run_name: string | null, tags?: {[key: string]: string} | undefined};

export type MlflowIntegrationResponse = {type?: 'mlflow' | undefined, mlflow: MlflowConfigResponse};

export type HfExportConfigResponse = {output_repo_name: string};

export type HfExportIntegrationResponse = {type?: (string | 'hf') | undefined, hf: HfExportConfigResponse};

export type FineTuningParamsInput = {model_name: string, hyperparameters?: (SupervisedHParameters | null) | undefined, method?: (SupervisedMethodConfig | SpecDraftMethodConfig | null) | undefined, from_checkpoint?: (HFCheckpointParametersRequestInput | FTCheckpointParameters | HFCheckpointParametersResponse | null) | undefined, integrations?: (Array<WandbIntegrationRequest | MlflowIntegrationRequest | HfExportIntegrationRequest | WandbIntegrationResponse | MlflowIntegrationResponse | HfExportIntegrationResponse> | null) | undefined, seed?: (number | null) | undefined, suffix?: (string | null) | undefined, tags?: (Array<string> | null) | undefined, trained_steps?: (number | null) | undefined, total_steps?: (number | null) | undefined, trained_tokens?: (number | null) | undefined, estimated_finish?: (number | null) | undefined};

export type FineTuningTextMapping = {type?: 'text' | undefined, text: MappingRecordColumn};

export type FineTuningPromptsMapping = {type?: 'prompts' | undefined, prompt: MappingRecordColumn, completion: MappingRecordColumn};

export type FineTuningMessagesMapping = {type?: 'messages' | undefined, messages: MappingRecordColumn};

export type FineTuningPretokenizedMapping = {type?: 'pretokenized' | undefined, input_ids: MappingRecordColumn, labels?: (MappingRecordColumn | null) | undefined, attention_mask?: (MappingRecordColumn | null) | undefined};

export type OperationSrcTrainDataset = {id: string, version: string | null, mapping: FineTuningTextMapping | FineTuningPromptsMapping | FineTuningMessagesMapping | FineTuningPretokenizedMapping, type?: 'train' | undefined};

export type OperationSrcValidationDataset = {id: string, version: string | null, mapping: FineTuningTextMapping | FineTuningPromptsMapping | FineTuningMessagesMapping | FineTuningPretokenizedMapping, type?: 'validation' | undefined};

export type CreateFineTuningOperationRequest = {type?: 'finetuning' | undefined, params: FineTuningParamsInput, src: Array<OperationSrcTrainDataset | OperationSrcValidationDataset>, ai_project_id?: (string | null) | undefined};

export type FileToDatasetParams = {file_id: string, dataset_name?: (string | null) | undefined};

export type CreateFileToDatasetOperationRequest = {type?: 'file_to_dataset' | undefined, params: FileToDatasetParams, ai_project_id?: (string | null) | undefined};

export type DatasetToFilePurpose = 'fine_tuning' | 'inference';

export type DatasetToFileParams = {purpose: DatasetToFilePurpose, model?: (string | null) | undefined};

export type OperationDatasetToFileDataset = {id: string, version: string};

export type DatasetToFileBatchInferenceMapping = {type?: 'text_messages' | undefined, messages: MappingRecordColumn, max_tokens?: (MappingRecordColumn | null) | undefined, custom_id?: (MappingRecordColumn | null) | undefined};

export type OperationSrcDatasetToFile = {dataset: OperationDatasetToFileDataset, mapping: FineTuningTextMapping | FineTuningPromptsMapping | FineTuningMessagesMapping | DatasetToFileBatchInferenceMapping};

export type CreateDatasetToFileOperationRequest = {type?: 'dataset_to_file' | undefined, params: DatasetToFileParams, src: Array<OperationSrcDatasetToFile>, dst?: Array<OperationDstDataset> | undefined, ai_project_id?: (string | null) | undefined};

export type MergeParams = {dataset_name?: (string | null) | undefined};

export type CreateMergeOperationRequest = {type?: 'merge' | undefined, params?: MergeParams | undefined, src: Array<OperationSrcDataset>, ai_project_id?: (string | null) | undefined};

export type CreateApiV1OperationsArgs = {body: CreateYQLOperationRequest | CreateMaterializeHistoryOperationRequest | CreateBatchInferenceOperationRequest | CreateFineTuningOperationRequest | CreateFileToDatasetOperationRequest | CreateDatasetToFileOperationRequest | CreateMergeOperationRequest};

export type CreateApiV1OperationsOptions = {baseUrl?: string | undefined, HTTPBearer?: string | undefined, fetch?: typeof fetch | undefined};

export const operationType = z.enum(["yql", "materialize_history", "batch_inference", "finetuning", "file_to_dataset", "dataset_to_file", "merge"]);

export const yQLParams = z.object({query: z.string()});

export const materializeHistoryParams = z.object({start_datetime: z.string(), end_datetime: z.string(), dataset_name: z.string(), model_name: z.string().nullable().optional()});

export const batchInferenceParams = z.object({model: z.string(), completion_window: z.string()});

export const supervisedHParameters = z.object({batch_size: z.number().int().gte(1).lte(64).optional(), learning_rate: z.number().gte(0).optional(), n_epochs: z.number().int().gte(1).lte(20).optional(), warmup_ratio: z.number().gte(0).lte(1).optional(), weight_decay: z.number().gte(0).optional(), packing: z.boolean().optional(), max_grad_norm: z.number().gt(0).optional(), context_length: z.number().int().gte(8192).lte(131072).optional(), lora: z.boolean().optional(), lora_r: z.number().int().gte(8).lte(128).optional(), lora_alpha: z.number().int().gte(8).optional(), lora_dropout: z.number().gte(0).lte(1).optional()}).and(z.record(z.string(), z.unknown()));

export const supervisedConfig = z.object({hyperparameters: supervisedHParameters});

export const supervisedMethodConfig = z.object({type: z.literal("supervised").optional(), supervised: supervisedConfig});

export const specDraftLossParams = z.object({type: z.enum(["kl", "lk_alpha", "lk_hybrid"]).optional()});

export const specDraftHParameters = z.object({batch_size: z.number().int().gte(1).lte(64).optional(), learning_rate: z.number().gt(0).lt(1).optional(), n_epochs: z.number().int().gte(1).lte(20).optional(), warmup_ratio: z.number().gte(0).lte(1).optional(), weight_decay: z.number().gte(0).optional(), packing: z.boolean().optional(), max_grad_norm: z.number().gt(0).optional(), context_length: z.number().int().gte(8192).lte(131072).optional(), architecture: z.string().optional(), num_decoding_heads: z.number().int().gte(1).lte(10).optional(), loss: specDraftLossParams.optional()}).and(z.record(z.string(), z.unknown()));

export const specDraftConfig = z.object({hyperparameters: specDraftHParameters});

export const specDraftMethodConfig = z.object({type: z.literal("spec-draft").optional(), spec_draft: specDraftConfig});

export const hFCheckpointParametersRequestOutput = z.object({type: z.union([z.string(), z.literal("hf")]).optional(), repo: z.string(), revision: z.string().nullable().optional(), token: z.string().nullable().optional()});

export const fTCheckpointParameters = z.object({type: z.union([z.string(), z.literal("ft")]).optional()});

export const hFCheckpointParametersResponse = z.object({type: z.union([z.string(), z.literal("hf")]).optional(), repo: z.string(), revision: z.string().nullable().optional()});

export const wandbConfigRequest = z.object({project: z.string(), name: z.string().nullable().optional(), entity: z.string().nullable().optional(), tags: z.array(z.string()).optional(), api_key: z.string()});

export const wandbIntegrationRequest = z.object({type: z.union([z.string(), z.literal("wandb")]).optional(), wandb: wandbConfigRequest});

export const mlflowConfigRequest = z.object({tracking_uri: z.string(), experiment_name: z.string().nullable(), run_name: z.string().nullable(), tags: z.record(z.string(), z.string()).optional(), certificate_file: z.string(), certificate_file_id: z.string().nullable().optional(), username: z.string(), password: z.string()});

export const mlflowIntegrationRequest = z.object({type: z.literal("mlflow").optional(), mlflow: mlflowConfigRequest});

export const hfExportConfigRequest = z.object({output_repo_name: z.string(), api_token: z.string()});

export const hfExportIntegrationRequest = z.object({type: z.union([z.string(), z.literal("hf")]).optional(), hf: hfExportConfigRequest});

export const wandbConfigResponse = z.object({project: z.string(), name: z.string().nullable().optional(), entity: z.string().nullable().optional(), tags: z.array(z.string()).optional()});

export const wandbIntegrationResponse = z.object({type: z.union([z.string(), z.literal("wandb")]).optional(), wandb: wandbConfigResponse});

export const mlflowConfigResponse = z.object({tracking_uri: z.string(), experiment_name: z.string().nullable(), run_name: z.string().nullable(), tags: z.record(z.string(), z.string()).optional()});

export const mlflowIntegrationResponse = z.object({type: z.literal("mlflow").optional(), mlflow: mlflowConfigResponse});

export const hfExportConfigResponse = z.object({output_repo_name: z.string()});

export const hfExportIntegrationResponse = z.object({type: z.union([z.string(), z.literal("hf")]).optional(), hf: hfExportConfigResponse});

export const fineTuningParamsOutput = z.object({model_name: z.string(), hyperparameters: supervisedHParameters.nullable().optional(), method: z.union([supervisedMethodConfig, specDraftMethodConfig]).nullable().optional(), from_checkpoint: z.union([hFCheckpointParametersRequestOutput, fTCheckpointParameters, hFCheckpointParametersResponse]).nullable().optional(), integrations: z.array(z.union([wandbIntegrationRequest, mlflowIntegrationRequest, hfExportIntegrationRequest, wandbIntegrationResponse, mlflowIntegrationResponse, hfExportIntegrationResponse])).nullable().optional(), seed: z.number().int().nullable().optional(), suffix: z.string().max(64).nullable().optional(), tags: z.array(z.string()).nullable().optional(), trained_steps: z.number().int().nullable().optional(), total_steps: z.number().int().nullable().optional(), trained_tokens: z.number().int().nullable().optional(), estimated_finish: z.number().int().nullable().optional()});

export const fileToDatasetParams = z.object({file_id: z.string(), dataset_name: z.string().nullable().optional()});

export const datasetToFilePurpose = z.enum(["fine_tuning", "inference"]);

export const datasetToFileParams = z.object({purpose: datasetToFilePurpose, model: z.string().nullable().optional()});

export const mergeParams = z.object({dataset_name: z.string().nullable().optional()});

export const mappingRecordText = z.object({type: z.literal("text").optional(), value: z.string().optional()});

export const mappingRecordColumn = z.object({type: z.literal("column").optional(), name: z.string().optional()});

export const mappingRecordNumber = z.object({type: z.literal("number").optional(), value: z.number().int().optional()});

export const mappingRecordConstant = z.object({type: z.literal("constant").optional(), value: z.string().optional()});

export const batchInferenceMessagesMappingResponse = z.object({type: z.string().optional(), messages: z.union([mappingRecordText, mappingRecordColumn]), max_tokens: z.union([mappingRecordNumber, mappingRecordColumn, mappingRecordText, mappingRecordConstant]).nullable().optional(), custom_id: mappingRecordColumn.nullable().optional()});

export const batchInferenceInputMapping = z.object({type: z.string().optional(), input: mappingRecordColumn, encoding_format: z.union([mappingRecordText, mappingRecordColumn]).nullable().optional(), dimensions: z.union([mappingRecordNumber, mappingRecordColumn]).nullable().optional(), custom_id: mappingRecordColumn.nullable().optional()});

export const operationSrcDataset = z.object({id: z.string(), version: z.string().nullable(), mapping: z.union([batchInferenceMessagesMappingResponse, batchInferenceInputMapping]).nullable().optional()});

export const fineTuningTextMapping = z.object({type: z.literal("text").optional(), text: mappingRecordColumn});

export const fineTuningPromptsMapping = z.object({type: z.literal("prompts").optional(), prompt: mappingRecordColumn, completion: mappingRecordColumn});

export const fineTuningMessagesMapping = z.object({type: z.literal("messages").optional(), messages: mappingRecordColumn});

export const fineTuningPretokenizedMapping = z.object({type: z.literal("pretokenized").optional(), input_ids: mappingRecordColumn, labels: mappingRecordColumn.nullable().optional(), attention_mask: mappingRecordColumn.nullable().optional()});

export const operationSrcTrainDataset = z.object({id: z.string(), version: z.string().nullable(), mapping: z.union([fineTuningTextMapping, fineTuningPromptsMapping, fineTuningMessagesMapping, fineTuningPretokenizedMapping]), type: z.literal("train").optional()});

export const operationSrcValidationDataset = z.object({id: z.string(), version: z.string().nullable(), mapping: z.union([fineTuningTextMapping, fineTuningPromptsMapping, fineTuningMessagesMapping, fineTuningPretokenizedMapping]), type: z.literal("validation").optional()});

export const operationDatasetToFileDataset = z.object({id: z.string(), version: z.string()});

export const datasetToFileBatchInferenceMapping = z.object({type: z.literal("text_messages").optional(), messages: mappingRecordColumn, max_tokens: mappingRecordColumn.nullable().optional(), custom_id: mappingRecordColumn.nullable().optional()});

export const operationSrcDatasetToFile = z.object({dataset: operationDatasetToFileDataset, mapping: z.union([fineTuningTextMapping, fineTuningPromptsMapping, fineTuningMessagesMapping, datasetToFileBatchInferenceMapping])});

export const operationDstDataset = z.object({id: z.string(), version: z.string().nullable()});

export const operationStatus = z.enum(["queued", "running", "succeeded", "failed", "cancelled", "unknown"]);

export const createOperationResponse = z.object({id: z.string(), type: operationType, params: z.union([yQLParams, materializeHistoryParams, batchInferenceParams, fineTuningParamsOutput, fileToDatasetParams, datasetToFileParams, mergeParams]), src: z.array(z.union([operationSrcDataset, operationSrcTrainDataset, operationSrcValidationDataset, operationSrcDatasetToFile])), dst: z.array(operationDstDataset), status: operationStatus, created_at: z.number().int(), in_progress_at: z.number().int().nullable().optional(), finished_at: z.number().int().nullable().optional(), ai_project_id: z.string().nullable().optional()});

export const validateResponse = async <Output>(
  schema: {
    readonly '~standard': {
      readonly validate: (value: unknown) => ({ readonly value: Output; readonly issues?: undefined } | { readonly issues: ReadonlyArray<{ readonly message: string; readonly path?: ReadonlyArray<PropertyKey | { readonly key: PropertyKey }> | undefined }> }) | Promise<{ readonly value: Output; readonly issues?: undefined } | { readonly issues: ReadonlyArray<{ readonly message: string; readonly path?: ReadonlyArray<PropertyKey | { readonly key: PropertyKey }> | undefined }> }>
    }
  },
  value: unknown
): Promise<Output> => {
  const result = await schema['~standard'].validate(value)

  if (result.issues !== undefined) {
    const detail = result.issues
      .map(({ message, path }) =>
        path?.length
          ? `${path.map((segment) => String(typeof segment === 'object' ? segment.key : segment)).join('.')}: ${message}`
          : message
      )
      .join('; ')

    throw new Error(`Response does not match the schema: ${detail}`, { cause: { issues: result.issues } })
  }

  return result.value
};

/**
 * Run operation
 * 
 * POST /v1/operations
 */
export const createApiV1Operations = async (args: CreateApiV1OperationsArgs, options?: CreateApiV1OperationsOptions) => {
const {baseUrl = "https://api.tokenfactory.nebius.com", fetch: fetchFn = fetch, HTTPBearer} = options ?? {}

const query = new URLSearchParams()
const headers = new Headers()
headers.set('content-type', "application/json")
headers.set('accept', "application/json")

if (HTTPBearer !== undefined) headers.set('authorization', `Bearer ${HTTPBearer}`)

const url = `${baseUrl.replace(/\/+$/, '')}/v1/operations${query.size > 0 ? `?${query}` : ''}`

const res = await fetchFn(url, {method: "POST", headers, body: JSON.stringify(args.body)})

if (!res.ok) throw new Error(`Request failed with status ${res.status}`, {cause: {status: res.status, body: await res.text()}})

return validateResponse(createOperationResponse, await res.json())
};