import { useState, useRef } from 'react';
import Divider from '@mui/material/Divider';
import Accordion from '@mui/material/Accordion';
import AccordionSummary from '@mui/material/AccordionSummary';
import AccordionDetails from '@mui/material/AccordionDetails';
import Card from '@mui/material/Card';
import Table from '@mui/material/Table';
import TableBody from '@mui/material/TableBody';
import TableCell from '@mui/material/TableCell';
import TableContainer from '@mui/material/TableContainer';
import TableHead from '@mui/material/TableHead';
import TableRow from '@mui/material/TableRow';
import Button from '@mui/material/Button';
import CardContent from '@mui/material/CardContent';
import CardActions from '@mui/material/CardActions';
import Collapse from '@mui/material/Collapse';
import { ExpandMore } from './ExpandMore';
import JsonView from '@uiw/react-json-view';
import React from 'react';
import { Box } from '@mui/material';
import { useTheme } from '@mui/material/styles';
import { SxProps, Theme } from '@mui/material';
import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
import LocationSearchingIcon from '@mui/icons-material/LocationSearching';

import { ErrorOutline, InfoOutline, Memory, Psychology, /* Stream, */ } from '@mui/icons-material';

import { StyledMarkdown } from './StyledMarkdown';

import { VectorVisualizer } from './VectorVisualizer';
import { SetSnackType } from './Snack';
import { CopyBubble } from './CopyBubble';
import { Scrollable } from './Scrollable';
import { BackstoryElementProps } from './BackstoryTab';
import { ChatMessage, ChatSession, ChatMessageMetaData, ChromaDBGetResponse, ApiActivityType, ChatMessageUser, ChatMessageError, ChatMessageStatus, ChatSenderType } from 'types/types';

const getStyle = (theme: Theme, type: ApiActivityType | ChatSenderType | "error"): any => {
  const defaultRadius = '16px';
  const defaultStyle = {
    padding: theme.spacing(1, 2),
    fontSize: '0.875rem',
    alignSelf: 'flex-start',
    maxWidth: '100%',
    minWidth: '100%',
    height: 'fit-content',
    '& > *': {
      color: 'inherit',
      overflow: 'hidden',
      m: 0,
    },
    '& > :last-child': {
      mb: 0,
      m: 0,
      p: 0,
    },
  };

  const styles: any = {
    assistant: {
      ...defaultStyle,
      backgroundColor: theme.palette.primary.main,
      border: `1px solid ${theme.palette.secondary.main}`,
      borderRadius: `${defaultRadius} ${defaultRadius} ${defaultRadius} 0`,
      color: theme.palette.primary.contrastText,
    },
    content: {
      ...defaultStyle,
      backgroundColor: '#F5F2EA',
      border: `1px solid ${theme.palette.custom.highlight}`,
      borderRadius: 0,
      alignSelf: 'center',
      color: theme.palette.text.primary,
      padding: '8px 8px',
      marginBottom: '0px',
      boxShadow: '0 2px 4px rgba(0, 0, 0, 0.05)',
      fontSize: '0.9rem',
      lineHeight: '1.3',
      fontFamily: theme.typography.fontFamily,
    },
    error: {
      ...defaultStyle,
      backgroundColor: '#F8E7E7',
      border: `1px solid #D83A3A`,
      borderRadius: defaultRadius,
      maxWidth: '90%',
      minWidth: '90%',
      alignSelf: 'center',
      color: '#8B2525',
      padding: '10px 16px',
      boxShadow: '0 1px 3px rgba(216, 58, 58, 0.15)',
    },
    'fact-check': 'qualifications',
    generating: 'status',
    'job-description': 'content',
    'job-requirements': 'qualifications',
    information: {
      ...defaultStyle,
      backgroundColor: '#BFD8D8',
      border: `1px solid ${theme.palette.secondary.main}`,
      borderRadius: defaultRadius,
      color: theme.palette.text.primary,
      opacity: 0.95,
    },
    preparing: 'status',
    processing: 'status',
    qualifications: {
      ...defaultStyle,
      backgroundColor: theme.palette.primary.light,
      border: `1px solid ${theme.palette.secondary.main}`,
      borderRadius: `${defaultRadius} ${defaultRadius} ${defaultRadius} 0`,
      color: theme.palette.primary.contrastText,
    },
    resume: 'content',
    searching: 'status',
    status: {
      ...defaultStyle,
      backgroundColor: 'rgba(74, 122, 125, 0.15)',
      border: `1px solid ${theme.palette.secondary.light}`,
      borderRadius: '4px',
      maxWidth: '75%',
      minWidth: '75%',
      alignSelf: 'center',
      color: theme.palette.secondary.dark,
      fontWeight: 500,
      fontSize: '0.95rem',
      padding: '8px 12px',
      opacity: 0.9,
      transition: 'opacity 0.3s ease-in-out',
    },
    streaming: 'response',
    system: {
      ...defaultStyle,
      backgroundColor: '#EDEAE0',
      border: `1px dashed ${theme.palette.custom.highlight}`,
      borderRadius: defaultRadius,
      maxWidth: '90%',
      minWidth: '90%',
      alignSelf: 'center',
      color: theme.palette.text.primary,
      fontStyle: 'italic',
    },
    thinking: 'status',
    user: {
      ...defaultStyle,
      backgroundColor: theme.palette.background.default,
      border: `1px solid ${theme.palette.custom.highlight}`,
      borderRadius: `${defaultRadius} ${defaultRadius} 0 ${defaultRadius}`,
      alignSelf: 'flex-end',
      color: theme.palette.primary.main,
    },
  };

  // Resolve string references in styles
  for (const [key, value] of Object.entries(styles)) {
    if (typeof value === 'string') {
      styles[key] = styles[value];
    }
  }

  if (!(type in styles)) {
    console.log(`Style does not exist for: ${type}`);
  }

  return styles[type];
}

