package azkaban.utils;

import com.google.common.collect.ImmutableSet;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.List;
import java.util.Set;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * Utility class for getting system memory information
 *
 * Note:
 * This check is designed for Linux only.
 */
class OsMemoryUtil {
  private static final Logger logger = LoggerFactory.getLogger(OsMemoryUtil.class);

  // This file is used by Linux. It doesn't exist on Mac for example.
  private static final String MEM_INFO_FILE = "/proc/meminfo";

  private static final Set<String> MEM_KEYS = ImmutableSet.of("MemFree", "Buffers", "Cached", "SwapFree");

  /**
   * Includes OS cache and free swap.
   * @return the total free memory size of the OS. 0 if there is an error or the OS doesn't support this memory check.
   */
  long getOsTotalFreeMemorySize() {
    if (!Files.isRegularFile(Paths.get(MEM_INFO_FILE))) {
      // Mac doesn't support /proc/meminfo for example.
      return 0;
    }

    List<String> lines;
    // The file /proc/meminfo is assumed to contain only ASCII characters.
    // The assumption is that the file is not too big. So it is simpler to read the whole file into memory.
    try {
      lines = Files.readAllLines(Paths.get(MEM_INFO_FILE), StandardCharsets.UTF_8);
    } catch (IOException e) {
      String errMsg = "Failed to open mem info file: " + MEM_INFO_FILE;
      logger.error(errMsg, e);
      return 0;
    }
    return getOsTotalFreeMemorySizeFromStrings(lines);
  }

  /**
   *
   * @param lines text lines from the procinfo file
   * @return the total size of free memory in kB. 0 if there is an error.
   */
  long getOsTotalFreeMemorySizeFromStrings(List<String> lines) {
    long totalFree = 0;
    int count = 0;

    for (String line : lines) {
      for (String keyName : MEM_KEYS) {
        if (line.startsWith(keyName)) {
          count++;
          long size = parseMemoryLine(line);
          if (size == 0) {
            return 0;
          }
          totalFree += size;
        }
      }
    }

    int length = MEM_KEYS.size();
    if (count != length) {
      String errMsg = String.format("Expect %d keys in the meminfo file. Got %d. content: %s", length, count, lines);
      logger.error(errMsg);
      totalFree = 0;
    }
    return totalFree;
  }

  /**
   * Example file:
   * $ cat /proc/meminfo
   *   MemTotal:       65894008 kB
   *   MemFree:        59400536 kB
   *   Buffers:          409348 kB
   *   Cached:          4290236 kB
   *   SwapCached:            0 kB
   *
   * Make the method package private to make unit testing easier.
   * Otherwise it can be made private.

   * @param line the text for a memory usage statistics we are interested in
   * @return size of the memory. unit kB. 0 if there is an error.
   */
  long parseMemoryLine(String line) {
    int idx1 = line.indexOf(":");
    int idx2 = line.lastIndexOf("kB");
    String sizeString = line.substring(idx1 + 1, idx2 - 1).trim();
    try {
      return Long.parseLong(sizeString);
    } catch (NumberFormatException e) {
      String err = "Failed to parse the meminfo file. Line: " + line;
      logger.error(err);
      return 0;
    }
  }
}