import { Box, Button, Collapse, Typography } from '@mui/material'
import { ArrowRight, ChevronRight, ErrorOutline } from '@mui/icons-material'
import type { ChatToolGroupProps } from '../types'
import { ChatToolTrace, ChatToolTraceDetails } from './chat-tool-trace'
import { styles } from './styles'
import { getToolLabel } from './get-tool-label'
import { useState } from 'react'

function useController<T>(value: T, setter?: (value: T) => void) {
  const [_value, _setter] = useState<T>(value)
  if (setter) {
    return [value, setter] as const
  } else {
    return [_value, _setter] as const
  }
}

export function ChatToolGroup({
  tools,
  expanded: _expanded = false,
  onExpandedChange: _setExpanded,
  expandedTools: _expandedTools = {},
  onToolExpandedChange: _setExpandedTools,
  labels = {},
  sx,
}: ChatToolGroupProps) {
  const [expanded, setExpanded] = useController(_expanded, _setExpanded)
  const [expandedTools, setExpandedTools] = useController(
    _expandedTools,
    _setExpandedTools,
  )

  const numErrors = tools.filter((t) => t.status === 'error').length
  const runningTools = tools.filter((t) => t.status === 'running')
  const nonRunningTools = tools.filter((t) => t.status !== 'running')
  const toolsUsedLabel = labels.toolsUsed ?? 'tools used'
  const errorLabel = labels.error ?? 'Error'

  if (tools.length === 1) {
    return (
      <ChatToolTrace
        tool={tools[0]!}
        expanded={expanded}
        onExpandedChange={setExpanded}
        labels={labels}
        sx={sx}
      />
    )
  }

  return (
    <Box width='100%' sx={sx} className='PsChat--tool-group'>
      <Button
        size='small'
        variant='text'
        onClick={() => setExpanded(!expanded)}
        aria-expanded={expanded}
        sx={{
          ...styles.traceHeader,
          pr: numErrors > 0 ? 0.5 : 0,
        }}
      >
        <Typography variant='caption' fontWeight={600} color='text.secondary'>
          {tools.length} {toolsUsedLabel}
        </Typography>
        <ArrowRight
          sx={{
            ...styles.traceChevron,
            transform: expanded ? 'rotate(90deg)' : 'rotate(0deg)',
          }}
        />
        {numErrors > 0 && (
          <Typography variant='code3' sx={styles.errorBadge}>
            {numErrors}
            <ErrorOutline sx={{ fontSize: 12 }} color='error' />
          </Typography>
        )}
      </Button>
      <Collapse in={expanded} unmountOnExit>
        <Box
          component='ul'
          sx={{
            ...styles.traceDetailsWrapper,
            p: 0,
            listStyle: 'none',
          }}
        >
          {nonRunningTools.map((tool) => {
            const isExpanded = expandedTools[tool.id] ?? false
            return (
              <Box
                key={tool.id}
                component='li'
                sx={styles.groupListItem}
                aria-expanded={isExpanded}
              >
                <Button
                  variant='text'
                  color='inherit'
                  fullWidth
                  onClick={() =>
                    setExpandedTools({
                      ...expandedTools,
                      [tool.id]: !isExpanded,
                    })
                  }
                  aria-expanded={isExpanded}
                  sx={styles.groupHeader}
                >
                  <Typography
                    flexGrow={1}
                    flexShrink={0}
                    variant='caption'
                    color='text.secondary'
                  >
                    {getToolLabel(tool)}
                  </Typography>
                  {tool.status === 'error' && (
                    <Typography variant='code3' color='error.main'>
                      {errorLabel}
                    </Typography>
                  )}
                  <ChevronRight
                    sx={{
                      ...styles.traceChevron,
                      transform: isExpanded ? 'rotate(90deg)' : 'rotate(0deg)',
                    }}
                  />
                </Button>
                <Collapse in={isExpanded} unmountOnExit>
                  <Box px={1.5} py={1}>
                    <ChatToolTraceDetails tool={tool} labels={labels} />
                  </Box>
                </Collapse>
              </Box>
            )
          })}
        </Box>
      </Collapse>
      {runningTools.map((tool) => (
        <ChatToolTrace
          key={tool.id}
          tool={tool}
          labels={labels}
          expanded={expandedTools[tool.id]}
          onExpandedChange={() =>
            setExpandedTools({
              ...expandedTools,
              [tool.id]: !expandedTools[tool.id],
            })
          }
        />
      ))}
    </Box>
  )
}