ComfyUI_frontend/src/platform/assets/services/assetService.ts

import { fromZodError } from 'zod-validation-error'

import { st } from '@/i18n'
import { assetResponseSchema } from '@/platform/assets/schemas/assetSchema'
import type {
  AssetItem,
  AssetMetadata,
  AssetResponse,
  ModelFile,
  ModelFolder
} from '@/platform/assets/schemas/assetSchema'
import { api } from '@/scripts/api'
import { useModelToNodeStore } from '@/stores/modelToNodeStore'

/**
 * Maps CivitAI validation error codes to localized error messages
 */
function getLocalizedErrorMessage(errorCode: string): string {
  const errorMessages: Record<string, string> = {
    FILE_TOO_LARGE: st('assetBrowser.errorFileTooLarge', 'File too large'),
    FORMAT_NOT_ALLOWED: st(
      'assetBrowser.errorFormatNotAllowed',
      'Format not allowed'
    ),
    UNSAFE_PICKLE_SCAN: st(
      'assetBrowser.errorUnsafePickleScan',
      'Unsafe pickle scan'
    ),
    UNSAFE_VIRUS_SCAN: st(
      'assetBrowser.errorUnsafeVirusScan',
      'Unsafe virus scan'
    ),
    MODEL_TYPE_NOT_SUPPORTED: st(
      'assetBrowser.errorModelTypeNotSupported',
      'Model type not supported'
    )
  }
  return (
    errorMessages[errorCode] ||
    st('assetBrowser.errorUnknown', 'Unknown error') ||
    'Unknown error'
  )
}

const ASSETS_ENDPOINT = '/assets'
const EXPERIMENTAL_WARNING = `EXPERIMENTAL: If you are seeing this please make sure "Comfy.Assets.UseAssetAPI" is set to "false" in your ComfyUI Settings.\n`
const DEFAULT_LIMIT = 500

export const MODELS_TAG = 'models'
export const MISSING_TAG = 'missing'

/**
 * Validates asset response data using Zod schema
 */
function validateAssetResponse(data: unknown): AssetResponse {
  const result = assetResponseSchema.safeParse(data)
  if (result.success) return result.data

  const error = fromZodError(result.error)
  throw new Error(
    `${EXPERIMENTAL_WARNING}Invalid asset response against zod schema:\n${error}`
  )
}

/**
 * Private service for asset-related network requests
 * Not exposed globally - used internally by ComfyApi
 */
