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 CardContent from '@mui/material/CardContent';
import CardActions from '@mui/material/CardActions';
import Collapse from '@mui/material/Collapse';
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 '../NewApp/Components/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 BackstoryMessage = {
  // Only two required fields
  role: MessageRoles,
  content: string,
  // Rest are optional
  prompt?: string;
  preamble?: {};
  status?: string;
  remaining_time?: number;
  full_content?: string;
  response?: string;     // Set when status === 'done', 'partial', or 'error'
  chunk?: string;        // Used when status === 'streaming'
  timestamp?: number;
  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,
  connectionBase: string,
  setSnack: SetSnackType,
}

type MessageList = BackstoryMessage[];

interface MessageProps extends BackstoryElementProps {
  sx?: SxProps<Theme>,
  message: BackstoryMessage,
  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 && 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>
    }
    {
      rag.map((collection: any) => (
        <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.query_embedding.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>
  </>);
};

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

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

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

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

  const formattedContent = message.content.trim();
  if (formattedContent === "") {
    return (<></>);
  }

  return (
    <ChatBubble
      className={`${className || ""} Message Message-${message.role}`}
      {...message}
      expanded={expanded}
      onExpand={onExpand}
      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' }}>
        <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 streaming={message.role === "streaming"} content={formattedContent} {...backstoryProps} />
        </Scrollable>
      </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) && <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={message.expanded}
              aria-label="show more"
            >
              <ExpandMoreIcon />
            </ExpandMore>
          </Box>
        )}
      </CardActions>
      {message.metadata && <>
        <Collapse in={metaExpanded} timeout="auto" unmountOnExit>
          <CardContent>
            <MessageMeta messageProps={props} metadata={message.metadata} />
          </CardContent>
        </Collapse>
      </>}
    </ChatBubble>
  );
};

export type {
  MessageProps, 
  MessageList,
  BackstoryMessage,
  MessageMetaData,
  MessageRoles,
};

export  {
  Message,
  MessageMeta,
};

