// Copyright 2017 Google 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.

// This program dumps current host data to the standard output.
// Output can be text or json if the `--json` flag is passed.

#include <assert.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "cpu_features_macros.h"

#if defined(CPU_FEATURES_ARCH_X86)
#include "cpuinfo_x86.h"
#elif defined(CPU_FEATURES_ARCH_ARM)
#include "cpuinfo_arm.h"
#elif defined(CPU_FEATURES_ARCH_AARCH64)
#include "cpuinfo_aarch64.h"
#elif defined(CPU_FEATURES_ARCH_MIPS)
#include "cpuinfo_mips.h"
#elif defined(CPU_FEATURES_ARCH_PPC)
#include "cpuinfo_ppc.h"
#endif

// Design principles
// -----------------
// We build a tree structure containing all the data to be displayed.
// Then depending on the output type (text or json) we walk the tree and display
// the data accordingly.

// We use a bump allocator to allocate strings and nodes of the tree,
// Memory is not intented to be reclaimed.
typedef struct {
  char* ptr;
  size_t size;
} BumpAllocator;

__attribute__((aligned(8))) char gGlobalBuffer[64 * 1024];
BumpAllocator gBumpAllocator = {.ptr = gGlobalBuffer,
                                .size = sizeof(gGlobalBuffer)};

static void internal_error() {
  fputs("internal error\n", stderr);
  exit(EXIT_FAILURE);
}

// Update the available memory left in the BumpAllocator.
static void* BA_Bump(size_t size) {
  size = (size + 8 - 1) / 8 * 8;  // Align size to next 8B boundary.
  if (gBumpAllocator.size < size) internal_error();
  void* ptr = gBumpAllocator.ptr;
  gBumpAllocator.size -= size;
  gBumpAllocator.ptr += size;
  assert((uintptr_t)(ptr) % 8 == 0 && "memory must be aligned");
  return ptr;
}

// The type of the nodes in the tree.
typedef enum {
  NT_INVALID,
  NT_INT,
  NT_MAP,
  NT_MAP_ENTRY,
  NT_ARRAY,
  NT_ARRAY_ELEMENT,
  NT_STRING,
} NodeType;

// The node in the tree.
typedef struct Node {
  NodeType type;
  unsigned integer;
  const char* string;
  struct Node* value;
  struct Node* next;
} Node;

// Creates an initialized Node.
static Node* BA_CreateNode(NodeType type) {
  Node* tv = (Node*)BA_Bump(sizeof(Node));
  assert(tv);
  *tv = (Node){.type = type};
  return tv;
}

// Adds an integer node.
static Node* CreateInt(int value) {
  Node* tv = BA_CreateNode(NT_INT);
  tv->integer = value;
  return tv;
}

// Adds a string node.
// `value` must outlive the tree.
static Node* CreateConstantString(const char* value) {
  Node* tv = BA_CreateNode(NT_STRING);
  tv->string = value;
  return tv;
}

// Adds a map node.
static Node* CreateMap() { return BA_CreateNode(NT_MAP); }

// Adds an array node.
static Node* CreateArray() { return BA_CreateNode(NT_ARRAY); }

// Adds a formatted string node.
static Node* CreatePrintfString(const char* format, ...) {
  va_list arglist;
  va_start(arglist, format);
  char* const ptr = gBumpAllocator.ptr;
  const int written = vsnprintf(ptr, gBumpAllocator.size, format, arglist);
  va_end(arglist);
  if (written < 0 || written >= gBumpAllocator.size) internal_error();
  return CreateConstantString((char*)BA_Bump(written));
}

// Adds a string node.
static Node* CreateString(const char* value) {
  return CreatePrintfString("%s", value);
}

// Adds a map entry node.
static void AddMapEntry(Node* map, const char* key, Node* value) {
  assert(map && map->type == NT_MAP);
  Node* current = map;
  while (current->next) current = current->next;
  current->next = (Node*)BA_Bump(sizeof(Node));
  *current->next = (Node){.type = NT_MAP_ENTRY, .string = key, .value = value};
}

