import Mustache from "mustache";

import type { PromptTemplate } from "../types/templating";

type GetTemplateVariableArgs = {
  template: PromptTemplate;
};
/**
 * Parse out the template variables of a prompt
 * @param {GetTemplateVariableArgs} args
 * @returns {string[]} a list of prompt template variables
 */
export function getTemplateVariables(args: GetTemplateVariableArgs): string[] {
  const { template } = args;
  if (typeof template === "string") {
    return getTemplateVariablesFromString(template);
  }
  return template.reduce<string[]>((acc, message) => {
    if (message.content !== undefined && typeof message.content === "string") {
      return [...acc, ...getTemplateVariablesFromString(message.content)];
    }
    return acc;
  }, []);
}
/**
 * Parse out the template variables of a string template
 * @param template - The template to get the variables from
 * @returns {string[]} a list of prompt template variables
 */
function getTemplateVariablesFromString(template: string): string[] {
  const templateSpans = Mustache.parse(template);
  return templateSpans.reduce<string[]>((acc, templateSpan) => {
    const [spanType, value] = templateSpan;
    if (spanType === "name" && typeof value === "string") {
      acc = [...acc, value];
    }
    return acc;
  }, []);
}