import Handlebars from 'handlebars';

export enum MessageRole {
  SYSTEM = 'system',
  USER = 'user',
  ASSISTANT = 'assistant',
}

export interface MessageObject {
  role: MessageRole;
  content: string;
}

export type Message = string | MessageObject;

const isValidMessageObject = (maybeMessage: unknown): maybeMessage is Message => {
  if (typeof maybeMessage !== 'object') {
    return false;
  }

  if (maybeMessage === null) {
    return false;
  }

  const messageObject = maybeMessage as Record<string, unknown>;

  return (
    Object.values(MessageRole).includes(messageObject.role as MessageRole)
    && typeof messageObject.content === 'string'
  );
};

export const normalizeChatMessage = (messageRaw: Message | Message[]): MessageObject[] => {
  if (typeof messageRaw === 'string') {
    return [
      {
        role: MessageRole.USER,
        content: messageRaw,
      },
    ];
  }

  if (Array.isArray(messageRaw)) {
    return messageRaw.map((message) => {
      if (typeof message === 'string') {
        return {
          role: MessageRole.USER,
          content: message,
        };
      }

      if (isValidMessageObject(message)) {
        return message;
      }

      throw new Error('Invalid message format');
    });
  }

  if (isValidMessageObject(messageRaw)) {
    return [messageRaw];
  }

  throw new Error('Invalid message format');
};

export const buildChatFromTranscript = (transcript: string) => {
  const parameterized = Handlebars.create().compile(transcript, {
    noEscape: true,
    ignoreStandalone: true,
    strict: true,
    preventIndent: true,
  });

  return (params: Record<string, unknown>) => {
    const compiled = parameterized(params);
    const prompts = compiled.split('\n---\n');
    return prompts.map((prompt) => {
      const lines = prompt.trim().split('\n\n');
      let lastRole = MessageRole.USER;
      return lines.filter((s) => s.trim().length > 0).map((lineRaw) => {
        const line = lineRaw.replace(/\n/g, ' ');
        const lineCheckRole = line.toLowerCase();
        if (lineCheckRole.startsWith('system:')) {
          lastRole = MessageRole.SYSTEM;
          return {
            role: MessageRole.SYSTEM,
            content: line.slice('system:'.length).trim(),
          };
        }

        if (lineCheckRole.startsWith('user:')) {
          lastRole = MessageRole.USER;
          return {
            role: MessageRole.USER,
            content: line.slice('user:'.length).trim(),
          };
        }

        if (lineCheckRole.startsWith('assistant:')) {
          lastRole = MessageRole.ASSISTANT;
          return {
            role: MessageRole.ASSISTANT,
            content: line.slice('assistant:'.length).trim(),
          };
        }

        return {
          role: lastRole,
          content: line.trim(),
        };
      });
    });
  };
};
