import { useState, useRef, useEffect } 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 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 { SxProps, Theme } from '@mui/material';
import JsonView from '@uiw/react-json-view';

import { ChatBubble } from './ChatBubble';
import { StyledMarkdown } from './StyledMarkdown';

import { VectorVisualizer } from './VectorVisualizer';
import { SetSnackType } from './Snack';
import { CopyBubble } from './CopyBubble';
import { Scrollable } from './Scrollable';
import { BackstoryElementProps } from './BackstoryTab';

type MessageRoles =
  'assistant' |
  'content' |
  'error' |
  'fact-check' |
  'info' |
  'job-description' |
  'job-requirements' |
  'processing' |
  'qualifications' |
  'resume' |
  'status' |
  'streaming' |
  'system' |
  'thinking' |
  'user';

type MessageData = {
  role: MessageRoles,
  content: string,
  status?: string, // streaming, done, error...
  response?: string,
  disableCopy?: boolean,
  user?: string,
  title?: string,
  origin?: string,
  display?: string, /* Messages generated on the server for filler should not be shown */
  id?: string,
  isProcessing?: boolean,
  actions?: string[],
  metadata?: MessageMetaData,
  expanded?: boolean,
  expandable?: boolean,
};

interface MessageMetaData {
  query?: {
    query_embedding: number[];
    vector_embedding: number[];
  },
  origin: string,
  rag: any,
  tools?: {
    tool_calls: 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 extends BackstoryElementProps {
  sx?: SxProps<Theme>,
  message: MessageData,
  expanded?: boolean,
  onExpand?: (open: boolean) => void,
  className?: string,
};

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


const MessageMeta = (props: MessageMetaProps) => {
  const {
    /* MessageData */
    rag,
    tools,
    eval_count,
    eval_duration,
    prompt_eval_count,
    prompt_eval_duration,
  } = 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 (<>
    {
      prompt_eval_duration !== 0 && eval_duration !== 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">{prompt_eval_count}</TableCell>
                <TableCell align="right">{Math.round(prompt_eval_duration / 10 ** 7) / 100}</TableCell>
                <TableCell align="right">{Math.round(prompt_eval_count * 10 ** 9 / 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">{eval_count}</TableCell>
                <TableCell align="right">{Math.round(eval_duration / 10 ** 7) / 100}</TableCell>
                <TableCell align="right">{Math.round(eval_count * 10 ** 9 / 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">{prompt_eval_count + eval_count}</TableCell>
                <TableCell align="right">{Math.round((prompt_eval_duration + eval_duration) / 10 ** 7) / 100}</TableCell>
                <TableCell align="right">{Math.round((prompt_eval_count + eval_count) * 10 ** 9 / (prompt_eval_duration + eval_duration))}</TableCell>
              </TableRow>
            </TableBody>
          </Table>
        </TableContainer>
      </>
    }
    {
      tools !== undefined && 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>
                  <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>
                </Box>)
            }
        </AccordionDetails>
      </Accordion>
    }
    {
      rag.map((rag: any) => (
        <Accordion key={rag.name}>
          <AccordionSummary expandIcon={<ExpandMoreIcon />}>
            <Box sx={{ fontSize: "0.8rem" }}>
              Top RAG {rag.ids.length} matches from '{rag.name}' collection against embedding vector of {rag.query_embedding.length} dimensions
            </Box>
          </AccordionSummary>
          <AccordionDetails>
            <Box sx={{ fontSize: "0.8rem" }}>
              UMAP Vector Visualization of '{rag.name}' RAG
            </Box>
            <VectorVisualizer inline {...props.messageProps} {...props.metadata} rag={rag} />
            {rag.ids.map((id: number, index: number) => <Box key={index}>
              <Divider />
              <Box sx={{ whiteSpace: "nowrap", fontSize: "0.75rem", p: 0, m: 0, pt: 0.5 }}>Doc ID: {rag.ids[index]}</Box>
              <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" }}>Distance: {Math.round(rag.distances[index] * 100) / 100}</div>
                  <div style={{ whiteSpace: "nowrap" }}>Type: {rag.metadatas[index].doc_type}</div>
                  <div style={{ whiteSpace: "nowrap" }}>Chunk Len: {rag.documents[index].length}</div>
                </div>
                <div style={{ display: "flex", padding: "3px", flexGrow: 1, border: "1px solid #E0E0E0", maxHeight: "5rem", overflow: "auto" }}>{rag.documents[index]}</div>
              </Box>
            </Box>
            )}
          </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>
  </>);
};

const Message = (props: MessageProps) => {
  const { message, submitQuery, sx, className, onExpand, expanded } = props;
  const [metaExpanded, setMetaExpanded] = useState<boolean>(props.expanded || false);
  const textFieldRef = useRef(null);

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

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

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

  const formattedContent = message.content.trim() || "Waiting for LLM to spool up...";

  return (
    <ChatBubble
      className={className || "Message"}
      {...message}
      onExpand={onExpand}
      expanded={expanded}
      sx={{
        display: "flex",
        flexDirection: "column",
        pb: message.metadata ? 0 : "8px",
        m: 0,
        mt: 1,
        marginBottom: "0px !important", // Remove whitespace from expanded Accordion
        // overflowX: "auto"
        ...sx,
      }}>
      <CardContent ref={textFieldRef} sx={{ position: "relative", display: "flex", flexDirection: "column", overflowX: "auto", m: 0, p: 0, paddingBottom: '0px !important' }}>
        {message.role !== 'user' ?
          <Scrollable
            className="MessageContent"
            autoscroll
            fallbackThreshold={0.5}
            sx={{
              p: 0,
              m: 0,
              maxHeight: (message.role === "streaming") ? "20rem" : "unset",
              display: "flex",
              flexGrow: 1,
              overflow: "auto", /* Handles scrolling for the div */
            }}
          >
            <StyledMarkdown {...{ content: formattedContent, submitQuery }} />
          </Scrollable>
           :
          <Typography
            className="MessageContent"
            ref={textFieldRef}
            variant="body2"
            sx={{ display: "flex", color: 'text.secondary' }}>
            {message.content}
          </Typography>
        }
      </CardContent>
      <CardActions disableSpacing sx={{ display: "flex", flexDirection: "row", justifyContent: "space-between", alignItems: "center", width: "100%", p: 0, m: 0 }}>
        {(message.disableCopy === undefined || message.disableCopy === false) && ["assistant", "content"].includes(message.role) && <CopyBubble content={message.content} />}
        {message.metadata && (
          <Box sx={{ display: "flex", alignItems: "center", gap: 1 }}>
            <Button variant="text" onClick={handleMetaExpandClick} sx={{ color: "darkgrey", p: 0 }}>
              LLM information for this query
            </Button>
            <ExpandMore
              expand={metaExpanded}
              onClick={handleMetaExpandClick}
              aria-expanded={expanded}
              aria-label="show more"
            >
              <ExpandMoreIcon />
            </ExpandMore>
          </Box>
        )}
      </CardActions>
      {message.metadata && <>
        <Collapse in={expanded} timeout="auto" unmountOnExit>
          <CardContent>
            <MessageMeta messageProps={props} metadata={message.metadata} />
          </CardContent>
        </Collapse>
      </>}
    </ChatBubble>
  );
};

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

export  {
  Message,
  MessageMeta,
};