#include "cpp.h"

#include "asm/intel64/encode.h"
#include "asm/operators.h"
#include "bnf.h"
#include "cppbnf.h"
#include "debug.h"
#include "flowgraph/node.h"
#include "lexer.h"
#include "grammer.h"
#include "minicc.h"

#include <gtest/gtest.h>
#include <gmock/gmock.h>

#include <boost/core/demangle.hpp>

#include <functional>
#include <optional>
#include <unordered_set>
#include <unordered_map>
#include <filesystem>

using namespace Gram;

namespace fs = std::filesystem;

CPP::CPP()
{
 node_eval_map = getNodeEvalMap();
}


CPP::~CPP(){}

// Phase 1: Map physical character set to basic source character set
void CPP::source_charset_map()
{
 // TODO
}

// Phase 2: Escape backslashed line endings
void CPP::backslash_escape()
{
 // TODO
}

// Phase 3: Parse preprocessing tokens, TODO: discard comments
std::vector<Token> CPP::preprocessing_tokenize(const std::string& s)
{
 auto bnf{SubBNF(CPPBNF::GetCppBNFLex(), "preprocessing-token")};
 
 Lex::Lexer lexer(bnf, "preprocessing-token");

 return lexer.Lex(s);
}

// Phase 4: Preprocessing
void CPP::preprocess()
{
 // TODO
}

// Phase 5: Map chars and strings to execution charset
void CPP::execution_charset_map()
{
 // TODO
}

// Phase 6: Concatenate adjacent string literals
void CPP::concatenate_strings()
{
 // TODO
}

std::string CPP::valueOfNode(index_t node_index) const
{
 std::string result;
 std::optional<size_t> pos0;
 index_t last_index;

 std::vector<int32_t> todo(1, int32_t(node_index));

 while (!todo.empty()) {
  int32_t current_index = todo.back();
  todo.pop_back();

  // visit node if token
  if (ChildIdIsToken(current_index)) {
   if (!pos0) {
    pos0 = m_tokens[TokenIdFromChildId(current_index)].location.pos;
   }
   last_index = TokenIdFromChildId(current_index);
  } else {

   const TreeNode &node{m_nodes[current_index]};

   // iterate backwards in childs, to get depth-first search in tree, from the beginning
   std::for_each(node.child_ids.rbegin(), node.child_ids.rend(), [&](int32_t child){
    todo.push_back(child);
   });
  }
 }

 if (!pos0)
  throw std::runtime_error("ICE: Node value not available");

 return m_source.substr(*pos0, m_tokens[last_index].location.pos - *pos0) + m_tokens[last_index].value;
};

std::string CPP::typeOfNode(index_t node_id) const
{
 if (node_id >= m_nodes.size())
  throw std::runtime_error("CPP::typeOfNode(): node_id="s + std::to_string(node_id) + ", m_nodes.size()="s + std::to_string(m_nodes.size()));
 return m_nodes[node_id].type;
}

namespace {
 std::unordered_set<std::string> pp_types {
  "identifier",
  "pp-number",
  "character-literal",
  "user-defined-character-literal",
  "string-literal",
  "user-defined-string-literal",
  "preprocessing-op-or-punc"
 };

 std::unordered_set<std::string> keywords {
  "alignas",
  "alignof",
  "asm",
  "auto",
  "bool",
  "break",
  "case",
  "catch",
  "char",
  "char8_t",
  "char16_t",
  "char32_t",
  "class",
  "concept",
  "const",
  "consteval",
  "constexpr",
  "constinit",
  "const_cast",
  "continue",
  "co_await",
  "co_return",
  "co_yield",
  "decltype",
  "default",
  "delete",
  "do",
  "double",
  "dynamic_cast",
  "else",
  "enum",
  "explicit",
  "export",
  "extern",
  "false",
  "float",
  "for",
  "friend",
  "goto",
  "if",
  "inline",
  "int",
  "long",
  "mutable",
  "namespace",
  "new",
  "noexcept",
  "nullptr",
  "operator",
  "private",
  "protected",
  "public",
  "register",
  "reinterpret_cast",
  "requires",
  "return",
  "short",
  "signed",
  "sizeof",
  "static",
  "static_assert",
  "static_cast",
  "struct",
  "switch",
  "template",
  "this",
  "thread_local",
  "throw",
  "true",
  "try",
  "typedef",
  "typeid",
  "typename",
  "union",
  "unsigned",
  "using",
  "virtual",
  "void",
  "volatile",
  "wchar_t",
  "while",
 };
}

