const DEFAULT_SYSTEM_PROMPT = 'You are a helpful and honest assistant. Please, respond concisely and truthfully.';
const DEFAULT_MODEL = "mistralai/Mistral-7B-Instruct-v0.3";

/**
 * Adds a custom menu to Google Sheets to allow users to input their Hugging Face API key.
 */
function onOpen() {
  const ui = SpreadsheetApp.getUi();
  ui.createMenu('Hugging Sheets')
    .addItem('Enter your HF API Key', 'showApiKeyPrompt')
    .addToUi();
}

/**
 * Prompts the user to enter their Hugging Face API key and saves it in the script properties.
 */
function showApiKeyPrompt() {
  const ui = SpreadsheetApp.getUi();
  const response = ui.prompt('Enter your Hugging Face API Key:');
  if (response.getSelectedButton() == ui.Button.OK) {
    const apiKey = response.getResponseText().trim();
    if (apiKey) {
      PropertiesService.getScriptProperties().setProperty('HF_API_KEY', apiKey);
      ui.alert('API Key saved successfully!');
    } else {
      ui.alert('API Key cannot be empty.');
    }
  }
}

/**
 * Sends a request to the Hugging Face API with the specified prompt and model.
 * 
 * @param {string} prompt - The input prompt to be sent to the model.
 * @param {string} [model] - The model ID to query (e.g., 'mistralai/Mistral-7B-Instruct-v0.3').
 * @param {string} [systemPrompt=DEFAULT_SYSTEM_PROMPT] - The system prompt to customize the assistant's behavior.
 * @returns {object} JSON response from the Hugging Face API.
 * @throws {Error} If the API key is not set or if the API request fails.
 */
function queryHuggingFace(prompt, model = DEFAULT_MODEL, systemPrompt = DEFAULT_SYSTEM_PROMPT) {
  const apiKey = PropertiesService.getScriptProperties().getProperty('HF_API_KEY');
  if (!apiKey) {
    throw new Error('Please enter your Hugging Face API key using the menu.');
  }

  const formattedPrompt = `<s> [INST] ${systemPrompt} ${prompt} [/INST] </s>`;
  const payload = {
    "inputs": formattedPrompt
  };

  const payloadString = JSON.stringify(payload);

  const url = `https://api-inference.huggingface.co/models/${model}`;
  const options = {
    'method': 'post',
    'contentType': 'application/json',
    'headers': {
      'Authorization': `Bearer ${apiKey}`
    },
    'payload': payloadString
  };

  try {
    const response = UrlFetchApp.fetch(url, options);
    const json = JSON.parse(response.getContentText());
    return json;
  } catch (error) {
    throw new Error(`Failed to fetch data from Hugging Face API: ${error.message}`);
  }
}

/**
 * Custom Google Sheets formula to query the Hugging Face API and return the generated text.
 * Function to create the custom formula =HF(prompt, model, [systemPrompt])
 * 
 * @param {string} prompt - The input prompt to be sent to the model.
 * @param {string} [model] - The model ID to query (e.g., 'mistralai/Mistral-7B-Instruct-v0.3').
 * @param {string} [systemPrompt] - The system prompt to customize the assistant's behavior. Defaults to DEFAULT_SYSTEM_PROMPT.
 * @returns {string} The generated output text from the Hugging Face API, or an error message if the request fails.
 */
function HF(prompt, model = DEFAULT_MODEL, systemPrompt = DEFAULT_SYSTEM_PROMPT) {
  try {
    const response = queryHuggingFace(prompt, model, systemPrompt);
    if (response && response.length > 0 && response[0].generated_text) {
      const fullResponse = response[0].generated_text;
      // Extract the part of the response after the prompt
      const generatedOutput = fullResponse.split(`</s>`).pop().trim();
      return generatedOutput;
    } else {
      return 'Error: Invalid response structure from Hugging Face API.';
    }
  } catch (error) {
    return `Error: ${error.message}`;
  }
}

/**
 * Add the formula to Google Sheets.
 * 
 * @param {object} e - The event object.
 */
function onInstall(e) {
  onOpen(e);
  const formula = SpreadsheetApp.newUserDefinedFunctionBuilder()
    .setName('HF')
    .setFunction('HF')
    .build();
  SpreadsheetApp.installUserDefinedFunction(formula);
}