// TODO: header template

#include "clang/AST/OSLog.h"
#include "clang/AST/Attr.h"
#include "clang/AST/Decl.h"
#include "clang/AST/DeclCXX.h"
#include "clang/AST/ExprObjC.h"
#include "clang/AST/FormatString.h"
#include "clang/Basic/Builtins.h"
#include "llvm/ADT/SmallBitVector.h"
#include <optional>

using namespace clang;

using clang::analyze_os_log::OSLogBufferItem;
using clang::analyze_os_log::OSLogBufferLayout;

namespace {
class OSLogFormatStringHandler
    : public analyze_format_string::FormatStringHandler {
private:
  struct ArgData {
    const Expr *E = nullptr;
    std::optional<OSLogBufferItem::Kind> Kind;
    std::optional<unsigned> Size;
    std::optional<const Expr *> Count;
    std::optional<const Expr *> Precision;
    std::optional<const Expr *> FieldWidth;
    unsigned char Flags = 0;
    StringRef MaskType;
  };
  SmallVector<ArgData, 4> ArgsData;
  ArrayRef<const Expr *> Args;

  OSLogBufferItem::Kind
  getKind(analyze_format_string::ConversionSpecifier::Kind K) {
    switch (K) {
    case clang::analyze_format_string::ConversionSpecifier::sArg: // "%s"
      return OSLogBufferItem::StringKind;
    case clang::analyze_format_string::ConversionSpecifier::SArg: // "%S"
      return OSLogBufferItem::WideStringKind;
    case clang::analyze_format_string::ConversionSpecifier::PArg: { // "%P"
      return OSLogBufferItem::PointerKind;
    case clang::analyze_format_string::ConversionSpecifier::ObjCObjArg: // "%@"
      return OSLogBufferItem::ObjCObjKind;
    case clang::analyze_format_string::ConversionSpecifier::PrintErrno: // "%m"
      return OSLogBufferItem::ErrnoKind;
    default:
      return OSLogBufferItem::ScalarKind;
    }
    }
  }

public:
  OSLogFormatStringHandler(ArrayRef<const Expr *> Args) : Args(Args) {
    ArgsData.reserve(Args.size());
  }

  bool HandlePrintfSpecifier(const analyze_printf::PrintfSpecifier &FS,
                             const char *StartSpecifier, unsigned SpecifierLen,
                             const TargetInfo &) override {
    if (!FS.consumesDataArgument() &&
        FS.getConversionSpecifier().getKind() !=
            clang::analyze_format_string::ConversionSpecifier::PrintErrno)
      return true;

    ArgsData.emplace_back();
    unsigned ArgIndex = FS.getArgIndex();
    if (ArgIndex < Args.size())
      ArgsData.back().E = Args[ArgIndex];

    // First get the Kind
    ArgsData.back().Kind = getKind(FS.getConversionSpecifier().getKind());
    if (ArgsData.back().Kind != OSLogBufferItem::ErrnoKind &&
        !ArgsData.back().E) {
      // missing argument
      ArgsData.pop_back();
      return false;
    }

    switch (FS.getConversionSpecifier().getKind()) {
    case clang::analyze_format_string::ConversionSpecifier::sArg:   // "%s"
    case clang::analyze_format_string::ConversionSpecifier::SArg: { // "%S"
      auto &precision = FS.getPrecision();
      switch (precision.getHowSpecified()) {
      case clang::analyze_format_string::OptionalAmount::NotSpecified: // "%s"
        break;
      case clang::analyze_format_string::OptionalAmount::Constant: // "%.16s"
        ArgsData.back().Size = precision.getConstantAmount();
        break;
      case clang::analyze_format_string::OptionalAmount::Arg: // "%.*s"
        ArgsData.back().Count = Args[precision.getArgIndex()];
        break;
      case clang::analyze_format_string::OptionalAmount::Invalid:
        return false;
      }
      break;
    }
    case clang::analyze_format_string::ConversionSpecifier::PArg: { // "%P"
      auto &precision = FS.getPrecision();
      switch (precision.getHowSpecified()) {
      case clang::analyze_format_string::OptionalAmount::NotSpecified: // "%P"
        return false; // length must be supplied with pointer format specifier
      case clang::analyze_format_string::OptionalAmount::Constant: // "%.16P"
        ArgsData.back().Size = precision.getConstantAmount();
        break;
      case clang::analyze_format_string::OptionalAmount::Arg: // "%.*P"
        ArgsData.back().Count = Args[precision.getArgIndex()];
        break;
      case clang::analyze_format_string::OptionalAmount::Invalid:
        return false;
      }
      break;
    }
    default:
      if (FS.getPrecision().hasDataArgument()) {
        ArgsData.back().Precision = Args[FS.getPrecision().getArgIndex()];
      }
      break;
    }
    if (FS.getFieldWidth().hasDataArgument()) {
      ArgsData.back().FieldWidth = Args[FS.getFieldWidth().getArgIndex()];
    }

    if (FS.isSensitive())
      ArgsData.back().Flags |= OSLogBufferItem::IsSensitive;
    else if (FS.isPrivate())
      ArgsData.back().Flags |= OSLogBufferItem::IsPrivate;
    else if (FS.isPublic())
      ArgsData.back().Flags |= OSLogBufferItem::IsPublic;

    ArgsData.back().MaskType = FS.getMaskType();
    return true;
  }

  void computeLayout(ASTContext &Ctx, OSLogBufferLayout &Layout) const {
    Layout.Items.clear();
    for (auto &Data : ArgsData) {
      if (!Data.MaskType.empty()) {
        CharUnits Size = CharUnits::fromQuantity(8);
        Layout.Items.emplace_back(OSLogBufferItem::MaskKind, nullptr,
                                  Size, 0, Data.MaskType);
      }

      if (Data.FieldWidth) {
        CharUnits Size = Ctx.getTypeSizeInChars((*Data.FieldWidth)->getType());
        Layout.Items.emplace_back(OSLogBufferItem::ScalarKind, *Data.FieldWidth,
                                  Size, 0);
      }
      if (Data.Precision) {
        CharUnits Size = Ctx.getTypeSizeInChars((*Data.Precision)->getType());
        Layout.Items.emplace_back(OSLogBufferItem::ScalarKind, *Data.Precision,
                                  Size, 0);
      }
      if (Data.Count) {
        // "%.*P" has an extra "count" that we insert before the argument.
        CharUnits Size = Ctx.getTypeSizeInChars((*Data.Count)->getType());
        Layout.Items.emplace_back(OSLogBufferItem::CountKind, *Data.Count, Size,
                                  0);
      }
      if (Data.Size)
        Layout.Items.emplace_back(Ctx, CharUnits::fromQuantity(*Data.Size),
                                  Data.Flags);
      if (Data.Kind) {
        CharUnits Size;
        if (*Data.Kind == OSLogBufferItem::ErrnoKind)
          Size = CharUnits::Zero();
        else
          Size = Ctx.getTypeSizeInChars(Data.E->getType());
        Layout.Items.emplace_back(*Data.Kind, Data.E, Size, Data.Flags);
      } else {
        auto Size = Ctx.getTypeSizeInChars(Data.E->getType());
        Layout.Items.emplace_back(OSLogBufferItem::ScalarKind, Data.E, Size,
                                  Data.Flags);
      }
    }
  }
};
} // end anonymous namespace

