/*
 * Copyright 2017 Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <array>
#include <cstdint>
#include <memory>
#include <string>

#include <folly/FBString.h>
#include <folly/Range.h>
#include <folly/String.h>
#include <folly/io/IOBuf.h>
#include <folly/experimental/symbolizer/Elf.h>
#include <folly/experimental/symbolizer/ElfCache.h>
#include <folly/experimental/symbolizer/Dwarf.h>
#include <folly/experimental/symbolizer/StackTrace.h>

namespace folly {
namespace symbolizer {

class Symbolizer;

/**
 * Frame information: symbol name and location.
 */
struct SymbolizedFrame {
  SymbolizedFrame() { }

  void set(const std::shared_ptr<ElfFile>& file,
           uintptr_t address,
           Dwarf::LocationInfoMode mode);

  void clear() { *this = SymbolizedFrame(); }

  bool found = false;
  const char* name = nullptr;
  Dwarf::LocationInfo location;

  /**
   * Demangle the name and return it. Not async-signal-safe; allocates memory.
   */
  fbstring demangledName() const {
    return name ? demangle(name) : fbstring();
  }

 private:
  std::shared_ptr<ElfFile> file_;
};

template <size_t N>
struct FrameArray {
  FrameArray() { }

  size_t frameCount = 0;
  uintptr_t addresses[N];
  SymbolizedFrame frames[N];
};

/**
 * Get stack trace into a given FrameArray, return true on success (and
 * set frameCount to the actual frame count, which may be > N) and false
 * on failure.
 */
namespace detail {
template <size_t N>
bool fixFrameArray(FrameArray<N>& fa, ssize_t n) {
  if (n != -1) {
    fa.frameCount = n;
    for (size_t i = 0; i < fa.frameCount; ++i) {
      fa.frames[i].found = false;
    }
    return true;
  } else {
    fa.frameCount = 0;
    return false;
  }
}
}  // namespace detail

// Always inline these functions; they don't do much, and unittests rely
// on them never showing up in a stack trace.
template <size_t N>
FOLLY_ALWAYS_INLINE bool getStackTrace(FrameArray<N>& fa);

template <size_t N>
inline bool getStackTrace(FrameArray<N>& fa) {
  return detail::fixFrameArray(fa, getStackTrace(fa.addresses, N));
}
template <size_t N>
FOLLY_ALWAYS_INLINE bool getStackTraceSafe(FrameArray<N>& fa);

template <size_t N>
inline bool getStackTraceSafe(FrameArray<N>& fa) {
  return detail::fixFrameArray(fa, getStackTraceSafe(fa.addresses, N));
}

class Symbolizer {
 public:
  static constexpr Dwarf::LocationInfoMode kDefaultLocationInfoMode =
      Dwarf::LocationInfoMode::FAST;

  explicit Symbolizer(Dwarf::LocationInfoMode mode = kDefaultLocationInfoMode)
    : Symbolizer(nullptr, mode) {}

  explicit Symbolizer(ElfCacheBase* cache,
                      Dwarf::LocationInfoMode mode = kDefaultLocationInfoMode);

  /**
   * Symbolize given addresses.
   */
  void symbolize(const uintptr_t* addresses,
                 SymbolizedFrame* frames,
                 size_t frameCount);

  template <size_t N>
  void symbolize(FrameArray<N>& fa) {
    symbolize(fa.addresses, fa.frames, fa.frameCount);
  }

  /**
   * Shortcut to symbolize one address.
   */
  bool symbolize(uintptr_t address, SymbolizedFrame& frame) {
    symbolize(&address, &frame, 1);
    return frame.found;
  }

 private:
  ElfCacheBase* const cache_;
  const Dwarf::LocationInfoMode mode_;
};

/**
 * Format one address in the way it's usually printed by SymbolizePrinter.
 * Async-signal-safe.
 */
class AddressFormatter {
 public:
  AddressFormatter();

  /**
   * Format the address. Returns an internal buffer.
   */
  StringPiece format(uintptr_t address);

 private:
  static constexpr char bufTemplate[] = "    @ 0000000000000000";
  char buf_[sizeof(bufTemplate)];
};

/**
 * Print a list of symbolized addresses. Base class.
 */
class SymbolizePrinter {
 public:
  /**
   * Print one address, no ending newline.
   */
  void print(uintptr_t address, const SymbolizedFrame& frame);

  /**
   * Print one address with ending newline.
   */
  void println(uintptr_t address, const SymbolizedFrame& frame);

  /**
   * Print multiple addresses on separate lines.
   */
  void println(const uintptr_t* addresses,
               const SymbolizedFrame* frames,
               size_t frameCount);

  /**
   * Print a string, no endling newline.
   */
  void print(StringPiece sp) { doPrint(sp); }

  /**
   * Print multiple addresses on separate lines, skipping the first
   * skip addresses.
   */
  template <size_t N>
  void println(const FrameArray<N>& fa, size_t skip=0) {
    if (skip < fa.frameCount) {
      println(fa.addresses + skip, fa.frames + skip, fa.frameCount - skip);
    }
  }

