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 Box from '@mui/material/Box';
import Button from '@mui/material/Button';
import IconButton from '@mui/material/IconButton';
import CardContent from '@mui/material/CardContent';
import CardActions from '@mui/material/CardActions';
import Collapse from '@mui/material/Collapse';
import Typography from '@mui/material/Typography';
import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
import { ExpandMore } from './ExpandMore';
import ContentCopyIcon from '@mui/icons-material/ContentCopy';
import CheckIcon from '@mui/icons-material/Check';

import { ChatBubble } from './ChatBubble';
import { StyledMarkdown } from './StyledMarkdown';
import { Tooltip } from '@mui/material';

import { VectorVisualizer } from './VectorVisualizer';
import { SetSnackType } from './Snack';

type MessageRoles = 'info' | 'user' | 'assistant' | 'system' | 'status' | 'error' | 'content';

type MessageData = {
  role: MessageRoles,
  content: string,
  user?: string,
  title?: string,
  origin?: string,
  display?: string, /* Messages generated on the server for filler should not be shown */
  id?: string,
  isProcessing?: boolean,
  metadata?: MessageMetaProps
};

interface MessageMetaProps {
  query?: {
    query_embedding: number[];
    vector_embedding: number[];
  },
  origin: string,
  full_query?: string,
  rag: any,
  tools: any[],
  eval_count: number,
  eval_duration: number,
  prompt_eval_count: number,
  prompt_eval_duration: number,
  sessionId?: string,
  connectionBase: string,
  setSnack: SetSnackType,
}

type MessageList = MessageData[];

interface MessageProps {
  message?: MessageData,
  isFullWidth?: boolean,
  submitQuery?: (text: string) => void,
  sessionId?: string,
  connectionBase: string,
  setSnack: SetSnackType,
};

interface ChatQueryInterface {
  text: string,
  submitQuery?: (text: string) => void
}


const MessageMeta = ({ ...props }: MessageMetaProps) => {
  return (<>
    <Box sx={{ fontSize: "0.8rem", mb: 1 }}>
      Below is the LLM performance of this query. Note that if tools are called, the
      entire context is processed for each separate tool request by the LLM. This
      can dramatically increase the total time for a response.
    </Box>
    <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">{props.prompt_eval_count}</TableCell>
            <TableCell align="right">{Math.round(props.prompt_eval_duration / 10 ** 7) / 100}</TableCell>
            <TableCell align="right">{Math.round(props.prompt_eval_count * 10 ** 9 / props.prompt_eval_duration)}</TableCell>
          </TableRow>
          <TableRow key="response" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
            <TableCell component="th" scope="row">Response</TableCell>
            <TableCell align="right">{props.eval_count}</TableCell>
            <TableCell align="right">{Math.round(props.eval_duration / 10 ** 7) / 100}</TableCell>
            <TableCell align="right">{Math.round(props.eval_count * 10 ** 9 / props.eval_duration)}</TableCell>
          </TableRow>
          <TableRow key="total" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
            <TableCell component="th" scope="row">Total</TableCell>
            <TableCell align="right">{props.prompt_eval_count + props.eval_count}</TableCell>
            <TableCell align="right">{Math.round((props.prompt_eval_duration + props.eval_duration) / 10 ** 7) / 100}</TableCell>
            <TableCell align="right">{Math.round((props.prompt_eval_count + props.eval_count) * 10 ** 9 / (props.prompt_eval_duration + props.eval_duration))}</TableCell>
          </TableRow>
        </TableBody>
      </Table>
    </TableContainer>

    {
      props?.full_query !== undefined &&
      <Accordion>
        <AccordionSummary expandIcon={<ExpandMoreIcon />}>
          <Box sx={{ fontSize: "0.8rem" }}>
            Full Query
          </Box>
        </AccordionSummary>
        <AccordionDetails>
          <pre>{props.full_query}</pre>
        </AccordionDetails>
      </Accordion>
    }
    {
      props.tools !== undefined && props.tools.length !== 0 &&
      <Accordion sx={{ boxSizing: "border-box" }}>
        <AccordionSummary expandIcon={<ExpandMoreIcon />}>
          <Box sx={{ fontSize: "0.8rem" }}>
            Tools queried
          </Box>
        </AccordionSummary>
        <AccordionDetails>
          {props.tools.map((tool: any, index: number) => <Box key={index}>
            {index !== 0 && <Divider />}
            <Box sx={{ fontSize: "0.75rem", display: "flex", flexDirection: "column", mt: 0.5 }}>
              <div style={{ display: "flex", paddingRight: "1rem", whiteSpace: "nowrap" }}>
                {tool.tool}
              </div>
              <div style={{
                display: "flex",
                padding: "3px",
                whiteSpace: "pre-wrap",
                flexGrow: 1,
                border: "1px solid #E0E0E0",
                wordBreak: "break-all",
                maxHeight: "5rem",
                overflow: "auto"
              }}>
                {JSON.stringify(tool.result, null, 2)}
              </div>
            </Box>
          </Box>)}
        </AccordionDetails>
      </Accordion>
    }
    {
      props?.rag?.name !== undefined && <>
        <Accordion>
          <AccordionSummary expandIcon={<ExpandMoreIcon />}>
            <Box sx={{ fontSize: "0.8rem" }}>
              Top RAG {props.rag.ids.length} matches from '{props.rag.name}' collection against embedding vector of {props.rag.query_embedding.length} dimensions
            </Box>
          </AccordionSummary>
          <AccordionDetails>
            {props.rag.ids.map((id: number, index: number) => <Box key={index}>
              {index !== 0 && <Divider />}
              <Box sx={{ fontSize: "0.75rem", display: "flex", flexDirection: "row", mb: 0.5, mt: 0.5 }}>
                <div style={{ display: "flex", flexDirection: "column", paddingRight: "1rem", minWidth: "10rem" }}>
                  <div style={{ whiteSpace: "nowrap" }}>Doc ID: {props.rag.ids[index].slice(-10)}</div>
                  <div style={{ whiteSpace: "nowrap" }}>Similarity: {Math.round(props.rag.distances[index] * 100) / 100}</div>
                  <div style={{ whiteSpace: "nowrap" }}>Type: {props.rag.metadatas[index].doc_type}</div>
                  <div style={{ whiteSpace: "nowrap" }}>Chunk Len: {props.rag.documents[index].length}</div>
                </div>
                <div style={{ display: "flex", padding: "3px", flexGrow: 1, border: "1px solid #E0E0E0", maxHeight: "5rem", overflow: "auto" }}>{props.rag.documents[index]}</div>
              </Box>
            </Box>
            )}
          </AccordionDetails>
        </Accordion>
        <Accordion>
          <AccordionSummary expandIcon={<ExpandMoreIcon />}>
            <Box sx={{ fontSize: "0.8rem" }}>
              UMAP Vector Visualization of RAG
            </Box>
          </AccordionSummary>
          <AccordionDetails>
            <VectorVisualizer inline {...props} rag={props?.rag} />
          </AccordionDetails>
        </Accordion>
      </>
    }
  </>
  );
};

