import {z} from 'zod'

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

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

export type FileData = {asset_type?: 'file' | undefined, file_format?: FileExtension | undefined, preview_url?: (string | null) | undefined, version?: string | undefined};

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

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 FileFormat = 'safetensors' | 'pt';

export type CheckpointData = {asset_type?: 'checkpoint' | undefined, data_type?: DataType | undefined, engine: BaseEngine, file_format?: FileFormat | undefined};

export type LoraData = {asset_type?: 'lora' | undefined, checkpoint?: (string | null) | undefined, data_type?: DataType | undefined, engine: BaseEngine, file_format?: FileFormat | undefined, trigger_words?: (Array<string> | null) | undefined};

export type TextualInversionData = {asset_type?: 'textual_inversion' | undefined, data_type?: DataType | undefined, engine: BaseEngine, file_format?: FileFormat | undefined, trigger_words?: Array<string> | undefined};

export type VAEData = {asset_type?: 'vae' | undefined, data_type?: DataType | undefined, engine: BaseEngine, file_format?: FileFormat | undefined};

export type LatentData = {asset_type?: 'latent' | undefined, engine: BaseEngine, file_format?: FileFormat | undefined};

export type TransferApiType = 'presigned-url' | 'sts' | 'gcp-oauth' | 'any';

export type CreateAssetRequest = {asset_type: AssetType, data: FileData | CheckpointData | LoraData | TextualInversionData | VAEData | LatentData, description: string | null, hf_repo?: (string | null) | undefined, hf_token_secret?: (string | null) | undefined, is_public?: boolean | undefined, name: string, skip_validation?: boolean | undefined, transfer_api_type?: TransferApiType | undefined, url?: (string | null) | undefined};

export type CreateApiV1AssetsArgs = {body: CreateAssetRequest};

export type CreateApiV1AssetsOptions = {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 presignedUrlTransferApi = z.object({get_url: z.string(), put_url: z.string(), type: z.literal("presigned-url").optional()});

export const stsTransferApi = z.object({aws_access_key_id: z.string(), aws_region: z.string(), aws_secret_access_key: z.string(), aws_session_token: z.string().nullable().optional(), s3_bucket: z.string(), s3_endpoint_url: z.string().nullable().optional(), s3_key: z.string(), s3_url: z.string(), type: z.literal("sts").optional()});

export const gcsOauthTransferApi = z.object({gcp_region: z.string(), gs_bucket: z.string(), gs_key: z.string(), gs_url: z.string(), token: z.string(), type: z.literal("gcp-oauth").optional()});

export const noTransferApi = z.object({type: z.literal("none").optional()});

export const createAssetResponse = z.object({asset: asset, token: z.string().optional(), transfer_api: z.discriminatedUnion("type", [presignedUrlTransferApi, stsTransferApi, gcsOauthTransferApi, noTransferApi])});

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

/**
 * Create Asset
 * 
 * POST /v1/assets
 */
export const createApiV1Assets = async (args: CreateApiV1AssetsArgs, options?: CreateApiV1AssetsOptions) => {
const {baseUrl = "https://api.octoai.cloud", fetch: fetchFn = fetch} = options ?? {}

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

const url = `${baseUrl.replace(/\/+$/, '')}/v1/assets${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(createAssetResponse, await res.json())
};