  virtual ~SymbolizePrinter() { }

  enum Options {
    // Skip file and line information
    NO_FILE_AND_LINE = 1 << 0,

    // As terse as it gets: function name if found, address otherwise
    TERSE = 1 << 1,

    // Always colorize output (ANSI escape code)
    COLOR = 1 << 2,

    // Colorize output only if output is printed to a TTY (ANSI escape code)
    COLOR_IF_TTY = 1 << 3,

    // Skip frame address information
    NO_FRAME_ADDRESS = 1 << 4,
  };

  // NOTE: enum values used as indexes in kColorMap.
  enum Color { DEFAULT, RED, GREEN, YELLOW, BLUE, CYAN, WHITE, PURPLE, NUM };
  void color(Color c);

 protected:
  explicit SymbolizePrinter(int options, bool isTty = false)
    : options_(options),
      isTty_(isTty) {
  }

  const int options_;
  const bool isTty_;

 private:
  void printTerse(uintptr_t address, const SymbolizedFrame& frame);
  virtual void doPrint(StringPiece sp) = 0;

  static constexpr std::array<const char*, Color::NUM> kColorMap = {{
      "\x1B[0m",
      "\x1B[31m",
      "\x1B[32m",
      "\x1B[33m",
      "\x1B[34m",
      "\x1B[36m",
      "\x1B[37m",
      "\x1B[35m",
  }};
};

/**
 * Print a list of symbolized addresses to a stream.
 * Not reentrant. Do not use from signal handling code.
 */
class OStreamSymbolizePrinter : public SymbolizePrinter {
 public:
  explicit OStreamSymbolizePrinter(std::ostream& out, int options=0);
 private:
  void doPrint(StringPiece sp) override;
  std::ostream& out_;
};

/**
 * Print a list of symbolized addresses to a file descriptor.
 * Ignores errors. Async-signal-safe.
 */
class FDSymbolizePrinter : public SymbolizePrinter {
 public:
  explicit FDSymbolizePrinter(int fd, int options=0,
                              size_t bufferSize=0);
  ~FDSymbolizePrinter() override;
  void flush();
 private:
  void doPrint(StringPiece sp) override;

  const int fd_;
  std::unique_ptr<IOBuf> buffer_;
};

/**
 * Print a list of symbolized addresses to a FILE*.
 * Ignores errors. Not reentrant. Do not use from signal handling code.
 */
class FILESymbolizePrinter : public SymbolizePrinter {
 public:
  explicit FILESymbolizePrinter(FILE* file, int options=0);
 private:
  void doPrint(StringPiece sp) override;
  FILE* const file_ = nullptr;
};

/**
 * Print a list of symbolized addresses to a std::string.
 * Not reentrant. Do not use from signal handling code.
 */
class StringSymbolizePrinter : public SymbolizePrinter {
 public:
  explicit StringSymbolizePrinter(int options=0) : SymbolizePrinter(options) { }

  std::string str() const { return buf_.toStdString(); }
  const fbstring& fbstr() const { return buf_; }
  fbstring moveFbString() { return std::move(buf_); }

 private:
  void doPrint(StringPiece sp) override;
  fbstring buf_;
};

/**
 * Use this class to print a stack trace from a signal handler, or other place
 * where you shouldn't allocate memory on the heap, and fsync()ing your file
 * descriptor is more important than performance.
 *
 * Make sure to create one of these on startup, not in the signal handler, as
 * the constructo allocates on the heap, whereas the other methods don't.  Best
 * practice is to just leak this object, rather than worry about destruction
 * order.
 *
 * These methods aren't thread safe, so if you could have signals on multiple
 * threads at the same time, you need to do your own locking to ensure you don't
 * call these methods from multiple threads.  They are signal safe, however.
 */
class StackTracePrinter {
 public:
  static constexpr size_t kDefaultMinSignalSafeElfCacheSize = 500;

  explicit StackTracePrinter(
      size_t minSignalSafeElfCacheSize = kDefaultMinSignalSafeElfCacheSize,
      int fd = STDERR_FILENO);

  /**
   * Only allocates on the stack and is signal-safe but not thread-safe.  Don't
   * call printStackTrace() on the same StackTracePrinter object from multiple
   * threads at the same time.
   */
  FOLLY_NOINLINE void printStackTrace(bool symbolize);

  void print(StringPiece sp) {
    printer_.print(sp);
  }

  // Flush printer_, also fsync, in case we're about to crash again...
  void flush();

 private:
  static constexpr size_t kMaxStackTraceDepth = 100;

  int fd_;
  SignalSafeElfCache elfCache_;
  FDSymbolizePrinter printer_;
  std::unique_ptr<FrameArray<kMaxStackTraceDepth>> addresses_;
};

}  // namespace symbolizer
}  // namespace folly
