lib/messageUtils.ts

import toast from 'react-hot-toast';
import type { ChatMessage } from '@/lib/db/types';
import { Message } from 'ai';

const PAIRS: Record<string, string> = {
  '┍': '┑',
  '┝': '┥',
  '├': '┤',
  '┕': '┙',
};

const MIDDLE_STARTER = '┝';
const MIDDLE_SEPARATOR = '┿';

const ANSWERS_PREFIX = 'answers';
const INPUT_PREFIX = 'input';

export const generateAnswersImageMarkdown = (index: number, url: string) => {
  return `![${ANSWERS_PREFIX}-${index}](${url})`;
};

export const generateInputImageMarkdown = (url: string, index = 0) => {
  if (url.toLowerCase().endsWith('.mp4')) {
    const prefix = 'input-video';
    return `![${INPUT_PREFIX}-${index}](<${url}>)`;
  } else {
    const prefix = 'input';
    return `![${INPUT_PREFIX}-${index}](<${url}>)`;
  }
};

export const cleanInputMessage = (content: string) => {
  return content
    .replace(/!\[input-.*?\)/g, '')
    .replace(/<video[^>]*>.*?<\/video>/g, '');
};

export const cleanAnswerMessage = (content: string) => {
  return content.replace(/!\[answers.*?\.png\)/g, '');
};

const generateJSONArrayMarkdown = (
  message: string,
  payload: Array<Record<string, string | boolean>>,
) => {
  if (payload.length === 0) return '';
  const keys = Object.keys(payload[0]);
  message += '\n';
  message += '| ' + keys.join(' | ') + ' |' + '\n';
  message += new Array(keys.length + 1).fill('|').join(' :- ') + '\n';
  payload.forEach((obj: any) => {
    message +=
      '| ' +
      keys
        .map(key => {
          if (key === 'documentation') {
            const doc = `\`\`\`\n${obj[key]}\n\`\`\`\n`;
            return `<button data-details=${JSON.stringify(encodeURI(doc))}>Show</button>`;
          } else {
            return obj[key];
          }
        })
        .join(' | ') +
      ' |' +
      '\n';
  });
  message += '\n';
  return message;
};

const generateCodeExecutionMarkdown = (
  message: string,
  payload: {
    code: string;
    test: string;
    result?: string;
  },
) => {
  let Details = 'Code: \n';
  Details += `\`\`\`python\n${payload.code}\n\`\`\`\n`;
  Details += 'Test: \n';
  Details += `\`\`\`python\n${payload.test}\n\`\`\`\n`;
  if (payload.result) {
    Details += 'Execution result: \n';
    Details += `\`\`\`python\n${payload.result}\n\`\`\`\n`;
  }
  message += `<button data-details=${JSON.stringify(encodeURI(Details))}>View details</button> \n`;

  return message;
};

const generateFinalCodeMarkdown = (
  code: string,
  test: string,
  result: PrismaJson.FinalChatResult['payload']['result'],
) => {
  let message = 'Final Code: \n';
  message += `\`\`\`python\n${code}\n\`\`\`\n`;
  message += 'Final test: \n';
  message += `\`\`\`python\n${test}\n\`\`\`\n`;
  message += `Final result: \n`;
  const images = result.results.map(result => result.png).filter(png => !!png);
  if (images.length > 0) {
    message += `Visualization output:\n`;
    images.forEach((image, index) => {
      message += generateAnswersImageMarkdown(index, image!);
    });
  }
  if (result.logs.stderr.length > 0) {
    message += `Error output:\n`;
    message += `\`\`\`\n${result.logs.stderr.join('\n')}\n\`\`\`\n`;
  }
  if (result.logs.stdout.length > 0) {
    message += `Output:\n`;
    message += `\`\`\`\n${result.logs.stdout.join('\n')}\n\`\`\`\n`;
  }
  return message;
};

type PlansBody =
  | {
      type: 'plans';
      status: 'started';
    }
  | {
      type: 'plans';
      status: 'completed';
      payload: Array<Record<string, string>>;
    };

type ToolsBody =
  | {
      type: 'tools';
      status: 'started';
    }
  | {
      type: 'tools';
      status: 'completed';
      payload: Array<Record<string, string>>;
    };

type CodeBody =
  | {
      type: 'code';
      status: 'started';
    }
  | {
      type: 'code';
      status: 'running';
      payload: {
        code: string;
        test: string;
      };
    }
  | {
      type: 'code';
      status: 'completed' | 'failed';
      payload: {
        code: string;
        test: string;
        result: string;
      };
    };

// this will return if self_reflection flag is true
type ReflectionBody =
  | {
      type: 'self_reflection';
      status: 'started';
    }
  | {
      type: 'self_reflection';
      status: 'completed' | 'failed';
      payload: { feedback: string; success: boolean };
    };

type MessageBody =
  | PlansBody
  | ToolsBody
  | CodeBody
  | ReflectionBody
  | PrismaJson.FinalChatResult;