// Phase 7.a: Create tokens from preprocessing tokens
std::vector<Token> CPP::tokens_from_pptokens(const std::vector<Token>& pp_tokens)
{
 std::vector<Token> result;

 // "identifier" + value -> "identifier" + value, except identifiers from table 5.11, p.14 -> keyword as value, value
 // "pp-number" + value -> "literal" + value
 // "character-literal" -> "literal" + value
 // "user-defined-character-literal" -> "literal" + value
 // "string-literal" -> "literal" + value
 // "user-defined-string-literal" -> "literal" + value
 // "preprocessing-op-or-punc" -> value+value (operator,punctuator)
 
 for (auto& token: pp_tokens) {
  if (pp_types.find(token.type) != pp_types.end()) {
   if (token.type == "identifier") {
    if (keywords.find(token.value) != keywords.end())
     result.emplace_back(Token{token.value, token.value, token.location});
    else
     result.emplace_back(Token{"identifier"s, token.value, token.location});
   }
   else if (token.type == "preprocessing-op-or-punc")
    result.emplace_back(Token{token.value, token.value, token.location});
   else
    result.emplace_back(Token{"literal", token.value, token.location});
  } else
   throw std::runtime_error("Unhandled preprocessing token: "s + token.toString());
 }
 return result;
}

// Phase 7.b: Grammar Analysis
std::vector<Gram::TreeNode> CPP::analysis(const std::vector<Token>& tokens)
{
 auto bnf = SubBNF(CPPBNF::GetCppBNFGram(), "translation-unit");

 Gram::Compiler compiler(bnf, "translation-unit");

 return compiler.compile(tokens);
}

std::string CPP::locationOfNode(index_t node_index) const
{
 if (node_index >= m_nodes.size())
  throw std::runtime_error("ICE: locationOfNode(): Bad node_index "s + std::to_string(node_index) + " vs. "s + std::to_string(m_nodes.size()));
 
 for (const auto& child_id: m_nodes[node_index].child_ids) {
  if (ChildIdIsNode(child_id)) {
   std::string location{locationOfNode(child_id)};
   if (location.size() > 0)
    return location;
  } else { // ChildIdIsToken
   return m_tokens[TokenIdFromChildId(child_id)].location.toString();
  }
 }

 return ""; // not found
}

void CPP::compileError(index_t node_id, const std::string& msg) const
{
 std::string location{locationOfNode(node_id)};
 if (location.size() == 0)
  location = "Node "s + std::to_string(node_id) + "doesn't have any token";
 throw std::runtime_error("Compile error: "s + location + ": "s + msg);
}

std::string CPP::typeOfChild(int32_t child_id) const
{
 if (ChildIdIsToken(child_id)) {
  index_t token_id {TokenIdFromChildId(child_id)};
  return m_tokens[token_id].type;
 } else { // ChildIdIsNode
  return m_nodes[child_id].type;
 }
}

bool CPP::childTypesOfNodeMatch(index_t node_id, const std::vector<std::string>& pattern) const
{
 const std::vector<int32_t>& child_ids{m_nodes[node_id].child_ids};

 if (child_ids.size() != pattern.size())
  return false; // mismatch

 for (size_t i = 0; i < child_ids.size(); i++) {
  if (pattern[i] != "" && typeOfChild(child_ids[i]) != pattern[i])
   return false; // mismatch
 }

 return true; // match
}

bool CPP::childTypesOfChildMatch(index_t index, index_t child_index, const std::vector<std::string>& pattern) const
{
 if (index >= m_nodes.size())
  return false;

 if (child_index >= m_nodes[index].child_ids.size())
  return false;

 if (!ChildIdIsNode(m_nodes[index].child_ids[child_index]))
  return false;

 return childTypesOfNodeMatch(m_nodes[index].child_ids[child_index], pattern);
}

// Get value for specified node, at bnf rule index child_index
std::any CPP::getValue(index_t node_id, index_t child_index)
{
 size_t num_values_on_top {m_nodes[node_id].child_ids.size()};

 if (num_values_on_top > mValues.size())
  throw std::runtime_error("ICE: Expected at least "s + std::to_string(num_values_on_top) + " but only have "s + std::to_string(mValues.size()) + " in total"s);

 if (child_index >= num_values_on_top)
  throw std::runtime_error("ICE: Requested value at index "s + std::to_string(child_index) + " but only have "s + std::to_string(num_values_on_top) + " values on top"s);

 return *(mValues.end() - num_values_on_top + child_index);
}

