backstory/frontend/src/components/StyledMarkdown.tsx

import React from 'react';
import { MuiMarkdown } from 'mui-markdown';
import { useTheme } from '@mui/material/styles';
import { Link } from '@mui/material';
import { BackstoryQuery, BackstoryQueryInterface } from 'components/BackstoryQuery';
import Box from '@mui/material/Box';
import JsonView from '@uiw/react-json-view';
import { vscodeTheme } from '@uiw/react-json-view/vscode';
import { Mermaid } from 'components/Mermaid';
import { Scrollable } from 'components/Scrollable';
import { jsonrepair } from 'jsonrepair';
import { GenerateImage } from 'components/GenerateImage';

import './StyledMarkdown.css';
import { BackstoryElementProps } from './BackstoryTab';
import { CandidateQuestion, ChatQuery, ChatSession } from 'types/types';
import { ChatSubmitQueryInterface } from 'components/BackstoryQuery';

interface StyledMarkdownProps extends BackstoryElementProps {
  className?: string,
  content: string,
  streaming?: boolean,
  chatSession?: ChatSession,
  submitQuery?: ChatSubmitQueryInterface
};

const StyledMarkdown: React.FC<StyledMarkdownProps> = (props: StyledMarkdownProps) => {
  const { className, content, chatSession, submitQuery, sx, streaming } = props;
  const theme = useTheme();

  const overrides: any = {
    p: { component: (element: any) =>{
      return <div>{element.children}</div>
    }},
    pre: {
      component: (element: any) => {
        const { className } = element.children.props;
        const content = element.children?.props?.children || "";
        if (className === "lang-mermaid" && !streaming) {
          return <Mermaid className="Mermaid" chart={content} />;
        }
        if (className === "lang-markdown") {
          return <MuiMarkdown children={content} />;
        }
        if (className === "lang-json" && !streaming) {
          try {
            let fixed = JSON.parse(jsonrepair(content));
            return <Scrollable className="JsonViewScrollable">
              <JsonView
                className="JsonView"
                style={{
                  ...vscodeTheme,
                  fontSize: "0.8rem",
                  maxHeight: "10rem",
                  padding: "14px 0",
                  overflow: "hidden",
                  width: "100%",
                  minHeight: "max-content",
                  backgroundColor: "transparent",
                }}
                displayDataTypes={false}
                objectSortKeys={false}
                collapsed={1}
                shortenTextAfterLength={100}
                value={fixed}>
                <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>
            </Scrollable>
          } catch (e) {
            return <pre><code className="JsonRaw">{content}</code></pre>
          };
        }
        return <pre><code className={className || ''}>{element.children}</code></pre>;
      },
    },
    a: {
      component: Link,
      props: {
        onClick: (event: React.MouseEvent<HTMLAnchorElement>) => {
          const href = event.currentTarget.getAttribute('href');
          console.log("StyledMarkdown onClick:", href);
          if (href) {
            if (href.match(/^\//)) {
              event.preventDefault();
              window.history.replaceState({}, '', `${href}`);
            }
          }
        },
        sx: {
          wordBreak: "break-all",
          color: theme.palette.secondary.main,
          textDecoration: 'none',
          '&:hover': {
            color: theme.palette.custom.highlight,
            textDecoration: 'underline',
          }
        }
      }
    },
    BackstoryQuery: {
      component: (props: { query: string }) => {
        const queryString = props.query.replace(/(\w+):/g, '"$1":');
        try {
          const query = JSON.parse(queryString);
          const backstoryQuestion: CandidateQuestion = {
            question: queryString
          }

          return submitQuery ? <BackstoryQuery submitQuery={submitQuery} question={query} /> : query.question;
        } catch (e) {
          console.log("StyledMarkdown error:", queryString, e);
          return props.query;
        }
      },
    }
  };

  if (chatSession) {
    overrides.GenerateImage = {
      component: (props: { prompt: string }) => {
        const prompt = props.prompt.replace(/(\w+):/g, '"$1":');
        try {
          return <GenerateImage {...{ chatSession, prompt }} />
        } catch (e) {
          console.log("StyledMarkdown error:", prompt, e);
          return props.prompt;
        }
      }
    }
  }

  return <Box
    className={`MuiMarkdown ${className || ""}`}
    sx={{
      display: "flex",
      m: 0,
      p: 0,
      boxSizing: "border-box",
      flexGrow: 1,
      height: "auto",
      ...sx
    }}>
    <MuiMarkdown
      overrides={overrides}
      children={content}
    />
  </Box>;
};

export { StyledMarkdown };