#include "CompilerLogger.h"

#include <algorithm>
#include <fstream>
#include <iostream>
#include <set>
#include <vector>

#include "IRMutator.h"
#include "Util.h"

namespace Halide {
namespace Internal {

namespace {

// TODO: for now we are just going to ignore potential issues with
// static-initialization-order-fiasco, as CompilerLogger isn't currently used
// from any static-initialization execution scope.
std::unique_ptr<CompilerLogger> active_compiler_logger;

class ObfuscateNames : public IRMutator {
 using IRMutator::visit;

 std::map<std::string, std::string> remapping;

 Expr visit(const Call *op) override {
 auto args = mutate(op->args);
 std::string name = op->name;
 if (op->call_type == Call::Extern ||
 op->call_type == Call::ExternCPlusPlus ||
 op->call_type == Call::Halide ||
 op->call_type == Call::Image) {
 name = remap(op->name);
 }
 return Call::make(op->type, name, args, op->call_type, op->func,
 op->value_index, op->image, op->param);
 }

 Expr visit(const Let *op) override {
 std::string name = remap(op->name);
 Expr value = mutate(op->value);
 Expr body = mutate(op->body);
 return Let::make(name, std::move(value), std::move(body));
 }

 Expr visit(const Load *op) override {
 std::string name = remap(op->name);
 Expr index = mutate(op->index);
 Expr predicate = mutate(op->predicate);
 return Load::make(op->type, name, index, op->image, op->param,
 predicate, op->alignment);
 }

 Expr visit(const Variable *op) override {
 std::string name = remap(op->name);
 return Variable::make(op->type, name, op->image, op->param, op->reduction_domain);
 }

public:
 ObfuscateNames() = default;
 ObfuscateNames(const std::initializer_list<std::pair<const std::string, std::string>> &values)
 : remapping(values) {
 }

