import { useEffect, useState, useContext, useCallback, useMemo } from 'react'
import PropTypes from 'prop-types'
import { TreeSelect, Space, Button, Input } from 'antd'
import ReloadIcon from '../../Icons/ReloadIcon'
import { ApiServerContext } from '../context/ApiServerContext'
import { AuthContext } from '../context/AuthContext'
import ObjectProperty from './ObjectProperty'
import { getModelByName } from '../../../database/ObjectModels'
import merge from 'lodash/merge'
const { SHOW_CHILD } = TreeSelect

const ObjectSelect = ({
  type = 'unknown',
  showSearch = false,
  multiple = false,
  treeSelectProps = {},
  filter = {},
  masterFilter = {},
  value,
  disabled = false,
  ...rest
}) => {
  const { fetchObjectsByProperty } = useContext(ApiServerContext)
  const { token } = useContext(AuthContext)
  // --- State ---
  const [treeData, setTreeData] = useState([])
  const [objectPropertiesTree, setObjectPropertiesTree] = useState({})
  const [initialized, setInitialized] = useState(false)
  const [error, setError] = useState(false)
  const properties = useMemo(() => getModelByName(type).group || [], [type])
  const [objectList, setObjectList] = useState([])
  const [treeSelectValue, setTreeSelectValue] = useState(null)
  const [initialLoading, setInitialLoading] = useState(true)

  // Fetch the object properties tree from the API
  const handleFetchObjectsProperties = useCallback(
    async (customFilter = filter) => {
      try {
        const data = await fetchObjectsByProperty(type, {
          properties: properties,
          filter: customFilter,
          masterFilter
        })
        if (Array.isArray(data)) {
          setObjectPropertiesTree((prev) => merge([], prev, data))
        } else {
          setObjectPropertiesTree((prev) => merge({}, prev, data))
        }
        setInitialLoading(false)
        setError(false)
        return data
      } catch (err) {
        console.error(err)
        setError(true)
        return null
      }
    },
    [type, fetchObjectsByProperty, properties, filter, masterFilter]
  )

  // Convert the API response to AntD TreeSelect treeData
  const buildTreeData = useCallback(
    (data, pIdx = 0, parentKeys = [], filterPath = []) => {
      if (!data) return []
      if (Array.isArray(data)) {
        return data.map((object) => {
          setObjectList((prev) => {
            const filtered = prev.filter(
              (prevObject) => prevObject._id != object._id
            )
            return [...filtered, object]
          })
          return {
            title: (
              <ObjectProperty
                key={object._id}
                type='object'
                value={object}
                objectType={type}
                objectData={object}
                isEditing={false}
              />
            ),
            value: object._id,
            key: object._id,
            isLeaf: true,
            property: properties[pIdx - 1], // previous property
            parentKeys,
            filterPath
          }
        })
      }
      if (typeof data == 'object') {
        const property = properties[pIdx] || null
        return Object.entries(data)
          .map(([key, value]) => {
            if (property != null && typeof value === 'object') {
              const newFilterPath = filterPath.concat({ property, value: key })
              return {
                title: <ObjectProperty type={property} value={key} />,
                value: parentKeys.concat(key).join('-'),
                filterValue: key,
                key: parentKeys.concat(key).join('-'),
                property,
                parentKeys: parentKeys.concat(key || '-'),
                filterPath: newFilterPath,
                selectable: false,
                children: buildTreeData(
                  value,
                  pIdx + 1,
                  parentKeys.concat(key),
                  newFilterPath
                ),
                isLeaf: false
              }
            }
          })
          .filter(Boolean)
      }
    },
    [properties, type]
  )

  // --- loadData for async loading on expand ---
  const loadData = async (node) => {
    // node.property is the property name, node.value is the value
    if (!node.property) return
    // Build filter for this node by merging all parent property-value pairs
    const customFilter = { ...filter }
    if (Array.isArray(node.filterPath)) {
      node.filterPath.forEach(({ property, value }) => {
        customFilter[property] = value
      })
    }
    customFilter[node.property] = node.filterValue
    // Fetch children for this node
    const data = await handleFetchObjectsProperties(customFilter)
    if (!data) return
    // Extract only the children for the specific node that was expanded
    let nodeSpecificData = data
    if (typeof data === 'object' && !Array.isArray(data)) {
      // If the API returns an object with multiple keys, get only the data for this node
      nodeSpecificData = data[node.value] || {}
    }
    // Build new tree children only for this specific node
    const children = buildTreeData(
      nodeSpecificData,
      properties.indexOf(node.property) + 1,
      node.parentKeys || [],
      (node.filterPath || []).concat({
        property: node.property,
        value: node.filterValue
      })
    )
    // Update treeData with new children for this node only
    setTreeData((prevTreeData) => {
      // Helper to recursively update the correct node
      const updateNode = (nodes) =>
        nodes.map((n) => {
          if (n.key === node.key) {
            return { ...n, children, isLeaf: children.length === 0 }
          } else if (n.children) {
            return { ...n, children: updateNode(n.children) }
          }
          return n
        })
      return updateNode(prevTreeData)
    })
  }

  const onTreeSelectChange = (value) => {
    // value can be a string (single) or array (multiple)
    if (multiple) {
      // Multiple selection
      let selectedObjects = []
      if (Array.isArray(value)) {
        selectedObjects = value
          .map((id) => objectList.find((obj) => obj._id === id))
          .filter(Boolean)
      }
      setTreeSelectValue(value)
      if (rest.onChange) rest.onChange(selectedObjects)
    } else {
      // Single selection
      const selectedObject = objectList.find((obj) => obj._id === value)
      setTreeSelectValue(value)
      if (rest.onChange) rest.onChange(selectedObject)
    }
  }

  // Update treeData when objectPropertiesTree changes
  useEffect(() => {
    if (objectPropertiesTree && Object.keys(objectPropertiesTree).length > 0) {
      const newTreeData = buildTreeData(objectPropertiesTree)
      setTreeData((prev) => {
        if (JSON.stringify(prev) !== JSON.stringify(newTreeData)) {
          return newTreeData
        }
        return prev
      })
    }
  }, [objectPropertiesTree, properties, type, buildTreeData])

  useEffect(() => {
    if (value && typeof value === 'object' && value !== null && !initialized) {
      // Build a new filter from value's properties that are in the properties list
      const valueFilter = { ...filter }
      properties.forEach((prop) => {
        if (Object.prototype.hasOwnProperty.call(value, prop)) {
          const filterValue = value[prop]
          if (filterValue?.name) {
            valueFilter[prop] = filterValue.name
          } else if (Array.isArray(filterValue)) {
            valueFilter[prop] = filterValue.join(',')
          } else {
            valueFilter[prop] = filterValue
          }
        }
      })
      // Fetch with the new filter
      handleFetchObjectsProperties(valueFilter)
      setTreeSelectValue(value._id)
      setInitialized(true)
      return
    }
    if (!initialized && token != null) {
      handleFetchObjectsProperties()
      setInitialized(true)
    }
  }, [
    value,
    filter,
    properties,
    handleFetchObjectsProperties,
    initialized,
    token
  ])

  // --- Error UI ---
  if (error) {
    return (
      <Space.Compact style={{ width: '100%' }}>
        <Input value='Failed to load data.' status='error' disabled />
        <Button
          icon={<ReloadIcon />}
          onClick={() => {
            setError(false)
            setTreeData([])
            setInitialized(false)
          }}
          danger
        />
      </Space.Compact>
    )
  }

  if (initialLoading) {
    return <TreeSelect disabled loading placeholder='Loading...' />
  }

  // --- Main TreeSelect UI ---
  return (
    <TreeSelect
      treeDataSimpleMode={false}
      treeDefaultExpandAll={true}
      treeData={treeData}
      showSearch={showSearch}
      multiple={multiple}
      loadData={loadData}
      showCheckedStrategy={SHOW_CHILD}
      placeholder={`Select a ${getModelByName(type).label.toLowerCase()}...`}
      {...treeSelectProps}
      {...rest}
      value={treeSelectValue}
      onChange={onTreeSelectChange}
      disabled={disabled}
    />
  )
}

ObjectSelect.propTypes = {
  properties: PropTypes.arrayOf(PropTypes.string).isRequired,
  filter: PropTypes.object,
  masterFilter: PropTypes.object,
  useFilter: PropTypes.bool,
  value: PropTypes.any,
  onChange: PropTypes.func,
  showSearch: PropTypes.bool,
  multiple: PropTypes.bool,
  treeSelectProps: PropTypes.object,
  type: PropTypes.string.isRequired,
  disabled: PropTypes.bool
}

export default ObjectSelect