function createAssetService() {
  /**
   * Handles API response with consistent error handling and Zod validation
   */
  async function handleAssetRequest(
    url: string,
    context: string
  ): Promise<AssetResponse> {
    const res = await api.fetchApi(url)
    if (!res.ok) {
      throw new Error(
        `${EXPERIMENTAL_WARNING}Unable to load ${context}: Server returned ${res.status}. Please try again.`
      )
    }
    const data = await res.json()
    return validateAssetResponse(data)
  }
  /**
   * Gets a list of model folder keys from the asset API
   *
   * Logic:
   * 1. Extract directory names directly from asset tags
   * 2. Filter out blacklisted directories
   * 3. Return alphabetically sorted directories with assets
   *
   * @returns The list of model folder keys
   */
  async function getAssetModelFolders(): Promise<ModelFolder[]> {
    const data = await handleAssetRequest(
      `${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG}&limit=${DEFAULT_LIMIT}`,
      'model folders'
    )

    // Blacklist directories we don't want to show
    const blacklistedDirectories = new Set(['configs'])

    // Extract directory names from assets that actually exist, exclude missing assets
    const discoveredFolders = new Set<string>(
      data?.assets
        ?.filter((asset) => !asset.tags.includes(MISSING_TAG))
        ?.flatMap((asset) => asset.tags)
        ?.filter(
          (tag) => tag !== MODELS_TAG && !blacklistedDirectories.has(tag)
        ) ?? []
    )

    // Return only discovered folders in alphabetical order
    const sortedFolders = Array.from(discoveredFolders).toSorted()
    return sortedFolders.map((name) => ({ name, folders: [] }))
  }

  /**
   * Gets a list of models in the specified folder from the asset API
   * @param folder The folder to list models from, such as 'checkpoints'
   * @returns The list of model filenames within the specified folder
   */
  async function getAssetModels(folder: string): Promise<ModelFile[]> {
    const data = await handleAssetRequest(
      `${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG},${folder}&limit=${DEFAULT_LIMIT}`,
      `models for ${folder}`
    )

    return (
      data?.assets
        ?.filter(
          (asset) =>
            !asset.tags.includes(MISSING_TAG) && asset.tags.includes(folder)
        )
        ?.map((asset) => ({
          name: asset.name,
          pathIndex: 0
        })) ?? []
    )
  }

  /**
   * Checks if a widget input should use the asset browser based on both input name and node comfyClass
   *
   * @param nodeType - The ComfyUI node comfyClass (e.g., 'CheckpointLoaderSimple', 'LoraLoader')
   * @param widgetName - The name of the widget to check (e.g., 'ckpt_name')
   * @returns true if this input should use asset browser
   */
  function isAssetBrowserEligible(
    nodeType: string | undefined,
    widgetName: string
  ): boolean {
    if (!nodeType || !widgetName) return false
    return (
      useModelToNodeStore().getRegisteredNodeTypes()[nodeType] === widgetName
    )
  }

  /**
   * Gets assets for a specific node type by finding the matching category
   * and fetching all assets with that category tag
   *
   * @param nodeType - The ComfyUI node type (e.g., 'CheckpointLoaderSimple')
   * @returns Promise<AssetItem[]> - Full asset objects with preserved metadata
   */
  async function getAssetsForNodeType(nodeType: string): Promise<AssetItem[]> {
    if (!nodeType || typeof nodeType !== 'string') {
      return []
    }

    // Find the category for this node type using efficient O(1) lookup
    const modelToNodeStore = useModelToNodeStore()
    const category = modelToNodeStore.getCategoryForNodeType(nodeType)

    if (!category) {
      return []
    }

    // Fetch assets for this category using same API pattern as getAssetModels
    const data = await handleAssetRequest(
      `${ASSETS_ENDPOINT}?include_tags=${MODELS_TAG},${category}&limit=${DEFAULT_LIMIT}`,
      `assets for ${nodeType}`
    )

    // Return full AssetItem[] objects (don't strip like getAssetModels does)
    return (
      data?.assets?.filter(
        (asset) =>
          !asset.tags.includes(MISSING_TAG) && asset.tags.includes(category)
      ) ?? []
    )
  }

  /**
   * Gets complete details for a specific asset by ID
   * Calls the detail endpoint which includes user_metadata and all fields
   *
   * @param id - The asset ID
   * @returns Promise<AssetItem> - Complete asset object with user_metadata
   */
  async function getAssetDetails(id: string): Promise<AssetItem> {
    const res = await api.fetchApi(`${ASSETS_ENDPOINT}/${id}`)
    if (!res.ok) {
      throw new Error(
        `${EXPERIMENTAL_WARNING}Unable to load asset details for ${id}: Server returned ${res.status}. Please try again.`
      )
    }
    const data = await res.json()

    // Validate the single asset response against our schema
    const result = assetResponseSchema.safeParse({ assets: [data] })
    if (result.success && result.data.assets?.[0]) {
      return result.data.assets[0]
    }

    const error = result.error
      ? fromZodError(result.error)
      : 'Unknown validation error'
    throw new Error(
      `${EXPERIMENTAL_WARNING}Invalid asset response against zod schema:\n${error}`
    )
  }

  /**
   * Gets assets filtered by a specific tag
   *
   * @param tag - The tag to filter by (e.g., 'models', 'input')
   * @param includePublic - Whether to include public assets (default: true)
   * @param options - Pagination options
   * @param options.limit - Maximum number of assets to return (default: 500)
   * @param options.offset - Number of assets to skip (default: 0)
   * @returns Promise<AssetItem[]> - Full asset objects filtered by tag, excluding missing assets
   */
  async function getAssetsByTag(
    tag: string,
    includePublic: boolean = true,
    {
      limit = DEFAULT_LIMIT,
      offset = 0
    }: { limit?: number; offset?: number } = {}
  ): Promise<AssetItem[]> {
    const queryParams = new URLSearchParams({
      include_tags: tag,
      limit: limit.toString(),
      include_public: includePublic ? 'true' : 'false'
    })

    if (offset > 0) {
      queryParams.set('offset', offset.toString())
    }

    const data = await handleAssetRequest(
      `${ASSETS_ENDPOINT}?${queryParams.toString()}`,
      `assets for tag ${tag}`
    )

    return (
      data?.assets?.filter((asset) => !asset.tags.includes(MISSING_TAG)) ?? []
    )
  }

  /**
   * Deletes an asset by ID
   * Only available in cloud environment
   *
   * @param id - The asset ID (UUID)
   * @returns Promise<void>
   * @throws Error if deletion fails
   */
  async function deleteAsset(id: string): Promise<void> {
    const res = await api.fetchApi(`${ASSETS_ENDPOINT}/${id}`, {
      method: 'DELETE'
    })

    if (!res.ok) {
      throw new Error(
        `Unable to delete asset ${id}: Server returned ${res.status}`
      )
    }
  }

  /**
   * Retrieves metadata from a download URL without downloading the file
   *
   * @param url - Download URL to retrieve metadata from (will be URL-encoded)
   * @returns Promise with metadata including content_length, final_url, filename, etc.
   * @throws Error if metadata retrieval fails
   */
  async function getAssetMetadata(url: string): Promise<AssetMetadata> {
    const encodedUrl = encodeURIComponent(url)
    const res = await api.fetchApi(
      `${ASSETS_ENDPOINT}/remote-metadata?url=${encodedUrl}`
    )

    if (!res.ok) {
      const errorData = await res.json().catch(() => ({}))
      throw new Error(
        getLocalizedErrorMessage(errorData.code || 'UNKNOWN_ERROR')
      )
    }

    const data: AssetMetadata = await res.json()
    if (data.validation?.is_valid === false) {
      throw new Error(
        getLocalizedErrorMessage(
          data.validation?.errors?.[0]?.code || 'UNKNOWN_ERROR'
        )
      )
    }

    return data
  }

  /**
   * Uploads an asset by providing a URL to download from
   *
   * @param params - Upload parameters
   * @param params.url - HTTP/HTTPS URL to download from
   * @param params.name - Display name (determines extension)
   * @param params.tags - Optional freeform tags
   * @param params.user_metadata - Optional custom metadata object
   * @param params.preview_id - Optional UUID for preview asset
   * @returns Promise<AssetItem & { created_new: boolean }> - Asset object with created_new flag
   * @throws Error if upload fails
   */
  async function uploadAssetFromUrl(params: {
    url: string
    name: string
    tags?: string[]
    user_metadata?: Record<string, any>
    preview_id?: string
  }): Promise<AssetItem & { created_new: boolean }> {
    const res = await api.fetchApi(ASSETS_ENDPOINT, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json'
      },
      body: JSON.stringify(params)
    })

    if (!res.ok) {
      throw new Error(
        st(
          'assetBrowser.errorUploadFailed',
          'Failed to upload asset. Please try again.'
        )
      )
    }

    return await res.json()
  }

  return {
    getAssetModelFolders,
    getAssetModels,
    isAssetBrowserEligible,
    getAssetsForNodeType,
    getAssetDetails,
    getAssetsByTag,
    deleteAsset,
    getAssetMetadata,
    uploadAssetFromUrl
  }
}

export const assetService = createAssetService()