/*
 * Copyright 2016 WebAssembly Community Group participants
 *
 * 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.
 */

#include <algorithm>
#include <iomanip>
#include <pass.h>
#include <support/colors.h>
#include <wasm.h>
#include <wasm-binary.h>

namespace wasm {

using namespace std;

// Prints metrics between optimization passes.
struct Metrics : public WalkerPass<PostWalker<Metrics, UnifiedExpressionVisitor<Metrics>>> {
  bool byFunction;

  Metrics(bool byFunction) : byFunction(byFunction) {}

  static Metrics *lastMetricsPass;

  map<const char *, int> counts;

  void visitExpression(Expression* curr) {
    auto name = getExpressionName(curr);
    counts[name]++;
  }

  void doWalkModule(Module* module) {
    // global things

    for (auto& curr : module->functionTypes) {
      visitFunctionType(curr.get());
    }
    for (auto& curr : module->imports) {
      visitImport(curr.get());
    }
    for (auto& curr : module->exports) {
      visitExport(curr.get());
    }
    for (auto& curr : module->globals) {
      walkGlobal(curr.get());
    }
    walkTable(&module->table);
    walkMemory(&module->memory);

    // add functions
    counts["[funcs]"] = module->functions.size();
    // add memory and table
    if (module->memory.exists) {
      Index size = 0;
      for (auto& segment: module->memory.segments) {
        size += segment.data.size();
      }
      counts["[memory-data]"] = size;
    }
    if (module->table.exists) {
      Index size = 0;
      for (auto& segment: module->table.segments) {
        size += segment.data.size();
      }
      counts["[table-data]"] = size;
    }

    if (byFunction) {
      // print global
      printCounts("global");
      // compute binary info, so we know function sizes
      BufferWithRandomAccess buffer;
      WasmBinaryWriter writer(module, buffer);
      writer.write();
      // print for each function
      for (Index i = 0; i < module->functions.size(); i++) {
        auto* func = module->functions[i].get();
        counts.clear();
        walkFunction(func);
        counts["[vars]"] = func->getNumVars();
        counts["[binary-bytes]"] = writer.tableOfContents.functionBodies[i].size;
        printCounts(std::string("func: ") + func->name.str);
      }
      // can't comapre detailed info between passes yet
      lastMetricsPass = nullptr;
    } else {
      // add function info
      size_t vars = 0;
      for (auto& func : module->functions) {
        walkFunction(func.get());
        vars += func->getNumVars();
      }
      counts["[vars]"] = vars;
      // print
      printCounts("total");
      // compare to next time
      lastMetricsPass = this;
    }
  }

  void printCounts(std::string title) {
    ostream &o = cout;
    vector<const char*> keys;
    // add total
    int total = 0;
    for (auto i : counts) {
      keys.push_back(i.first);
      total += i.second;
    }
    keys.push_back("[total]");
    counts["[total]"] = total;
    // sort
    sort(keys.begin(), keys.end(), [](const char* a, const char* b) -> bool {
      return strcmp(b, a) > 0;
    });
    o << title << "\n";
    for (auto* key : keys) {
      auto value = counts[key];
      o << " " << left << setw(15) << key << ": " << setw(8)
        << value;
      if (lastMetricsPass) {
        if (lastMetricsPass->counts.count(key)) {
          int before = lastMetricsPass->counts[key];
          int after = value;
          if (after - before) {
            if (after > before) {
              Colors::red(o);
            } else {
              Colors::green(o);
            }
            o << right << setw(8);
            o << showpos << after - before << noshowpos;
            Colors::normal(o);
          }
        }
      }
      o << "\n";
    }
  }
};

Pass *createMetricsPass() {
  return new Metrics(false);
}

Pass *createFunctionMetricsPass() {
  return new Metrics(true);
}

Metrics *Metrics::lastMetricsPass;

} // namespace wasm