import {z} from 'zod'

export type TextDoc = {text: string};

export type EmbeddingsV2Request = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, normalized?: (boolean | null) | undefined, truncate?: (boolean | null) | undefined, model: 'jina-embeddings-v2-base-en' | 'jina-embeddings-v2-base-zh' | 'jina-embeddings-v2-base-de' | 'jina-embeddings-v2-base-es' | 'jina-embeddings-v2-base-code', input: string | TextDoc | Array<string | TextDoc>};

export type EmbeddingsV3Request = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, normalized?: (boolean | null) | undefined, truncate?: (boolean | null) | undefined, model: 'jina-embeddings-v3', input: string | TextDoc | Array<string | TextDoc>, task?: ('retrieval.query' | 'retrieval.passage' | 'text-matching' | 'classification' | 'separation' | null) | undefined, late_chunking?: (boolean | null) | undefined, dimensions?: (number | null) | undefined};

export type ImageDoc = {image: string};

export type PDFDoc = {pdf: string};

export type EmbeddingsV4Request = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, model: 'jina-embeddings-v4', input: string | TextDoc | ImageDoc | PDFDoc | Array<string | TextDoc | ImageDoc>, truncate?: (boolean | null) | undefined, late_chunking?: (boolean | null) | undefined, dimensions?: (number | null) | undefined, task?: ('text-matching' | 'retrieval.query' | 'retrieval.passage' | 'code.query' | 'code.passage' | null) | undefined, return_multivector?: (boolean | null) | undefined, return_tokenized_input?: (boolean | null) | undefined};

export type VideoDoc = {video: string};

export type AudioDoc = {audio: string};

export type MergedContentGroup = {content: Array<TextDoc | ImageDoc | VideoDoc | AudioDoc>};

export type EmbeddingsV5Request = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, normalized?: (boolean | null) | undefined, truncate?: (boolean | null) | undefined, model: 'jina-embeddings-v5-text-nano' | 'jina-embeddings-v5-text-small' | 'jina-embeddings-v5-omni-small' | 'jina-embeddings-v5-omni-nano', input: string | TextDoc | ImageDoc | VideoDoc | AudioDoc | PDFDoc | Array<string | TextDoc | ImageDoc | VideoDoc | AudioDoc | MergedContentGroup>, task?: ('retrieval.query' | 'retrieval.passage' | 'text-matching' | 'clustering' | 'classification' | null) | undefined, dimensions?: (number | null) | undefined};

export type CodeEmbeddings500MRequest = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, model: 'jina-code-embeddings-0.5b', input: string | TextDoc | Array<string | TextDoc>, truncate?: (boolean | null) | undefined, dimensions?: (number | null) | undefined, task?: ('nl2code.query' | 'nl2code.passage' | 'qa.query' | 'qa.passage' | 'code2code.query' | 'code2code.passage' | 'code2nl.query' | 'code2nl.passage' | 'code2completion.query' | 'code2completion.passage' | null) | undefined};

export type CodeEmbeddings1500MRequest = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, model: 'jina-code-embeddings-1.5b', input: string | TextDoc | Array<string | TextDoc>, truncate?: (boolean | null) | undefined, dimensions?: (number | null) | undefined, task?: ('nl2code.query' | 'nl2code.passage' | 'qa.query' | 'qa.passage' | 'code2code.query' | 'code2code.passage' | 'code2nl.query' | 'code2nl.passage' | 'code2completion.query' | 'code2completion.passage' | null) | undefined};

export type ClipV1Request = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, normalized?: (boolean | null) | undefined, truncate?: (boolean | null) | undefined, model: 'jina-clip-v1', input: string | TextDoc | ImageDoc | PDFDoc | Array<string | TextDoc | ImageDoc>};

export type ClipV2Request = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, normalized?: (boolean | null) | undefined, truncate?: (boolean | null) | undefined, model: 'jina-clip-v2', input: string | TextDoc | ImageDoc | PDFDoc | Array<string | TextDoc | ImageDoc>, task?: ('retrieval.query' | null) | undefined, dimensions?: (number | null) | undefined};

export type ColbertV1Request = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, model: 'jina-colbert-v1-en', input: string | TextDoc | Array<string | TextDoc>, input_type?: ('query' | 'document' | null) | undefined};

export type ColbertV2Request = {embedding_type?: ('float' | 'base64' | 'binary' | 'ubinary' | Array<'float' | 'base64' | 'binary' | 'ubinary'> | null) | undefined, model: 'jina-colbert-v2', input: string | TextDoc | Array<string | TextDoc>, input_type?: ('query' | 'document' | null) | undefined, dimensions?: (64 | 96 | 128 | null) | undefined};

export type CreateApiV1EmbeddingsArgs = {body: EmbeddingsV2Request | EmbeddingsV3Request | EmbeddingsV4Request | EmbeddingsV5Request | CodeEmbeddings500MRequest | CodeEmbeddings1500MRequest | ClipV1Request | ClipV2Request | ColbertV1Request | ColbertV2Request};

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

export const embeddingUsage = z.object({total_tokens: z.number().int(), prompt_tokens: z.number().int(), image_tokens: z.number().int().nullable().optional(), audio_tokens: z.number().int().nullable().optional(), video_tokens: z.number().int().nullable().optional()});

export const eLSERV2EmbeddingUsage = z.object({total_tokens: z.number().int(), input_tokens: z.number().int(), output_tokens: z.number().int()});

export const baseUsage = z.object({total_tokens: z.number().int()});

export const singleEmbeddingData = z.object({object: z.string().optional(), index: z.number().int(), embedding: z.union([z.string(), z.array(z.number()), z.record(z.string(), z.number())])});

export const multiEmbeddingData = z.object({object: z.string().optional(), index: z.number().int(), embeddings: z.union([z.array(z.string()), z.array(z.array(z.number()))]), tokenized_input: z.array(z.string()).nullable().optional()});

export const embeddingResponse = z.object({model: z.string(), object: z.string().optional(), usage: z.union([embeddingUsage, eLSERV2EmbeddingUsage, baseUsage]), data: z.union([z.array(singleEmbeddingData), z.array(multiEmbeddingData), z.array(z.record(z.string(), z.number()))])});

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
};

/**
 * Embeddings
 * 
 * POST /v1/embeddings
 */
export const createApiV1Embeddings = async (args: CreateApiV1EmbeddingsArgs, options?: CreateApiV1EmbeddingsOptions) => {
const {baseUrl = "https://api.jina.ai", 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/embeddings${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(embeddingResponse, await res.json())
};