std::string CPP::getType(index_t node_id, index_t child_index)
{
 if (node_id >= m_nodes.size())
  throw std::runtime_error("ICE: Requested node at index "s + std::to_string(node_id) + " but only have "s + std::to_string(m_nodes.size()) + " nodes"s);
 
 if (child_index >= m_nodes[node_id].child_ids.size())
  throw std::runtime_error("ICE: "s + m_nodes[node_id].type + ": Requested child at index "s + std::to_string(child_index) + " but only have "s + std::to_string(m_nodes[node_id].child_ids.size()) + " childs"s);

 int32_t child_id {m_nodes[node_id].child_ids[child_index]};
 return typeOfChild(child_id);
}

std::string CPP::ruleString(index_t node_id)
{
 if (node_id >= m_nodes.size())
  throw std::runtime_error("ICE: Requested node at index "s + std::to_string(node_id) + " but only have "s + std::to_string(m_nodes.size()) + " nodes"s);

 std::string result{m_nodes[node_id].type + ":"};

 for (const auto child_id: m_nodes[node_id].child_ids) {
  result += " " + typeOfChild(child_id);
 }

 return result;
}

std::unordered_map<std::string, std::function<std::any(index_t)>> CPP::getNodeEvalMap()
{
 return {
  { "primary-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"literal"}))
     return getValue(index, 0);
    if (childTypesOfNodeMatch(index, {"(", "expression", ")"}))
     return getValue(index, 1);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "postfix-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"primary-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "unary-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"postfix-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "cast-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"unary-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "pm-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"cast-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "multiplicative-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"multiplicative-expression", "*", "pm-expression"})) {
     if (getValue(index, 0).type() != typeid(FlowGraph::Graph))
      throw std::runtime_error("ICE: multiplicative-expression: Bad data type for argument 1: "s + demangle(getValue(index, 0).type()));
     if (getValue(index, 2).type() != typeid(FlowGraph::Graph))
      throw std::runtime_error("ICE: multiplicative-expression: Bad data type for argument 3: "s + demangle(getValue(index, 2).type()));

     FlowGraph::Graph value0{std::any_cast<FlowGraph::Graph>(getValue(index, 0))};
     std::shared_ptr<FlowGraph::Node> lastOp0{value0.lastOp()};

     FlowGraph::Graph value1{std::any_cast<FlowGraph::Graph>(getValue(index, 2))};
     std::shared_ptr<FlowGraph::Node> lastOp1{value1.lastOp()};

     FlowGraph::Graph result{value0};
     result.append(value1);

     FlowGraph::Data destination{FlowGraph::MakeTemporaryInt(result.scope())};

     std::shared_ptr<FlowGraph::Node> node {std::make_shared<FlowGraph::BinaryOperation>(FlowGraph::BinaryOperationType::Multiply,
                                                                                         destination,
                                                                                         lastOp0->destination(), lastOp1->destination())};

     result.append(node);

     return result;
    }
    if (childTypesOfNodeMatch(index, {"pm-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "additive-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"additive-expression", "+", "multiplicative-expression"})) {
     if (getValue(index, 0).type() != typeid(FlowGraph::Graph))
      throw std::runtime_error("ICE: additive-expression: Bad data type for argument 1: "s + demangle(getValue(index, 0).type()));
     if (getValue(index, 2).type() != typeid(FlowGraph::Graph))
      throw std::runtime_error("ICE: additive-expression: Bad data type for argument 3: "s + demangle(getValue(index, 2).type()));

     FlowGraph::Graph value0{std::any_cast<FlowGraph::Graph>(getValue(index, 0))};
     std::shared_ptr<FlowGraph::Node> lastOp0{value0.lastOp()};

     FlowGraph::Graph value1{std::any_cast<FlowGraph::Graph>(getValue(index, 2))};
     std::shared_ptr<FlowGraph::Node> lastOp1{value1.lastOp()};

     FlowGraph::Graph result{value0};
     result.append(value1);

     FlowGraph::Data destination{FlowGraph::MakeTemporaryInt(result.scope())};

     std::shared_ptr<FlowGraph::Node> node {std::make_shared<FlowGraph::BinaryOperation>(FlowGraph::BinaryOperationType::Add,
                                                                                         destination,
                                                                                         lastOp0->destination(), lastOp1->destination())};

     result.append(node);

     return result;
    }
    if (childTypesOfNodeMatch(index, {"multiplicative-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "shift-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"additive-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "compare-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"shift-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "relational-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"compare-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "equality-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"relational-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "and-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"equality-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "exclusive-or-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"and-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "inclusive-or-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"exclusive-or-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "logical-and-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"inclusive-or-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "logical-or-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"logical-and-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "conditional-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"logical-or-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "assignment-expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"conditional-expression"}))
     return getValue(index, 0);
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
  { "expression", [&](index_t index) -> std::any
   {
    if (childTypesOfNodeMatch(index, {"assignment-expression"})) {
     if (getValue(index, 0).type() == typeid(FlowGraph::Graph)) {
      FlowGraph::Graph graph {std::any_cast<FlowGraph::Graph>(getValue(index, 0))};
      mCPPContext.graph = graph;
      return graph;
     } else {
      throw std::runtime_error("ICE: expression: Unsupported argument type: "s + demangle(getValue(index, 0).type()));
     }
    }
    throw std::runtime_error("ICE: Unsupported childs: "s + ruleString(index)); // TODO
   }
  },
 };
}