const ChatQuery = ({ text, submitQuery }: ChatQueryInterface) => {
  if (submitQuery === undefined) {
    return (<Box>{text}</Box>);
  }
  return (
    <Button variant="outlined" sx={{
      color: theme => theme.palette.custom.highlight, // Golden Ochre (#D4A017)
      borderColor: theme => theme.palette.custom.highlight,
      m: 1
    }}
      size="small" onClick={(e: any) => { submitQuery(text); }}>
      {text}
    </Button>
  );
}

const Message = ({ message, submitQuery, isFullWidth, sessionId, setSnack, connectionBase }: MessageProps) => {
  const [expanded, setExpanded] = useState<boolean>(false);
  const [copied, setCopied] = useState(false);
  const textFieldRef = useRef(null);

  const handleCopy = () => {
    if (message === undefined || message.content === undefined) {
      return;
    }

    navigator.clipboard.writeText(message.content.trim()).then(() => {
      setCopied(true);
      setTimeout(() => setCopied(false), 2000); // Reset after 2 seconds
    });
  };

  const handleExpandClick = () => {
    setExpanded(!expanded);
  };

  if (message === undefined) {
    return (<></>);
  }

  if (message.content === undefined) {
    console.info("Message content is undefined");
    return (<></>);
  }

  const formattedContent = message.content.trim();

  return (
    <ChatBubble
      className="Message"
      isFullWidth={isFullWidth}
      role={message.role}
      title={message.title}
      sx={{
        display: "flex",
        flexDirection: "column",
        pb: message.metadata ? 0 : "8px",
        m: 0,
        mb: 1,
        mt: 1,
        overflowX: "auto"
      }}>
      <CardContent ref={textFieldRef} sx={{ position: "relative", display: "flex", flexDirection: "column", overflowX: "auto", m: 0, p: 0 }}>
        <Tooltip title="Copy to clipboard" placement="top" arrow>
          <IconButton
            onClick={handleCopy}
            sx={{
              position: 'absolute',
              top: 0,
              right: 0,
              width: 24,
              height: 24,
              opacity: 0.75,
              bgcolor: 'background.paper',
              '&:hover': { bgcolor: 'action.hover', opacity: 1 },
            }}
            size="small"
            color={copied ? "success" : "default"}
          >
            {copied ? <CheckIcon sx={{ width: 16, height: 16 }} /> : <ContentCopyIcon sx={{ width: 16, height: 16 }} />}
          </IconButton>
        </Tooltip>

        {message.role !== 'user' ?
          <StyledMarkdown
            className="MessageContent"
            sx={{ display: "flex", color: 'text.secondary' }}
            {...{ content: formattedContent, submitQuery }} />
           :
          <Typography
            className="MessageContent"
            ref={textFieldRef}
            variant="body2"
            sx={{ display: "flex", color: 'text.secondary' }}>
            {message.content}
          </Typography>
        }
      </CardContent>
      {message.metadata && <>
        <CardActions disableSpacing sx={{ justifySelf: "flex-end" }}>
          <Button variant="text" onClick={handleExpandClick} sx={{ color: "darkgrey", p: 1, flexGrow: 0, justifySelf: "flex-end" }}>LLM information for this query</Button>
          <ExpandMore
            expand={expanded}
            onClick={handleExpandClick}
            aria-expanded={expanded}
            aria-label="show more"
          >
            <ExpandMoreIcon />
          </ExpandMore>
        </CardActions>
        <Collapse in={expanded} timeout="auto" unmountOnExit>
          <CardContent>
            <MessageMeta {...{ ...message.metadata, sessionId, connectionBase, setSnack }} />
          </CardContent>
        </Collapse>
      </>}
    </ChatBubble>
  );
};

export type {
  MessageProps, 
  MessageList,
  ChatQueryInterface,
  MessageMetaProps,
  MessageData,
  MessageRoles
};

export  {
  Message,
  ChatQuery,
  MessageMeta
};