// Adds an array element node.
static void AddArrayElement(Node* array, Node* value) {
  assert(array && array->type == NT_ARRAY);
  Node* current = array;
  while (current->next) current = current->next;
  current->next = (Node*)BA_Bump(sizeof(Node));
  *current->next = (Node){.type = NT_ARRAY_ELEMENT, .value = value};
}

static int cmp(const void* p1, const void* p2) {
  return strcmp(*(const char* const*)p1, *(const char* const*)p2);
}

#define DEFINE_ADD_FLAGS(HasFeature, FeatureName, FeatureType, LastEnum) \
  static void AddFlags(Node* map, const FeatureType* features) {         \
    size_t i;                                                            \
    const char* ptrs[LastEnum] = {0};                                    \
    size_t count = 0;                                                    \
    for (i = 0; i < LastEnum; ++i) {                                     \
      if (HasFeature(features, i)) {                                     \
        ptrs[count] = FeatureName(i);                                    \
        ++count;                                                         \
      }                                                                  \
    }                                                                    \
    qsort((void*)ptrs, count, sizeof(char*), cmp);                       \
    Node* const array = CreateArray();                                   \
    for (i = 0; i < count; ++i)                                          \
      AddArrayElement(array, CreateConstantString(ptrs[i]));             \
    AddMapEntry(map, "flags", array);                                    \
  }

#if defined(CPU_FEATURES_ARCH_X86)
DEFINE_ADD_FLAGS(GetX86FeaturesEnumValue, GetX86FeaturesEnumName, X86Features,
                 X86_LAST_)
#elif defined(CPU_FEATURES_ARCH_ARM)
DEFINE_ADD_FLAGS(GetArmFeaturesEnumValue, GetArmFeaturesEnumName, ArmFeatures,
                 ARM_LAST_)
#elif defined(CPU_FEATURES_ARCH_AARCH64)
DEFINE_ADD_FLAGS(GetAarch64FeaturesEnumValue, GetAarch64FeaturesEnumName,
                 Aarch64Features, AARCH64_LAST_)
#elif defined(CPU_FEATURES_ARCH_MIPS)
DEFINE_ADD_FLAGS(GetMipsFeaturesEnumValue, GetMipsFeaturesEnumName,
                 MipsFeatures, MIPS_LAST_)
#elif defined(CPU_FEATURES_ARCH_PPC)
DEFINE_ADD_FLAGS(GetPPCFeaturesEnumValue, GetPPCFeaturesEnumName, PPCFeatures,
                 PPC_LAST_)
#endif

// Prints a json string with characters escaping.
static void printJsonString(const char* str) {
  putchar('"');
  for (; str && *str; ++str) {
    switch (*str) {
      case '\"':
      case '\\':
      case '/':
      case '\b':
      case '\f':
      case '\n':
      case '\r':
      case '\t':
        putchar('\\');
    }
    putchar(*str);
  }
  putchar('"');
}

// Walks a Node and print it as json.
static void printJson(const Node* current) {
  assert(current);
  switch (current->type) {
    case NT_INT:
      printf("%d", current->integer);
      break;
    case NT_STRING:
      printJsonString(current->string);
      break;
    case NT_ARRAY:
      putchar('[');
      if (current->next) printJson(current->next);
      putchar(']');
      break;
    case NT_MAP:
      putchar('{');
      if (current->next) printJson(current->next);
      putchar('}');
      break;
    case NT_MAP_ENTRY:
      printf("\"%s\":", current->string);
      printJson(current->value);
      if (current->next) {
        putchar(',');
        printJson(current->next);
      }
      break;
    case NT_ARRAY_ELEMENT:
      printJson(current->value);
      if (current->next) {
        putchar(',');
        printJson(current->next);
      }
      break;
  }
}

// Walks a Node and print it as text.
static void printTextField(const Node* current) {
  switch (current->type) {
    case NT_INT:
      printf("%3d (0x%02X)", current->integer, current->integer);
      break;
    case NT_STRING:
      fputs(current->string, stdout);
      break;
    case NT_ARRAY:
      if (current->next) printTextField(current->next);
      break;
    case NT_MAP:
      if (current->next) printJson(current->next);
      break;
    case NT_MAP_ENTRY:
      printf("%-15s : ", current->string);
      printTextField(current->value);
      if (current->next) {
        putchar('\n');
        printTextField(current->next);
      }
      break;
    case NT_ARRAY_ELEMENT:
      printTextField(current->value);
      if (current->next) {
        putchar(',');
        printTextField(current->next);
      }
      break;
  }
}

