import {
  Box,
  CircularProgress,
  Typography,
  Collapse,
  Button,
} from '@mui/material'
import { ArrowRight, CheckCircle, Error } from '@mui/icons-material'
import type { ChatToolTraceProps } from '../types'
import { ChatToolCodeArea } from './chat-tool-code-area'
import { styles } from './styles'
import { ChatThinking } from './chat-thinking'
import { getToolLabel } from './get-tool-label'
import { McpTool } from '@carto/meridian-ds/custom-icons'

function TraceStatusIndicator({
  status,
  labels,
}: {
  status: string
  labels: ChatToolTraceProps['labels']
}) {
  switch (status) {
    case 'complete':
      return (
        <Box sx={{ display: 'flex', alignItems: 'center', gap: 0.5 }}>
          <CheckCircle sx={{ fontSize: 12 }} color='success' />
          <Typography variant='caption'>
            {labels?.success ?? 'Success'}
          </Typography>
        </Box>
      )
    case 'error':
      return (
        <Box sx={{ display: 'flex', alignItems: 'center', gap: 0.5 }}>
          <Error sx={{ fontSize: 12 }} color='error' />
          <Typography variant='code3'>{labels?.error ?? 'Error'}</Typography>
        </Box>
      )
    case 'running':
      return <CircularProgress size={14} />
    default:
      return null
  }
}

export function ChatToolTraceDetails({
  tool,
  labels = {},
}: Pick<ChatToolTraceProps, 'tool' | 'labels'>) {
  const isError = tool.status === 'error'
  const toolExecutedLabel = labels.toolExecuted ?? 'Tool executed'
  const referenceLabel = labels.reference ?? 'Reference:'
  const durationLabel = labels.duration ?? 'Duration:'
  const statusLabel = labels.status ?? 'Status:'
  const inputArgumentsLabel = labels.inputArguments ?? 'Input arguments:'
  const outputLabel = labels.output ?? 'Output:'

  return (
    <Box>
      {tool.reference && (
        <Box sx={styles.traceField}>
          <Typography
            variant='code3'
            fontWeight={600}
            color='text.secondary'
            sx={styles.traceFieldLabel}
          >
            {referenceLabel}
          </Typography>
          <Box sx={styles.traceReference}>
            <McpTool color='success' />
            <Typography variant='code3' fontWeight={600}>
              {tool.reference}
            </Typography>
          </Box>
        </Box>
      )}

      {tool.duration != null && (
        <Box sx={styles.traceField}>
          <Typography
            variant='code3'
            fontWeight={600}
            color='text.secondary'
            sx={styles.traceFieldLabel}
          >
            {durationLabel}
          </Typography>
          <Typography variant='code3'>{tool.duration}s</Typography>
        </Box>
      )}

      <Box sx={styles.traceField}>
        <Typography
          variant='code3'
          fontWeight={600}
          color='text.secondary'
          sx={styles.traceFieldLabel}
        >
          {statusLabel}
        </Typography>
        <TraceStatusIndicator status={tool.status} labels={labels} />
      </Box>

      {tool.inputArguments && (
        <Box>
          <Typography
            variant='code3'
            fontWeight={600}
            color='text.secondary'
            sx={styles.traceField}
          >
            {inputArgumentsLabel}
          </Typography>
          <ChatToolCodeArea
            sx={{ mt: 0.5, mb: tool.output ? 1.5 : 0 }}
            content={tool.inputArguments}
            title={`${toolExecutedLabel}: ${inputArgumentsLabel}`}
            isError={isError}
          />
        </Box>
      )}

      {tool.output && (
        <Box>
          <Typography
            variant='code3'
            fontWeight={600}
            color='text.secondary'
            sx={styles.traceField}
          >
            {outputLabel}
          </Typography>
          <ChatToolCodeArea
            sx={{ mt: 0.5 }}
            content={tool.output}
            title={`${toolExecutedLabel}: ${outputLabel}`}
            isError={isError}
          />
        </Box>
      )}
    </Box>
  )
}

export function ChatToolTrace({
  tool,
  expanded,
  onExpandedChange,
  labels = {},
  sx,
}: ChatToolTraceProps) {
  const isRunning = tool.status === 'running'
  const toolExecutedLabel = labels.toolExecuted ?? 'Tool executed'

  if (isRunning) {
    return <ChatThinking>{getToolLabel(tool)}</ChatThinking>
  }

  return (
    <Box
      width='100%'
      sx={sx}
      data-tool={tool.name}
      className={`PsChat--tool-trace`}
    >
      <Button
        size='small'
        variant='text'
        onClick={() => onExpandedChange?.(!expanded)}
        aria-expanded={expanded}
        sx={styles.traceHeader}
      >
        <Typography variant='caption' fontWeight={600} color='text.secondary'>
          {toolExecutedLabel}
        </Typography>
        <ArrowRight
          sx={{
            ...styles.traceChevron,
            transform: expanded ? 'rotate(90deg)' : 'rotate(0deg)',
          }}
        />
      </Button>
      <Collapse in={expanded} unmountOnExit>
        <Box sx={styles.traceDetailsWrapper}>
          <ChatToolTraceDetails tool={tool} labels={labels} />
        </Box>
      </Collapse>
    </Box>
  )
}