 std::string remap(const std::string &var_name) {
 std::string anon_name = "anon" + std::to_string(remapping.size());
 auto p = remapping.insert({var_name, std::move(anon_name)});
 return p.first->second;
 }
};

} // namespace

std::unique_ptr<CompilerLogger> set_compiler_logger(std::unique_ptr<CompilerLogger> compiler_logger) {
 std::unique_ptr<CompilerLogger> result = std::move(active_compiler_logger);
 active_compiler_logger = std::move(compiler_logger);
 return result;
}

CompilerLogger *get_compiler_logger() {
 return active_compiler_logger.get();
}

JSONCompilerLogger::JSONCompilerLogger(
 const std::string &generator_name,
 const std::string &function_name,
 const std::string &autoscheduler_name,
 const Target &target,
 const std::string &generator_args,
 bool obfuscate_exprs)
 : generator_name(generator_name),
 function_name(function_name),
 autoscheduler_name(autoscheduler_name),
 target(target),
 generator_args(generator_args),
 obfuscate_exprs(obfuscate_exprs) {
}

void JSONCompilerLogger::record_matched_simplifier_rule(const std::string &rulename, Expr expr) {
 matched_simplifier_rules[rulename].emplace_back(std::move(expr));
}

void JSONCompilerLogger::record_non_monotonic_loop_var(const std::string &loop_var, Expr expr) {
 non_monotonic_loop_vars[loop_var].emplace_back(std::move(expr));
}
void JSONCompilerLogger::record_failed_to_prove(Expr failed_to_prove, Expr original_expr) {
 failed_to_prove_exprs.emplace_back(failed_to_prove, original_expr);
}

void JSONCompilerLogger::record_object_code_size(uint64_t bytes) {
 object_code_size += bytes;
}

void JSONCompilerLogger::record_compilation_time(Phase phase, double duration) {
 compilation_time[phase] += duration;
}

void JSONCompilerLogger::obfuscate() {
 {
 std::map<std::string, std::vector<Expr>> n;
 for (const auto &it : matched_simplifier_rules) {
 std::string rule = it.first;
 for (const auto &e : it.second) {
 ObfuscateNames obfuscater;
 n[rule].emplace_back(obfuscater.mutate(e));
 }
 }
 matched_simplifier_rules = n;
 }
 {
 std::vector<std::pair<Expr, Expr>> n;
 for (const auto &it : failed_to_prove_exprs) {
 // Note that use a separate obfuscater for each pair, so each
 // shares identifiers only with each other; this makes it simpler
 // to post-process output from multiple unrelated Generators
 // and combine Exprs with similar shapes.
 ObfuscateNames obfuscater;
 auto failed_to_prove = obfuscater.mutate(it.first);
 auto original_expr = obfuscater.mutate(it.second);
 n.emplace_back(std::move(failed_to_prove), std::move(original_expr));
 }
 failed_to_prove_exprs = n;
 }
}

namespace {

std::ostream &emit_eol(std::ostream &o, bool comma = true) {
 o << (comma ? ",\n" : "\n");
 return o;
}

std::ostream &emit_key(std::ostream &o, int indent, const std::string &key) {
 std::string spaces(indent, ' ');
 o << spaces << "\"" << key << "\" : ";
 return o;
}

std::ostream &emit_object_key_open(std::ostream &o, int indent, const std::string &key) {
 std::string spaces(indent, ' ');
 o << spaces << "\"" << key << "\" : {\n";
 return o;
}

std::ostream &emit_object_key_close(std::ostream &o, int indent, bool comma = true) {
 std::string spaces(indent, ' ');
 o << spaces << "}";
 emit_eol(o, comma);
 return o;
}

template<typename VALUE>
std::ostream &emit_value(std::ostream &o, const VALUE &value) {
 o << value;
 return o;
}

template<>
std::ostream &emit_value<std::string>(std::ostream &o, const std::string &value) {
 std::string v = value;
 v = replace_all(v, "\\", "\\\\");
 v = replace_all(v, "\"", "\\\"");
 o << "\"" << v << "\"";
 return o;
}

template<typename VALUE>
std::ostream &emit_key_value(std::ostream &o, int indent, const std::string &key, const VALUE &value, bool comma = true) {
 emit_key(o, indent, key);
 emit_value(o, value);
 emit_eol(o, comma);
 return o;
}

std::ostream &emit_optional_key_value(std::ostream &o, int indent, const std::string &key, const std::string &value, bool comma = true) {
 if (!value.empty()) {
 return emit_key_value(o, indent, key, value, comma);
 }
 return o;
}

template<typename T>
std::ostream &emit_pairs(std::ostream &o, int indent, const std::string &key, const T &pairs, bool comma = true) {
 std::string spaces(indent, ' ');

 emit_key(o, indent, key);
 o << "{\n";
 int commas_to_emit = (int)pairs.size() - 1;
 for (const auto &it : pairs) {
 const bool comma = (commas_to_emit > 0);
 emit_key_value(o, indent + 1, it.first, it.second, comma);
 commas_to_emit--;
 }
 o << spaces << "}";
 emit_eol(o, comma);
 return o;
}

template<typename T>
std::ostream &emit_list(std::ostream &o, int indent, const T &list, bool comma = true) {
 std::string spaces(indent, ' ');
 std::string spaces_in(indent + 1, ' ');

 o << spaces << "[\n";
 int commas_to_emit = (int)list.size() - 1;
 for (const auto &it : list) {
 o << spaces_in;
 emit_value(o, it);
 emit_eol(o, commas_to_emit-- > 0);
 }
 o << spaces << "]";
 emit_eol(o, comma);
 return o;
}

std::string expr_to_string(const Expr &e) {
 std::ostringstream s;
 s << e;
 return s.str();
}

std::set<std::string> exprs_to_strings(const std::vector<Expr> &exprs) {
 std::set<std::string> strings;
 for (const auto &e : exprs) {
 strings.insert(expr_to_string(e));
 }
 return strings;
}

} // namespace

std::ostream &JSONCompilerLogger::emit_to_stream(std::ostream &o) {
 if (obfuscate_exprs) {
 obfuscate();
 }

 // Output in JSON form

 o << "{\n";

 int indent = 1;
 emit_optional_key_value(o, indent, "generator_name", generator_name);
 emit_optional_key_value(o, indent, "function_name", function_name);
 emit_optional_key_value(o, indent, "autoscheduler_name", autoscheduler_name);
 emit_optional_key_value(o, indent, "target", target == Target() ? "" : target.to_string());
 emit_optional_key_value(o, indent, "generator_args", generator_args);

 if (object_code_size) {
 emit_key_value(o, indent, "object_code_size", object_code_size);
 }

 // If these are present, emit them, even if value is zero
 if (compilation_time.count(Phase::HalideLowering)) {
 emit_key_value(o, indent, "compilation_time_halide_lowering", compilation_time[Phase::HalideLowering]);
 }
 if (compilation_time.count(Phase::LLVM)) {
 emit_key_value(o, indent, "compilation_time_llvm", compilation_time[Phase::LLVM]);
 }

 if (!matched_simplifier_rules.empty()) {
 emit_object_key_open(o, indent, "matched_simplifier_rules");

 int commas_to_emit = (int)matched_simplifier_rules.size() - 1;
 for (const auto &it : matched_simplifier_rules) {
 const auto &loop_var = it.first;
 emit_key(o, indent + 1, loop_var);
 emit_eol(o, false);
 emit_list(o, indent + 1, exprs_to_strings(it.second), (commas_to_emit-- > 0));
 }

 emit_object_key_close(o, indent);
 }

 if (!non_monotonic_loop_vars.empty()) {
 emit_object_key_open(o, indent, "non_monotonic_loop_vars");

 int commas_to_emit = (int)non_monotonic_loop_vars.size() - 1;
 for (const auto &it : non_monotonic_loop_vars) {
 const auto &loop_var = it.first;
 emit_key(o, indent + 1, loop_var);
 emit_eol(o, false);
 emit_list(o, indent + 1, exprs_to_strings(it.second), (commas_to_emit-- > 0));
 }

 emit_object_key_close(o, indent);
 }

 if (!failed_to_prove_exprs.empty()) {
 emit_object_key_open(o, indent, "failed_to_prove");

 // We'll do deduplication here, during stringification.
 std::map<std::string, std::set<std::string>> sorted;
 for (const auto &it : failed_to_prove_exprs) {
 const auto failed_to_prove_str = expr_to_string(it.first);
 const auto original_expr_str = expr_to_string(it.second);
 sorted[failed_to_prove_str].insert(original_expr_str);
 }

 int commas_to_emit = (int)sorted.size() - 1;
 for (const auto &it : sorted) {
 emit_key(o, indent + 1, it.first);
 emit_eol(o, false);
 emit_list(o, indent + 1, it.second, (commas_to_emit-- > 0));
 }

 emit_object_key_close(o, indent);
 }

 // Emit this last as a simple way to dodge the trailing-comma nonsense
 o << " \"version\": \"HalideJSONCompilerLoggerV1\"\n";
 o << "}\n";

 return o;
}

} // namespace Internal
} // namespace Halide