static void printTextRoot(const Node* current) {
  if (current->type == NT_MAP && current->next) printTextField(current->next);
}

static void showUsage(const char* name) {
  printf(
      "\n"
      "Usage: %s [options]\n"
      "      Options:\n"
      "      -h | --help     Show help message.\n"
      "      -j | --json     Format output as json instead of plain text.\n"
      "\n",
      name);
}

static Node* CreateTree() {
  Node* root = CreateMap(gBumpAllocator);
#if defined(CPU_FEATURES_ARCH_X86)
  char brand_string[49];
  const X86Info info = GetX86Info();
  FillX86BrandString(brand_string);
  AddMapEntry(root, "arch", CreateString("x86"));
  AddMapEntry(root, "brand", CreateString(brand_string));
  AddMapEntry(root, "family", CreateInt(info.family));
  AddMapEntry(root, "model", CreateInt(info.model));
  AddMapEntry(root, "stepping", CreateInt(info.stepping));
  AddMapEntry(root, "uarch",
              CreateString(
                  GetX86MicroarchitectureName(GetX86Microarchitecture(&info))));
  AddFlags(root, &info.features);
#elif defined(CPU_FEATURES_ARCH_ARM)
  const ArmInfo info = GetArmInfo();
  AddMapEntry(root, "arch", CreateString("ARM"));
  AddMapEntry(root, "implementer", CreateInt(info.implementer));
  AddMapEntry(root, "architecture", CreateInt(info.architecture));
  AddMapEntry(root, "variant", CreateInt(info.variant));
  AddMapEntry(root, "part", CreateInt(info.part));
  AddMapEntry(root, "revision", CreateInt(info.revision));
  AddFlags(root, &info.features);
#elif defined(CPU_FEATURES_ARCH_AARCH64)
  const Aarch64Info info = GetAarch64Info();
  AddMapEntry(root, "arch", CreateString("aarch64"));
  AddMapEntry(root, "implementer", CreateInt(info.implementer));
  AddMapEntry(root, "variant", CreateInt(info.variant));
  AddMapEntry(root, "part", CreateInt(info.part));
  AddMapEntry(root, "revision", CreateInt(info.revision));
  AddFlags(root, &info.features);
#elif defined(CPU_FEATURES_ARCH_MIPS)
  const MipsInfo info = GetMipsInfo();
  AddMapEntry(root, "arch", CreateString("mips"));
  AddFlags(root, &info.features);
#elif defined(CPU_FEATURES_ARCH_PPC)
  const PPCInfo info = GetPPCInfo();
  const PPCPlatformStrings strings = GetPPCPlatformStrings();
  AddMapEntry(root, "arch", CreateString("ppc"));
  AddMapEntry(root, "platform", CreateString(strings.platform));
  AddMapEntry(root, "model", CreateString(strings.model));
  AddMapEntry(root, "machine", CreateString(strings.machine));
  AddMapEntry(root, "cpu", CreateString(strings.cpu));
  AddMapEntry(root, "instruction", CreateString(strings.type.platform));
  AddMapEntry(root, "microarchitecture",
              CreateString(strings.type.base_platform));
  AddFlags(root, &info.features);
#endif
  return root;
}

int main(int argc, char** argv) {
  const Node* const root = CreateTree(&gBumpAllocator);
  bool outputJson = false;
  int i = 1;
  for (; i < argc; ++i) {
    const char* arg = argv[i];
    if (strcmp(arg, "-j") == 0 || strcmp(arg, "--json") == 0) {
      outputJson = true;
    } else {
      showUsage(argv[0]);
      if (strcmp(arg, "-h") == 0 || strcmp(arg, "--help") == 0)
        return EXIT_SUCCESS;
      return EXIT_FAILURE;
    }
  }
  if (outputJson)
    printJson(root);
  else
    printTextRoot(root);
  putchar('\n');
  return EXIT_SUCCESS;
}