// precondition: stack contains child values c1, ..., cn on top -> to be popped
// postcondition: stack contains value on top -> to be pushed
void CPP::getValueOfNode(index_t index)
{
 size_t num_childs {m_nodes[index].child_ids.size()};

 if (mValues.size() < num_childs)
  throw std::runtime_error("ICE: Expected num_childs elements on Values stack at "s + locationOfNode(index));

 auto function_it{node_eval_map.find(m_nodes[index].type)};
 std::any result;
 if (function_it != node_eval_map.end()) {
  result = function_it->second(index);
 } else {
  //TODO: throw std::runtime_error("ICE: Node type not implemented: "s + m_nodes[index].type);
 }

 mValues.resize(mValues.size() - num_childs);

 mValues.push_back(result);
}

// pushes result onto stack
void CPP::getValueOfToken(index_t index)
{
 // TODO: also support the other tokens ...
 if (m_tokens[index].type == "literal") {
  // TODO: also support other types, different from Int
  FlowGraph::Data data{FlowGraph::MakeConstantInt(stoi(m_tokens[index].value))};
  FlowGraph::Graph graph{{std::make_shared<FlowGraph::DataNode>(data)}};
  mValues.push_back(graph);
 } else {
  mValues.push_back(std::any{});
 }
}

void CPP::visitRecursive(index_t node_id)
{
 const auto& childs {m_nodes[node_id].child_ids};
 for (const auto child: childs) {
  if (ChildIdIsNode(child)) {
   visitRecursive(child);
  } else {
   getValueOfToken(TokenIdFromChildId(child));
  }
 }
 getValueOfNode(node_id);
}

// Phase 7.c: Translate
void CPP::translate()
{
 if (m_nodes.size() == 0)
  throw std::runtime_error("ICE: Tree is empty");

 // TODO: this could be implemented iteratively, via a stack?
 visitRecursive(0);
}

// Phase 8: Instantiate objects
void CPP::instantiate()
{
 // TODO: template instantiation

 Asm::toMachineCode(mCPPContext.graph, mSegment);
}

// Phase 9: Link libraries
void CPP::link()
{
 // TODO
 
 // mSegment -> elf

 mCode = std::vector<uint8_t>{
            0x48, 0xc7, 0xc0, 0x3c, 0x00, 0x00, 0x00, // mov    $0x3c,%rax     # syscall 60
            0x48, 0x31, 0xff,                         // xor    %rdi,%rdi      # exit code 0
 } + mSegment.getCode() + std::vector<uint8_t>{       //                       # leave exit code in edi
            0x0f, 0x05,                               // syscall
 }
 ;
}

// phases of translation, according to standard
void CPP::compile(const std::string& source)
{
 m_source = source;

 source_charset_map(); // phase 1
 
 backslash_escape(); // phase 2
 
 auto pp_tokens = preprocessing_tokenize(m_source); // phase 3
 
 preprocess(); // phase 4
 
 execution_charset_map(); // phase 5
 
 concatenate_strings(); // phase 6
 
 m_tokens = tokens_from_pptokens(pp_tokens); // phase 7a
 m_nodes = analysis(m_tokens); // phase 7b
 translate(); // phase 7c
 
 instantiate(); // phase 8
 
 link(); // phase 9
}

std::vector<uint8_t> CPP::getCode()
{
 return mCode;
}

std::vector<uint8_t> CPP::getData()
{
 return {}; // TODO
}