const getIcon = (activityType: ApiActivityType | ChatSenderType | "error"): React.ReactNode | null => {
  const icons: any = {
    error: <ErrorOutline color="error" />,
    generating: <LocationSearchingIcon />,
    information: <InfoOutline color="info" />,
    preparing: <LocationSearchingIcon />,
    processing: <LocationSearchingIcon />,
    system: <Memory />,
    thinking: <Psychology />,
    tooling: <LocationSearchingIcon />,
  };
  return icons[activityType] || null;
}

interface MessageProps extends BackstoryElementProps {
  message: ChatMessageUser | ChatMessage | ChatMessageError | ChatMessageStatus,
  title?: string,
  chatSession?: ChatSession,
  className?: string,
  sx?: SxProps<Theme>,
  expandable?: boolean,
  expanded?: boolean,
  onExpand?: (open: boolean) => void,
};

interface MessageMetaProps {
  metadata: ChatMessageMetaData,
  messageProps: MessageProps
};

const MessageMeta = (props: MessageMetaProps) => {
  const {
    /* MessageData */
    ragResults = [],
    tools = null,
    evalCount = 0,
    evalDuration = 0,
    promptEvalCount = 0,
    promptEvalDuration = 0,
  } = props.metadata || {};
  const message: any = props.messageProps.message;

  let llm_submission: string = "<|system|>\n"
  llm_submission += message.system_prompt + "\n\n"
  llm_submission += message.context_prompt

  return (<>
    {
      promptEvalDuration !== 0 && evalDuration !== 0 && <>
        <TableContainer component={Card} className="PromptStats" sx={{ mb: 1 }}>
          <Table aria-label="prompt stats" size="small">
            <TableHead>
              <TableRow>
                <TableCell></TableCell>
                <TableCell align="right" >Tokens</TableCell>
                <TableCell align="right">Time (s)</TableCell>
                <TableCell align="right">TPS</TableCell>
              </TableRow>
            </TableHead>
            <TableBody>
              <TableRow key="prompt" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
                <TableCell component="th" scope="row">Prompt</TableCell>
                <TableCell align="right">{promptEvalCount}</TableCell>
                <TableCell align="right">{Math.round(promptEvalDuration / 10 ** 7) / 100}</TableCell>
                <TableCell align="right">{Math.round(promptEvalCount * 10 ** 9 / promptEvalDuration)}</TableCell>
              </TableRow>
              <TableRow key="response" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
                <TableCell component="th" scope="row">Response</TableCell>
                <TableCell align="right">{evalCount}</TableCell>
                <TableCell align="right">{Math.round(evalDuration / 10 ** 7) / 100}</TableCell>
                <TableCell align="right">{Math.round(evalCount * 10 ** 9 / evalDuration)}</TableCell>
              </TableRow>
              <TableRow key="total" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
                <TableCell component="th" scope="row">Total</TableCell>
                <TableCell align="right">{promptEvalCount + evalCount}</TableCell>
                <TableCell align="right">{Math.round((promptEvalDuration + evalDuration) / 10 ** 7) / 100}</TableCell>
                <TableCell align="right">{Math.round((promptEvalCount + evalCount) * 10 ** 9 / (promptEvalDuration + evalDuration))}</TableCell>
              </TableRow>
            </TableBody>
          </Table>
        </TableContainer>
      </>
    }
    {
      tools && tools.tool_calls && tools.tool_calls.length !== 0 &&
      <Accordion sx={{ boxSizing: "border-box" }}>
        <AccordionSummary expandIcon={<ExpandMoreIcon />}>
          <Box sx={{ fontSize: "0.8rem" }}>
            Tools queried
          </Box>
        </AccordionSummary>
        <AccordionDetails>
            {
              tools.tool_calls.map((tool: any, index: number) =>
                <Box key={index} sx={{ m: 0, p: 1, pt: 0, display: "flex", flexDirection: "column", border: "1px solid #e0e0e0" }}>
                  {index !== 0 && <Divider />}
                  <Box sx={{ fontSize: "0.75rem", display: "flex", flexDirection: "column", mt: 1, mb: 1, fontWeight: "bold" }}>
                    {tool.name}
                  </Box>
                  {tool.content !== "null" &&
                    <JsonView
                      displayDataTypes={false}
                      objectSortKeys={true}
                      collapsed={1} value={JSON.parse(tool.content)} style={{ fontSize: "0.8rem", maxHeight: "20rem", overflow: "auto" }}>
                      <JsonView.String
                        render={({ children, ...reset }) => {
                          if (typeof (children) === "string" && children.match("\n")) {
                            return <pre {...reset} style={{ display: "flex", border: "none", ...reset.style }}>{children}</pre>
                          }
                        }}
                      />
                    </JsonView>
                  }
                  {tool.content === "null" && "No response from tool call"}
                </Box>)
            }
        </AccordionDetails>
      </Accordion>
    }
    {
      ragResults.map((collection: ChromaDBGetResponse) => (
        <Accordion key={collection.name}>
          <AccordionSummary expandIcon={<ExpandMoreIcon />}>
            <Box sx={{ fontSize: "0.8rem" }}>
              Top {collection.ids?.length} RAG matches from {collection.size} entries using an embedding vector of {collection.queryEmbedding?.length} dimensions
            </Box>
          </AccordionSummary>
          <AccordionDetails>
            <VectorVisualizer inline
              {...props.messageProps} {...props.metadata}
              rag={collection} />
            {/* { ...rag, query: message.prompt }} /> */}
          </AccordionDetails>
        </Accordion>
      ))
    }

    <Accordion>
      <AccordionSummary expandIcon={<ExpandMoreIcon />}>
        <Box sx={{ fontSize: "0.8rem" }}>
          Full Response Details
        </Box>
      </AccordionSummary>
      <AccordionDetails>
        <Box sx={{ pb: 1 }}>Copy LLM submission: <CopyBubble content={llm_submission} /></Box>
        <JsonView displayDataTypes={false} objectSortKeys={true} collapsed={1} value={message} style={{ fontSize: "0.8rem", maxHeight: "20rem", overflow: "auto" }}>
          <JsonView.String
            render={({ children, ...reset }) => {
              if (typeof (children) === "string" && children.match("\n")) {
                return <pre {...reset} style={{ display: "inline", border: "none", ...reset.style }}>{children.trim()}</pre>
              }
            }}
          />
        </JsonView>
      </AccordionDetails>
    </Accordion>
  </>);
};