bool clang::analyze_os_log::computeOSLogBufferLayout(
    ASTContext &Ctx, const CallExpr *E, OSLogBufferLayout &Layout) {
  ArrayRef<const Expr *> Args(E->getArgs(), E->getArgs() + E->getNumArgs());

  const Expr *StringArg;
  ArrayRef<const Expr *> VarArgs;
  switch (E->getBuiltinCallee()) {
  case Builtin::BI__builtin_os_log_format_buffer_size:
    assert(E->getNumArgs() >= 1 &&
           "__builtin_os_log_format_buffer_size takes at least 1 argument");
    StringArg = E->getArg(0);
    VarArgs = Args.slice(1);
    break;
  case Builtin::BI__builtin_os_log_format:
    assert(E->getNumArgs() >= 2 &&
           "__builtin_os_log_format takes at least 2 arguments");
    StringArg = E->getArg(1);
    VarArgs = Args.slice(2);
    break;
  default:
    llvm_unreachable("non-os_log builtin passed to computeOSLogBufferLayout");
  }

  const StringLiteral *Lit = cast<StringLiteral>(StringArg->IgnoreParenCasts());
  assert(Lit && (Lit->isOrdinary() || Lit->isUTF8()));
  StringRef Data = Lit->getString();
  OSLogFormatStringHandler H(VarArgs);
  ParsePrintfString(H, Data.begin(), Data.end(), Ctx.getLangOpts(),
                    Ctx.getTargetInfo(), /*isFreeBSDKPrintf*/ false);

  H.computeLayout(Ctx, Layout);
  return true;
}