'use client'
import type { FC } from 'react'
import type {
  HandleAddCondition,
  HandleRemoveCondition,
  HandleToggleConditionLogicalOperator,
  HandleUpdateCondition,
  MetadataFilteringModeEnum,
} from '@/app/components/workflow/nodes/knowledge-retrieval/types'
import type { DataSet } from '@/models/datasets'
import { cn } from '@langgenius/dify-ui/cn'
import { intersectionBy } from 'es-toolkit/compat'
import { produce } from 'immer'
import * as React from 'react'
import { useCallback, useMemo } from 'react'
import { useTranslation } from 'react-i18next'
import { useContext } from 'use-context-selector'
import { v4 as uuid4 } from 'uuid'
import { ModelTypeEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
import { useModelListAndDefaultModelAndCurrentProviderAndModel } from '@/app/components/header/account-setting/model-provider-page/hooks'
import MetadataFilter from '@/app/components/workflow/nodes/knowledge-retrieval/components/metadata/metadata-filter'
import {
  ComparisonOperator,
  LogicalOperator,
  MetadataFilteringVariableType,
} from '@/app/components/workflow/nodes/knowledge-retrieval/types'
import {
  getMultipleRetrievalConfig,
  getSelectedDatasetsMode,
} from '@/app/components/workflow/nodes/knowledge-retrieval/utils'
import { useSelector as useAppContextSelector } from '@/context/app-context'
import ConfigContext from '@/context/debug-configuration'
import { AppModeEnum } from '@/types/app'
import { hasEditPermissionForDataset } from '@/utils/permission'
import FeaturePanel from '../base/feature-panel'
import OperationBtn from '../base/operation-btn'
import { useFormattingChangedDispatcher } from '../debug/hooks'
import CardItem from './card-item'
import ContextVar from './context-var'
import ParamsConfig from './params-config'

type Props = {
  readonly?: boolean
  hideMetadataFilter?: boolean
}
const DatasetConfig: FC<Props> = ({ readonly, hideMetadataFilter }) => {
  const { t } = useTranslation()
  const userProfile = useAppContextSelector(s => s.userProfile)
  const {
    mode,
    dataSets: dataSet,
    setDataSets: setDataSet,
    modelConfig,
    setModelConfig,
    showSelectDataSet,
    isAgent,
    datasetConfigs,
    datasetConfigsRef,
    setDatasetConfigs,
    setRerankSettingModalOpen,
  } = useContext(ConfigContext)
  const formattingChangedDispatcher = useFormattingChangedDispatcher()

  const hasData = dataSet.length > 0

  const {
    currentModel: currentRerankModel,
    currentProvider: currentRerankProvider,
  } = useModelListAndDefaultModelAndCurrentProviderAndModel(ModelTypeEnum.rerank)

  const onRemove = (id: string) => {
    const filteredDataSets = dataSet.filter(item => item.id !== id)
    setDataSet(filteredDataSets)
    const { datasets, retrieval_model, score_threshold_enabled, ...restConfigs } = datasetConfigs
    const {
      top_k,
      score_threshold,
      reranking_model,
      reranking_mode,
      weights,
      reranking_enable,
    } = restConfigs
    const oldRetrievalConfig = {
      top_k,
      score_threshold,
      reranking_model: (reranking_model && reranking_model.reranking_provider_name && reranking_model.reranking_model_name)
        ? {
            provider: reranking_model.reranking_provider_name,
            model: reranking_model.reranking_model_name,
          }
        : undefined,
      reranking_mode,
      weights,
      reranking_enable,
    }
    const retrievalConfig = getMultipleRetrievalConfig(oldRetrievalConfig, filteredDataSets, dataSet, {
      provider: currentRerankProvider?.provider,
      model: currentRerankModel?.model,
    })
    setDatasetConfigs({
      ...datasetConfigsRef.current,
      ...retrievalConfig,
      reranking_model: {
        reranking_provider_name: retrievalConfig?.reranking_model?.provider || '',
        reranking_model_name: retrievalConfig?.reranking_model?.model || '',
      },
      retrieval_model,
      score_threshold_enabled,
      datasets,
    })
    const {
      allExternal,
      allInternal,
      mixtureInternalAndExternal,
      mixtureHighQualityAndEconomic,
      inconsistentEmbeddingModel,
    } = getSelectedDatasetsMode(filteredDataSets)

    if (
      (allInternal && (mixtureHighQualityAndEconomic || inconsistentEmbeddingModel))
      || mixtureInternalAndExternal
      || allExternal
    ) {
      setRerankSettingModalOpen(true)
    }
    formattingChangedDispatcher()
  }

  const handleSave = (newDataset: DataSet) => {
    const index = dataSet.findIndex(item => item.id === newDataset.id)

    const newDatasets = [...dataSet.slice(0, index), newDataset, ...dataSet.slice(index + 1)]
    setDataSet(newDatasets)
    formattingChangedDispatcher()
  }

  const promptVariables = modelConfig.configs.prompt_variables
  const promptVariablesToSelect = promptVariables.map(item => ({
    name: item.name,
    type: item.type,
    value: item.key,
  }))
  const selectedContextVar = promptVariables?.find(item => item.is_context_var)
  const handleSelectContextVar = (selectedValue: string) => {
    const newModelConfig = produce(modelConfig, (draft) => {
      draft.configs.prompt_variables = modelConfig.configs.prompt_variables.map((item) => {
        return ({
          ...item,
          is_context_var: item.key === selectedValue,
        })
      })
    })
    setModelConfig(newModelConfig)
  }

  const formattedDataset = useMemo(() => {
    return dataSet.map((item) => {
      const datasetConfig = {
        createdBy: item.created_by,
        partialMemberList: item.partial_member_list || [],
        permission: item.permission,
      }
      return {
        ...item,
        editable: hasEditPermissionForDataset(userProfile?.id || '', datasetConfig),
      }
    })
  }, [dataSet, userProfile?.id])

  const metadataList = useMemo(() => {
    return intersectionBy(...formattedDataset.filter((dataset) => {
      return !!dataset.doc_metadata
    }).map((dataset) => {
      return dataset.doc_metadata!
    }), 'name')
  }, [formattedDataset])

  const handleMetadataFilterModeChange = useCallback((newMode: MetadataFilteringModeEnum) => {
    setDatasetConfigs(produce(datasetConfigsRef.current!, (draft) => {
      draft.metadata_filtering_mode = newMode
    }))
  }, [setDatasetConfigs, datasetConfigsRef])

  const handleAddCondition = useCallback<HandleAddCondition>(({ id, name, type }) => {
    let operator: ComparisonOperator = ComparisonOperator.is

    if (type === MetadataFilteringVariableType.number)
      operator = ComparisonOperator.equal

    const newCondition = {
      id: uuid4(),
      metadata_id: id, // Save metadata.id for reliable reference
      name,
      comparison_operator: operator,
    }

    const newInputs = produce(datasetConfigsRef.current!, (draft) => {
      if (draft.metadata_filtering_conditions) {
        draft.metadata_filtering_conditions.conditions.push(newCondition)
      }
      else {
        draft.metadata_filtering_conditions = {
          logical_operator: LogicalOperator.and,
          conditions: [newCondition],
        }
      }
    })
    setDatasetConfigs(newInputs)
  }, [setDatasetConfigs, datasetConfigsRef])

  const handleRemoveCondition = useCallback<HandleRemoveCondition>((id) => {
    const conditions = datasetConfigsRef.current!.metadata_filtering_conditions?.conditions || []
    const index = conditions.findIndex(c => c.id === id)
    const newInputs = produce(datasetConfigsRef.current!, (draft) => {
      if (index > -1)
        draft.metadata_filtering_conditions?.conditions.splice(index, 1)
    })
    setDatasetConfigs(newInputs)
  }, [setDatasetConfigs, datasetConfigsRef])

  const handleUpdateCondition = useCallback<HandleUpdateCondition>((id, newCondition) => {
    const conditions = datasetConfigsRef.current!.metadata_filtering_conditions?.conditions || []
    const index = conditions.findIndex(c => c.id === id)
    const newInputs = produce(datasetConfigsRef.current!, (draft) => {
      if (index > -1)
        draft.metadata_filtering_conditions!.conditions[index] = newCondition
    })
    setDatasetConfigs(newInputs)
  }, [setDatasetConfigs, datasetConfigsRef])

  const handleToggleConditionLogicalOperator = useCallback<HandleToggleConditionLogicalOperator>(() => {
    const oldLogicalOperator = datasetConfigsRef.current!.metadata_filtering_conditions?.logical_operator
    const newLogicalOperator = oldLogicalOperator === LogicalOperator.and ? LogicalOperator.or : LogicalOperator.and
    const newInputs = produce(datasetConfigsRef.current!, (draft) => {
      draft.metadata_filtering_conditions!.logical_operator = newLogicalOperator
    })
    setDatasetConfigs(newInputs)
  }, [setDatasetConfigs, datasetConfigsRef])

  const handleMetadataModelChange = useCallback((model: { provider: string, modelId: string, mode?: string }) => {
    const newInputs = produce(datasetConfigsRef.current!, (draft) => {
      draft.metadata_model_config = {
        provider: model.provider,
        name: model.modelId,
        mode: model.mode || AppModeEnum.CHAT,
        completion_params: draft.metadata_model_config?.completion_params || { temperature: 0.7 },
      }
    })
    setDatasetConfigs(newInputs)
  }, [setDatasetConfigs, datasetConfigsRef])

  const handleMetadataCompletionParamsChange = useCallback((newParams: Record<string, any>) => {
    const newInputs = produce(datasetConfigsRef.current!, (draft) => {
      draft.metadata_model_config = {
        ...draft.metadata_model_config!,
        completion_params: newParams,
      }
    })
    setDatasetConfigs(newInputs)
  }, [setDatasetConfigs, datasetConfigsRef])

  return (
    <FeaturePanel
      className="mt-2"
      title={t('feature.dataSet.title', { ns: 'appDebug' })}
      headerRight={(
        !readonly && (
          <div className="flex items-center gap-1">
            {!isAgent && <ParamsConfig disabled={!hasData} selectedDatasets={dataSet} />}
            <OperationBtn type="add" onClick={showSelectDataSet} />
          </div>
        )
      )}
      hasHeaderBottomBorder={!hasData}
      noBodySpacing
    >
      {hasData
        ? (
            <div className={cn('mt-1 grid grid-cols-1 px-3 pb-3', readonly && 'grid-cols-2 gap-1')}>
              {formattedDataset.map(item => (
                <CardItem
                  key={item.id}
                  config={item}
                  onRemove={onRemove}
                  onSave={handleSave}
                  editable={item.editable}
                  readonly={readonly}
                />
              ))}
            </div>
          )
        : (
            <div className="mt-1 px-3 pb-3">
              <div className="pt-2 pb-1 text-xs text-text-tertiary">{t('feature.dataSet.noData', { ns: 'appDebug' })}</div>
            </div>
          )}

      {!hideMetadataFilter && (
        <div className="border-t border-t-divider-subtle py-2">
          <MetadataFilter
            metadataList={metadataList}
            selectedDatasetsLoaded
            metadataFilterMode={datasetConfigs.metadata_filtering_mode}
            metadataFilteringConditions={datasetConfigs.metadata_filtering_conditions}
            handleAddCondition={handleAddCondition}
            handleMetadataFilterModeChange={handleMetadataFilterModeChange}
            handleRemoveCondition={handleRemoveCondition}
            handleToggleConditionLogicalOperator={handleToggleConditionLogicalOperator}
            handleUpdateCondition={handleUpdateCondition}
            metadataModelConfig={datasetConfigs.metadata_model_config}
            handleMetadataModelChange={handleMetadataModelChange}
            handleMetadataCompletionParamsChange={handleMetadataCompletionParamsChange}
            isCommonVariable
            availableCommonStringVars={promptVariablesToSelect.filter(item => item.type === MetadataFilteringVariableType.string || item.type === MetadataFilteringVariableType.select)}
            availableCommonNumberVars={promptVariablesToSelect.filter(item => item.type === MetadataFilteringVariableType.number)}
          />
        </div>
      )}

      {!readonly && mode === AppModeEnum.COMPLETION && dataSet.length > 0 && (
        <ContextVar
          value={selectedContextVar?.key}
          options={promptVariablesToSelect}
          onChange={handleSelectContextVar}
        />
      )}
    </FeaturePanel>
  )
}
export default React.memo(DatasetConfig)