interface MessageContainerProps {
  type: ApiActivityType | ChatSenderType | "error",
  metadataView?: React.ReactNode | null,
  messageView?: React.ReactNode | null,
  sx?: SxProps<Theme>,
  copyContent?: string,
};

const MessageContainer = (props: MessageContainerProps) => {
  const { type, sx, messageView, metadataView, copyContent } = props;
  const icon = getIcon(type);

  return <Box
    className={`Message Message-${type}`}
    sx={{
      display: "flex",
      flexDirection: "column",
      m: 0,
      mt: 1,
      marginBottom: "0px !important", // Remove whitespace from expanded Accordion
      gap: 1,
      ...sx,
    }}>
    <Box sx={{ display: "flex", flexDirection: 'row', alignItems: 'center', gap: 1 }}>
      {icon !== null && icon}
      {messageView}
    </Box>
    <Box flex={{ display: "flex", position: "relative", flexDirection: "row", justifyContent: "flex-end", alignItems: "center" }}>
      {copyContent && <CopyBubble content={copyContent} sx={{ position: "absolute", top: "11px", left: 0 }} />}
      {metadataView}
    </Box>
  </Box>;
};

const Message = (props: MessageProps) => {
  const { message, title, sx, className, chatSession, onExpand, expanded, expandable } = props;
  const [metaExpanded, setMetaExpanded] = useState<boolean>(false);
  const theme = useTheme();
  const type: ApiActivityType | ChatSenderType | "error" = ('activity' in message) ? message.activity : ('error' in message) ? 'error' : (message as ChatMessage).role;
  const style: any = getStyle(theme, type);

  const handleMetaExpandClick = () => {
    setMetaExpanded(!metaExpanded);
  };

  let content;
  if (typeof (message.content) === "string") {
    content = message.content.trim();
  } else {
    console.error(`message content is not a string`);
    return (<></>)
  }

  if (!content) {
    return (<></>)
  };

  const messageView = (
    <StyledMarkdown chatSession={chatSession} streaming={message.status === "streaming"} content={content} />
  );

  let metadataView = (<></>);
  const metadata: ChatMessageMetaData | null = ('metadata' in message) ? (message.metadata as ChatMessageMetaData || null) : null;
  if (metadata) {
    metadataView = (
      <Box sx={{ display: "flex", flexDirection: "column", width: "100%" }}>
        <Box sx={{ display: "flex", alignItems: "center", gap: 1, flexDirection: "row" }}>
          <Box sx={{ display: "flex", flexGrow: 1 }} />
          <Button variant="text" onClick={handleMetaExpandClick} sx={{ flexShrink: 1, color: "darkgrey", p: 0 }}>
            LLM information for this query
          </Button>
          <ExpandMore
            sx={{ flexShrink: 1 }}
            expand={metaExpanded}
            onClick={handleMetaExpandClick}
            aria-expanded={true /*message.expanded*/}
            aria-label="show more">
            <ExpandMoreIcon />
          </ExpandMore>
        </Box>
        <Collapse in={metaExpanded} timeout="auto" unmountOnExit>
          <CardContent>
            <MessageMeta messageProps={props} metadata={metadata} />
          </CardContent>
        </Collapse>
      </Box>);
  }

  const copyContent = (type === 'assistant') ? message.content : undefined;

  if (!expandable) {
    /* When not expandable, the styles are applied directly to MessageContainer */
    return (<>
      {messageView && <MessageContainer copyContent={copyContent} type={type} {...{ messageView, metadataView }} sx={{ ...style, ...sx }} />}
    </>);
  }

  // Determine if Accordion is controlled
  const isControlled = typeof expanded === 'boolean' && typeof onExpand === 'function';
  return (
    <Accordion
      expanded={isControlled ? expanded : undefined} // Omit expanded prop for uncontrolled
      defaultExpanded={expanded} // Default to collapsed for uncontrolled Accordion
      className={className}
      onChange={(_event, newExpanded) => { isControlled && onExpand && onExpand(newExpanded) }}
      sx={{ ...sx, ...style }}>
      <AccordionSummary
        expandIcon={<ExpandMoreIcon />}
        slotProps={{
          content: {
            sx: {
              display: 'flex',
              justifyItems: 'center',
              m: 0, p: 0,
              fontWeight: 'bold',
              fontSize: '1.1rem',
            },
          },
        }}>
        {title || ''}
      </AccordionSummary>
      <AccordionDetails sx={{ mt: 0, mb: 0, p: 0, pl: 2, pr: 2 }}>
        <MessageContainer copyContent={copyContent} type={type} {...{ messageView, metadataView }} />
      </AccordionDetails>
    </Accordion>
  );
}

export type {
  MessageProps,
};

export  {
  Message,
  MessageMeta,
};