const getMessageTitle = (json: MessageBody) => {
  switch (json.type) {
    case 'plans':
      if (json.status === 'started') {
        return '🎬 Start generating plans...\n';
      } else {
        return '✅ Going to run the following plan(s) in sequence:\n';
      }
    case 'tools':
      if (json.status === 'started') {
        return '🎬 Start retrieving tools...\n';
      } else {
        return '✅ Tools retrieved:\n';
      }
    case 'code':
      if (json.status === 'started') {
        return '🎬 Start generating code...\n';
      } else if (json.status === 'running') {
        return '🎬 Code generated, start execution... ';
      } else if (json.status === 'completed') {
        return '✅ Code executed successfully. ';
      } else {
        return '❌ Code execution failed. ';
      }
    case 'self_reflection':
      if (json.status === 'started') {
        return '🎬 Start self reflection...\n';
      } else if (json.status === 'completed') {
        return '✅ Self reflection completed: \n';
      } else {
        return '❌ Self reflection failed: \n';
      }
    case 'final_code':
      if (json.status === 'completed') {
        return '✅ The vision agent has concluded the chat, the last execution is successful. \n';
      } else {
        return '❌ The vision agent has concluded the chat, the last execution is failed. \n';
      }
    default:
      throw 'Not supported type';
  }
};
const parseLine = (json: MessageBody) => {
  const title = getMessageTitle(json);
  if (json.status === 'started') {
    return title;
  }
  switch (json.type) {
    case 'plans':
      return generateJSONArrayMarkdown(title, json.payload);
    case 'tools':
      return generateJSONArrayMarkdown(title, json.payload);
    case 'code':
      return generateCodeExecutionMarkdown(title, json.payload);
    case 'self_reflection':
      return generateJSONArrayMarkdown(title, [json.payload]);
    case 'final_code':
      return '';
    default:
      throw 'Not supported type';
  }
};

export const getFormattedMessage = ({
  content,
  role,
}: Pick<ChatMessage, 'role' | 'content'>) => {
  if (role === 'user') {
    return {
      content,
    };
  }
  const lines = content.split('\n');
  let formattedLogs = '';
  let finalResult = null;
  const jsons: MessageBody[] = [];
  for (let line of lines) {
    console.log(line);
    if (!line.trim()) {
      continue;
    }
    try {
      const json = JSON.parse(line) as MessageBody;
      if (json.type === 'final_code') {
        const result = JSON.parse(
          json.payload.result as unknown as string,
        ) as PrismaJson.FinalChatResult['payload']['result'];
        finalResult = generateFinalCodeMarkdown(
          json.payload.code,
          json.payload.test,
          result,
        );
      } else if (
        jsons.length > 0 &&
        json.type === jsons[jsons.length - 1].type &&
        json.status !== 'started'
      ) {
        jsons[jsons.length - 1] = json;
      } else {
        jsons.push(json);
      }
    } catch (e) {
      console.error((e as Error).message);
      console.error(line);
    }
  }
  jsons.forEach(json => (formattedLogs += parseLine(json)));
  return {
    content: formattedLogs + (finalResult ?? ''),
  };
};

export const convertDbMessageToMessage = (message: ChatMessage): Message => {
  return {
    id: message.id,
    role: message.role,
    createdAt: message.createdAt,
    content: message.result
      ? message.content + '\n' + JSON.stringify(message.result)
      : message.content,
  };
};

export const convertMessageToDbMessage = (
  message: Message,
): Pick<ChatMessage, 'content' | 'result' | 'role'> => {
  let result = null;
  const lines = message.content.split('\n');
  for (let line of lines) {
    try {
      const json = JSON.parse(line) as MessageBody;
      if (json.type == 'final_code') {
        result = json as PrismaJson.FinalChatResult;
        break;
      }
    } catch (e) {
      console.error((e as Error).message);
    }
  }
  return {
    role: message.role as ChatMessage['role'],
    content: message.content,
    result,
  };
};

export type CodeResult = {
  code: string;
  test: string;
  result: string;
};

export type ChunkBody = {
  type: 'plans' | 'tools' | 'code' | 'final_code';
  status: 'started' | 'completed' | 'failed' | 'running';
  payload: Array<Record<string, string>> | CodeResult;
};

/**
 * Formats the stream logs and returns an array of grouped sections.
 *
 * @param content - The content of the stream logs.
 * @returns An array of grouped sections and an optional final code result.
 */
export const formatStreamLogs = (
  content: string,
): [ChunkBody[], CodeResult?] => {
  const streamLogs = content.split('\n').filter(log => !!log);

  let parsedStreamLogs: ChunkBody[] = [];
  try {
    parsedStreamLogs = streamLogs.map(streamLog => JSON.parse(streamLog));
  } catch {
    toast.error('Error parsing stream logs');
    return [[], undefined];
  }

  // Merge consecutive logs of the same type to the latest status
  const groupedSections = parsedStreamLogs.reduce((acc, curr) => {
    if (acc.length > 0 && acc[acc.length - 1].type === curr.type) {
      acc[acc.length - 1] = curr;
    } else {
      acc.push(curr);
    }
    return acc;
  }, [] as ChunkBody[]);

  return [
    groupedSections.filter(section => section.type !== 'final_code'),
    groupedSections.find(section => section.type === 'final_code')
      ?.payload as CodeResult,
  ];
};