import {z} from 'zod'

export type DataType = 'fp16' | 'fp32' | 'int4' | 'int8';

export type AssetType = 'file' | 'checkpoint' | 'lora' | 'textual_inversion' | 'vae' | 'latent';

export type BaseEngine = 'text/llama-2-7b' | 'text/llama-2-13b' | 'text/llama-2-70b' | 'text/llama-3-8b' | 'text/llama-3-70b' | 'text/codellama-7b' | 'text/codellama-13b' | 'text/codellama-34b' | 'text/codellama-70b' | 'text/mistral-7b' | 'text/mistral-nemo' | 'text/mixtral-8x7b' | 'text/llama-3.1-8b' | 'text/llama-3.1-70b' | 'text/llama-3.1-405b' | 'image/controlnet-sd15' | 'image/controlnet-sdxl' | 'image/stable-diffusion-v1-5' | 'image/stable-diffusion-xl-v1-0' | 'image/stable-diffusion-v3-medium' | 'image/stable-diffusion-ssd' | 'image/flux-dev' | 'image/flux-schnell' | 'audio/tortoise-tts' | 'audio/rvc' | 'video/svd';

export type BaseEngineType = 'text' | 'image' | 'audio' | 'video';

export type FileFormat = 'safetensors' | 'pt';

export type FileExtension = 'unknown' | 'png' | 'jpg' | 'jpeg' | 'txt' | 'mp3' | 'wav' | 'jsonl';

export type GetApiV1AssetsArgs = {name?: (string | null) | undefined, is_public?: boolean | undefined, data_type?: (DataType | null) | undefined, asset_type?: Array<AssetType> | undefined, engine?: Array<BaseEngine> | undefined, engine_type?: (BaseEngineType | null) | undefined, limit?: number | undefined, offset?: number | undefined, get_preview_urls?: boolean | undefined, asset_ids?: (string | null) | undefined, owner?: (string | null) | undefined, file_format?: (Array<FileFormat | FileExtension>) | undefined};

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

export const assetType = z.enum(["file", "checkpoint", "lora", "textual_inversion", "vae", "latent"]);

export const fileExtension = z.enum(["unknown", "png", "jpg", "jpeg", "txt", "mp3", "wav", "jsonl"]);

export const fileData = z.object({asset_type: z.literal("file").optional(), file_format: fileExtension.optional(), preview_url: z.string().nullable().optional(), version: z.string().optional()});

export const dataType = z.enum(["fp16", "fp32", "int4", "int8"]);

export const baseEngine = z.enum(["text/llama-2-7b", "text/llama-2-13b", "text/llama-2-70b", "text/llama-3-8b", "text/llama-3-70b", "text/codellama-7b", "text/codellama-13b", "text/codellama-34b", "text/codellama-70b", "text/mistral-7b", "text/mistral-nemo", "text/mixtral-8x7b", "text/llama-3.1-8b", "text/llama-3.1-70b", "text/llama-3.1-405b", "image/controlnet-sd15", "image/controlnet-sdxl", "image/stable-diffusion-v1-5", "image/stable-diffusion-xl-v1-0", "image/stable-diffusion-v3-medium", "image/stable-diffusion-ssd", "image/flux-dev", "image/flux-schnell", "audio/tortoise-tts", "audio/rvc", "video/svd"]);

export const fileFormat = z.enum(["safetensors", "pt"]);

export const checkpointData = z.object({asset_type: z.literal("checkpoint").optional(), data_type: dataType.optional(), engine: baseEngine, file_format: fileFormat.optional()});

export const loraData = z.object({asset_type: z.literal("lora").optional(), checkpoint: z.string().nullable().optional(), data_type: dataType.optional(), engine: baseEngine, file_format: fileFormat.optional(), trigger_words: z.array(z.string()).nullable().optional()});

export const textualInversionData = z.object({asset_type: z.literal("textual_inversion").optional(), data_type: dataType.optional(), engine: baseEngine, file_format: fileFormat.optional(), trigger_words: z.array(z.string()).optional()});

export const vAEData = z.object({asset_type: z.literal("vae").optional(), data_type: dataType.optional(), engine: baseEngine, file_format: fileFormat.optional()});

export const latentData = z.object({asset_type: z.literal("latent").optional(), engine: baseEngine, file_format: fileFormat.optional()});

export const fileStructure = z.enum(["unknown", "single_file", "multiple_files"]);

export const status = z.enum(["ready_to_upload", "uploaded", "ready", "deleted", "rejected", "error", "uploading"]);

export const asset = z.object({asset_type: assetType, created_at: z.string(), data: z.discriminatedUnion("asset_type", [fileData, checkpointData, loraData, textualInversionData, vAEData, latentData]), description: z.string(), file_structure: fileStructure.optional(), hf_repo: z.string().nullable().optional(), id: z.string(), is_public: z.boolean().optional(), name: z.string(), size_bytes: z.number().int(), status: status, status_details: z.string(), status_info: z.object({}).optional(), tenant_uuid: z.string(), url: z.string().min(1).max(2083).nullable().optional()});

export const listAssetsResponse = z.object({data: z.array(asset), has_more: z.boolean(), total: z.number().int()});

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

/**
 * List Assets
 * 
 * GET /v1/assets
 */
export const getApiV1Assets = async (args: GetApiV1AssetsArgs, options?: GetApiV1AssetsOptions) => {
const {baseUrl = "https://api.octoai.cloud", fetch: fetchFn = fetch} = options ?? {}

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

if (args.name != null) query.append("name", String(args.name))
if (args.is_public != null) query.append("is_public", String(args.is_public))
if (args.data_type != null) query.append("data_type", String(args.data_type))
for (const item of args.asset_type ?? []) query.append("asset_type", String(item))
for (const item of args.engine ?? []) query.append("engine", String(item))
if (args.engine_type != null) query.append("engine_type", String(args.engine_type))
if (args.limit != null) query.append("limit", String(args.limit))
if (args.offset != null) query.append("offset", String(args.offset))
if (args.get_preview_urls != null) query.append("get_preview_urls", String(args.get_preview_urls))
if (args.asset_ids != null) query.append("asset_ids", String(args.asset_ids))
if (args.owner != null) query.append("owner", String(args.owner))
for (const item of args.file_format ?? []) query.append("file_format", String(item))

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

const res = await fetchFn(url, {method: "GET", headers})

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

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