import PropTypes from 'prop-types'
import { Card, Tree, Spin, Space, Button, message } from 'antd'
import { LoadingOutlined } from '@ant-design/icons'
import React, { useState, useEffect, useContext } from 'react'
import { useNavigate } from 'react-router-dom'
import SubJobState from './SubJobState'
import { SocketContext } from '../context/SocketContext'
import axios from 'axios'
import JobState from './JobState'
import JobIcon from '../../Icons/JobIcon'
import SubJobIcon from '../../Icons/SubJobIcon'
import ReloadIcon from '../../Icons/ReloadIcon'

import config from '../../../config'

const PrinterJobsTree = ({
  subJobs: initialSubJobs,
  loading: initialLoading
}) => {
  const [subJobs, setSubJobs] = useState(initialSubJobs || [])
  const [treeLoading, setTreeLoading] = useState(initialLoading)
  const [error, setError] = useState(null)
  const { socket } = useContext(SocketContext)
  const [messageApi] = message.useMessage()
  const [expandedKeys, setExpandedKeys] = useState([])
  const [treeData, setTreeData] = useState([])
  const navigate = useNavigate()

  const handleNodeClick = (selectedKeys) => {
    const key = selectedKeys[0]
    if (key.startsWith('job-')) {
      const jobId = key.replace('job-', '')
      navigate(`/dashboard/production/jobs/info?jobId=${jobId}`)
    }
  }

  const buildTreeData = (subJobsData) => {
    if (!subJobsData?.length) {
      setTreeData([])
      setExpandedKeys([])
      return
    }

    // Group subjobs by job
    const jobGroups = subJobsData.reduce((acc, subJob) => {
      const jobId = subJob.job._id
      if (!acc[jobId]) {
        acc[jobId] = {
          job: subJob.job,
          subJobs: []
        }
      }
      acc[jobId].subJobs.push(subJob)
      return acc
    }, {})

    // Create tree nodes for each job
    const jobNodes = Object.values(jobGroups).map(({ job, subJobs }) => {
      setExpandedKeys((prev) => [...prev, `job-${job._id}`])
      return {
        title: (
          <Space size={5}>
            <JobIcon />
            {'Job'}
            <JobState job={job} />
          </Space>
        ),
        key: `job-${job._id}`,
        children: subJobs.map((subJob) => ({
          title: (
            <Space>
              <SubJobIcon />
              {'Sub Job'}
              <SubJobState
                subJob={subJob}
                showProgress={false}
                showControls={false}
              />
            </Space>
          ),
          key: `subjob-${subJob._id}`,
          isLeaf: true
        }))
      }
    })

    setTreeData(jobNodes)
  }

  useEffect(() => {
    buildTreeData(subJobs)
  }, [subJobs])

  useEffect(() => {
    const initializeData = async () => {
      if (!initialSubJobs) {
        try {
          setTreeLoading(true)
          const response = await axios.get(`${config.backendUrl}/jobs`, {
            headers: { Accept: 'application/json' },
            withCredentials: true
          })
          if (response.data?.subJobs) {
            setSubJobs(response.data.subJobs)
          }
        } catch (err) {
          setError('Failed to fetch sub jobs')
          messageApi.error('Failed to fetch sub jobs')
        } finally {
          setTreeLoading(false)
        }
      } else {
        setSubJobs(initialSubJobs)
      }
    }

    initializeData()

    // Add socket.io event listener for subjob updates
    if (socket) {
      socket.on('notify_subjob_update', (updateData) => {
        if (updateData.subJobId) {
          setSubJobs((prevSubJobs) =>
            prevSubJobs.map((subJob) => {
              if (subJob._id === updateData._id) {
                return {
                  ...subJob,
                  state: updateData.state,
                  subJobId: updateData.subJobId
                }
              }
              return subJob
            })
          )
        }
      })
    }

    return () => {
      if (socket) {
        socket.off('notify_subjob_update')
      }
    }
  }, [initialSubJobs, socket])

  if (error) {
    return (
      <Space
        direction='vertical'
        style={{ width: '100%', textAlign: 'center' }}
      >
        <p>{error}</p>
        <Button icon={<ReloadIcon />} onClick={() => setError(null)}>
          Retry
        </Button>
      </Space>
    )
  }

  return (
    <Spin indicator={<LoadingOutlined />} spinning={treeLoading}>
      <Card>
        <Tree
          treeData={treeData}
          expandedKeys={expandedKeys}
          onExpand={setExpandedKeys}
          onSelect={handleNodeClick}
          showLine={true}
        />
      </Card>
    </Spin>
  )
}

PrinterJobsTree.propTypes = {
  subJobs: PropTypes.arrayOf(
    PropTypes.shape({
      state: PropTypes.object.isRequired,
      _id: PropTypes.string.isRequired,
      printer: PropTypes.string.isRequired,
      job: PropTypes.shape({
        state: PropTypes.object.isRequired,
        _id: PropTypes.string.isRequired,
        printers: PropTypes.arrayOf(PropTypes.string).isRequired,
        createdAt: PropTypes.string.isRequired,
        updatedAt: PropTypes.string.isRequired,
        startedAt: PropTypes.string.isRequired,
        gcodeFile: PropTypes.string.isRequired,
        quantity: PropTypes.number.isRequired,
        subJobs: PropTypes.arrayOf(PropTypes.string).isRequired
      }).isRequired,
      subJobId: PropTypes.string.isRequired,
      number: PropTypes.number.isRequired,
      createdAt: PropTypes.string.isRequired,
      updatedAt: PropTypes.string.isRequired
    })
  ),
  loading: PropTypes.bool
}

export default PrinterJobsTree