#include <iostream>
#include <filesystem>
#include <fstream>
#include <stdint.h>
#include <string>
#include <map>
#include <vector>
#include <iterator>
#include <algorithm>
#include <tree_sitter/api.h>
#include <tree_sitter/tree-sitter-c.h>
using namespace std;
namespace fs=filesystem;
const vector<bool> CCC_C_KEYYORD_HEAD {0,0,0};
const vector<bool> CCC_SPACE {0,1,1,1,0,0,1};
const vector<bool> CCC_PREPROCESSOR_CONDITIONAL_IF {0,0,1,0,0,0};
const vector<bool> CCC_PREPROCESSOR_CONDITIONAL_IFDEF {0,0,1,0,0,1};
const vector<bool> CCC_PREPROCESSOR_CONDITIONAL_IFNDEF {0,0,1,0,1,0};
const vector<bool> CCC_PREPROCESSOR_CONDITIONAL_ELSE {0,0,1,0,1,1};
const vector<bool> CCC_PREPROCESSOR_CONDITIONAL_ELIF {0,0,1,1,0,0};
const vector<bool> CCC_PREPROCESSOR_CONDITIONAL_ELIFDEF {0,0,1,1,0,1};
const vector<bool> CCC_PREPROCESSOR_CONDITIONAL_ELIFNDEF {0,0,1,1,1,0};
const vector<bool> CCC_PREPROCESSOR_CONDITIONAL_ENDIF {0,0,1,1,1,1};
const vector<bool> CCC_PREPROCESSOR_OTHER_DEFINE {0,1,0,0,0,0};
const vector<bool> CCC_PREPROCESSOR_OTHER_UNDEF {0,1,0,0,0,1};
const vector<bool> CCC_PREPROCESSOR_OTHER_INCLUDE {0,1,0,0,1,0};
const vector<bool> CCC_PREPROCESSOR_OTHER_ERROR {0,1,0,0,1,1};
const vector<bool> CCC_PREPROCESSOR_OTHER_WARNING {0,1,0,1,0,0};
const vector<bool> CCC_PREPROCESSOR_OTHER_PRAGMA {0,1,0,1,0,1};
const vector<bool> CCC_PREPROCESSOR_OTHER_LINE {0,1,0,1,1,0};
const vector<bool> CCC_QUOTE {0,1,0,1,1,1};
const vector<bool> CCC_DELIMITER_HEAD {0,1,1};
const vector<bool> CCC_OTHER_GRAMMAR_HEAD {1,0,0};
const vector<bool> CCC_MISCELLANEOUS_HEAD {1,0,1};
const vector<bool> CCC_REC_TABLE_REF_HEAD {1,1,0};
const vector<bool> CCC_STRING_ASCII {1,1,1,0};
const vector<bool> CCC_STRING_UTF8 {1,1,1,1};
const vector<bool> CCC_STRING_END_ASCII {0,0,0,0,0,0,0};
const vector<bool> CCC_STRING_END_UTF8 {0,0,0,0,0,0,0,0};
#define CCC_ADD_COMPOMENT(vec,tail) \
  do { \
    auto tmp=tail; \
    vec.insert(vec.end(),tmp.begin(),tmp.end()); \
  } while (0)
const vector<string> delimiter={
  "\n",
  "\t",
  "{",
  "}",
  "(",
  ")",
  "[",
  "]",
  " ",
  "{}",
  "()",
  "[]",
  "",
  ";",
  ",",
  "."
};
const vector<string> other_grammer={
  "!",
  "%",
  "'",
  "*",
  "+",
  "-",
  "/",
  ":",
  "<",
  ">",
  "=",
  "?",
  "^",
  "|",
  "&",
  "~"
};
const vector<string> miscellaneous={
  "+=",
  "-=",
  "*=",
  "/=",
  "%=",
  "&=",
  "|=",
  "^=",
  "<<=",
  ">>=",
  "++",
  "--",
  "<<",
  ">>",
  "==",
  "!=",
  "<=",
  ">=",
  "->",
  "...",
  "||",
  "&&",
  "NULL",
  "size_t",
  "uint8_t",
  "uint16_t",
  "uint32_t",
  "uint64_t",
  "int8_t",
  "int16_t",
  "int32_t",
  "int64_t"
};
const vector<string> c_keywords={
  "alignas",
  "alignof",
  "auto",
  "bool",
  "break",
  "case",
  "char",
  "const",
  "constexpr",
  "continue",
  "default",
  "do",
  "double",
  "else",
  "enum",
  "extern",
  "false",
  "float",
  "for",
  "goto",
  "if",
  "inline",
  "int",
  "long",
  "nullptr",
  "register",
  "restrict",
  "return",
  "short",
  "signed",
  "sizeof",
  "static",
  "static_assert",
  "struct",
  "switch",
  "thread_local",
  "true",
  "typedef",
  "typeof",
  "typeof_unequal",
  "union",
  "unsigned",
  "void",
  "volatile",
  "while"
};
struct symbol {
  string name;
  int score;
};
struct node {
  map<unsigned char,node*> children;
  int token_id=-1;
};
void insert(node* root,string str,int id) {
  node* curr=root;
  for (char c:str) {
    if (curr->children.find(c)==curr->children.end()) {
      curr->children[c]=new node();
    }
    curr=curr->children[c];
  }
  curr->token_id=id;
}
vector<TSNode> all_tokens;
void get_all_nodes(TSNode node,const string &source_code,map<string,int> &rec_map) {
  if (ts_node_child_count(node)==0) {
    all_tokens.push_back(node);
    string text=source_code.substr(ts_node_start_byte(node),ts_node_end_byte(node)-ts_node_start_byte(node));
    if (string(ts_node_type(node))=="string_content" || string(ts_node_type(node))=="system_lib_string" || string(ts_node_type(node))=="identifier" || string(ts_node_type(node))=="number_literal" || string(ts_node_type(node))=="type_identifier" || string(ts_node_type(node))=="field_identifier" || string(ts_node_type(node))=="escape_sequence" || string(ts_node_type(node))=="statement_identifier") {
      rec_map[text]++;
    }
    if (string(ts_node_type(node))=="primitive_type" && find(c_keywords.begin(),c_keywords.end(),text)==c_keywords.end()) {
      rec_map[text]++;
    }
  } else {
    uint32_t child_count=ts_node_child_count(node);
    for (uint32_t i=0;i<child_count;++i) {
      TSNode child=ts_node_child(node,i);
      get_all_nodes(child,source_code,rec_map);
    }
  }
}
vector<bool> byte_to_bits(unsigned char c) {
  vector<bool> out;
  for (int i=7;i>=0;i--) {
    bool enabled=(c>>i)&0x01;
    out.push_back(enabled);
  }
  return out;
}
vector<bool> ascii_to_bits(unsigned char c) {
  vector<bool> out;
  for (int i=6;i>=0;i--) {
    bool enabled=(c>>i)&0x01;
    out.push_back(enabled);
  }
  return out;
}
vector<bool> generate_c_keyword(size_t index) {
  vector<bool> out;
  CCC_ADD_COMPOMENT(out,CCC_C_KEYYORD_HEAD);
  for (int i=5;i>=0;i--) {
    bool enabled=(index>>i)&0x01;
    out.push_back(enabled);
  }
  return out;
}
vector<bool> generate_rec(size_t index,size_t total_recs) {
  vector<bool> out;
  size_t bits=0;
  while (total_recs) {
    total_recs>>=1;
    ++bits;
  }
  CCC_ADD_COMPOMENT(out,CCC_REC_TABLE_REF_HEAD);
  for (int i=bits;i>=0;i--) {
    bool enabled=(index>>i)&0x01;
    out.push_back(enabled);
  }
  return out;
}
vector<bool> generate_delimiter(size_t index) {
  vector<bool> out;
  CCC_ADD_COMPOMENT(out,CCC_DELIMITER_HEAD);
  for (int i=3;i>=0;i--) {
    bool enabled=(index>>i)&0x01;
    out.push_back(enabled);
  }
  return out;
}
vector<bool> generate_other_grammar(size_t index) {
  vector<bool> out;
  CCC_ADD_COMPOMENT(out,CCC_OTHER_GRAMMAR_HEAD);
  for (int i=3;i>=0;i--) {
    bool enabled=(index>>i)&0x01;
    out.push_back(enabled);
  }
  return out;
}
vector<bool> generate_miscellaneous(size_t index) {
  vector<bool> out;
  CCC_ADD_COMPOMENT(out,CCC_MISCELLANEOUS_HEAD);
  for (int i=4;i>=0;i--) {
    bool enabled=(index>>i)&0x01;
    out.push_back(enabled);
  }
  return out;
}
vector<bool> generate_string_content(string str) {
  vector<bool> out;
  bool is_utf8=false;
  for (auto c:str) {
    if (c>127) {
      is_utf8=true;
      break;
    }
  }
  if (is_utf8) {
    CCC_ADD_COMPOMENT(out,CCC_STRING_UTF8);
    for (auto c:str) {
      CCC_ADD_COMPOMENT(out,byte_to_bits(c));
    }
    CCC_ADD_COMPOMENT(out,CCC_STRING_END_UTF8);
  } else {
    CCC_ADD_COMPOMENT(out,CCC_STRING_ASCII);
    for (auto c:str) {
      CCC_ADD_COMPOMENT(out,ascii_to_bits(c));
    }
    CCC_ADD_COMPOMENT(out,CCC_STRING_END_ASCII);
  }
  return out;
}
vector<bool> process_all_nodes(vector<TSNode> *nodes,string code,vector<string> &rec_list) {
  vector<bool> out;
  for (int i=0;i<nodes->size();i++) {
    string type=string(ts_node_type(nodes->at(i)));
    if (type=="#if") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_CONDITIONAL_IF);
      cout<<"if"<<endl;
    } else if (type=="#ifdef") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_CONDITIONAL_IFDEF);
      cout<<"ifdef"<<endl;
    } else if (type=="#ifndef") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_CONDITIONAL_IFNDEF);
      cout<<"ifndef"<<endl;
    } else if (type=="#else") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_CONDITIONAL_ELSE);
      cout<<"else"<<endl;
    } else if (type=="#elif") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_CONDITIONAL_ELIF);
      cout<<"elif"<<endl;
    } else if (type=="#elifdef") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_CONDITIONAL_ELIFDEF);
      cout<<"elifdef"<<endl;
    } else if (type=="#elifndef") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_CONDITIONAL_ELIFNDEF);
      cout<<"elifndef"<<endl;
    } else if (type=="#endif") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_CONDITIONAL_ENDIF);
      cout<<"endif"<<endl;
    } else if (type=="#define") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_OTHER_DEFINE);
      cout<<"define"<<endl;
    } else if (type=="#undef") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_OTHER_UNDEF);
      cout<<"undef"<<endl;
    } else if (type=="#include") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_OTHER_INCLUDE);
      cout<<"include"<<endl;
    } else if (type=="#error") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_OTHER_ERROR);
      cout<<"error"<<endl;
    } else if (type=="#warning") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_OTHER_WARNING);
      cout<<"warning"<<endl;
    } else if (type=="#pragma") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_OTHER_PRAGMA);
      cout<<"pragma"<<endl;
    } else if (type=="#line") {
      CCC_ADD_COMPOMENT(out,CCC_PREPROCESSOR_OTHER_LINE);
      cout<<"line"<<endl;
    } else if (type=="string_content" || type=="system_lib_string" || type=="identifier" || type=="number_literal" || type=="type_identifier" || type=="field_identifier" || type=="preproc_arg" || type=="escape_sequence" || type=="character" || type=="statement_identifier") {
      string text=code.substr(ts_node_start_byte(nodes->at(i)),ts_node_end_byte(nodes->at(i))-ts_node_start_byte(nodes->at(i)));
      auto it=find(rec_list.begin(),rec_list.end(),text);
      if (it==rec_list.end()) {
        if (!text.empty()) {
          string text=code.substr(ts_node_start_byte(nodes->at(i)),ts_node_end_byte(nodes->at(i))-ts_node_start_byte(nodes->at(i)));
          CCC_ADD_COMPOMENT(out,generate_string_content(text));
          cout<<"string ("<<type<<"): "<<text<<endl;
        } else {
          auto it=find(delimiter.begin(),delimiter.end(),"");
          size_t index=distance(delimiter.begin(),it);
          CCC_ADD_COMPOMENT(out,generate_delimiter(index));
          cout<<"delimiter for empty string"<<endl;
        }
      } else {
        size_t index=distance(rec_list.begin(),it);
        CCC_ADD_COMPOMENT(out,generate_rec(index,rec_list.size()));
        cout<<"rec_table for string ("<<type<<"): "<<text<<endl;
      }
    } else if (type=="primitive_type") {
      string text=code.substr(ts_node_start_byte(nodes->at(i)),ts_node_end_byte(nodes->at(i))-ts_node_start_byte(nodes->at(i)));
      auto it=find(c_keywords.begin(),c_keywords.end(),text);
      if (it!=c_keywords.end()) {
        size_t index=distance(c_keywords.begin(),it);
        CCC_ADD_COMPOMENT(out,generate_c_keyword(index));
        cout<<"primitive_type: "<<text<<endl;
      } else {
        auto it=find(rec_list.begin(),rec_list.end(),text);
        if (it==rec_list.end()) {
          if (!text.empty()) {
            CCC_ADD_COMPOMENT(out,generate_string_content(text));
            cout<<"string ("<<type<<"): "<<text<<endl;
          } else {
            cout<<"Error: provided primitive is empty: "<<text;
            exit(-1);
          }
        } else {
          size_t index=distance(rec_list.begin(),it);
          CCC_ADD_COMPOMENT(out,generate_rec(index,rec_list.size()));
          cout<<"rec_table for string ("<<type<<"): "<<text<<endl;
        }
      }
    } else if (find(delimiter.begin(),delimiter.end(),type)!=delimiter.end()) {
      string text;
      if (type=="(" && i+1<nodes->size()) {
        if (string(ts_node_type(nodes->at(i+1)))==")") {
          text="()";
          i++;
        } else {
          text="(";
        }
      } else if (type=="[" && i+1<nodes->size()) {
        if (string(ts_node_type(nodes->at(i+1)))=="]") {
          text="[]";
          i++;
        } else {
          text="[";
        }
      } else if (type=="{" && i+1<nodes->size()) {
        if (string(ts_node_type(nodes->at(i+1)))=="}") {
          text="{}";
          i++;
        } else {
          text="{";
        }
      } else {
        text=type;
      }
      auto it=find(delimiter.begin(),delimiter.end(),text);
      if (it!=delimiter.end()) {
        size_t index=distance(delimiter.begin(),it);
        CCC_ADD_COMPOMENT(out,generate_delimiter(index));
        cout<<"delimiter: "<<text<<endl;
      } else {
        cout<<"Error: unknow delimiter, that shouldn't happen: "<<text<<endl;;
        exit(-1);
      }
    } else if (find(other_grammer.begin(),other_grammer.end(),type)!=other_grammer.end()) {
      auto it=find(other_grammer.begin(),other_grammer.end(),type);
      if (it!=other_grammer.end()) {
        size_t index=distance(other_grammer.begin(),it);
        CCC_ADD_COMPOMENT(out,generate_other_grammar(index));
        cout<<"other grammar: "<<type<<endl;
      } else {
        cout<<"Error: unknow other grammar symbol, that shouldn't happen: "<<type<<endl;;
        exit(-1);
      }
    } else if (find(c_keywords.begin(),c_keywords.end(),type)!=c_keywords.end()) {
      auto it=find(c_keywords.begin(),c_keywords.end(),type);
      if (it!=c_keywords.end()) {
        size_t index=distance(c_keywords.begin(),it);
        CCC_ADD_COMPOMENT(out,generate_c_keyword(index));
        cout<<"c keyword: "<<type<<endl;
      } else {
        cout<<"Error: unknow C keyword, that shouldn't happen: "<<type<<endl;;
        exit(-1);
      }
    } else if (find(miscellaneous.begin(),miscellaneous.end(),type)!=miscellaneous.end()) {
      auto it=find(miscellaneous.begin(),miscellaneous.end(),type);
      if (it!=miscellaneous.end()) {
        size_t index=distance(miscellaneous.begin(),it);
        CCC_ADD_COMPOMENT(out,generate_miscellaneous(index));
        cout<<"miscellaneous: "<<type<<endl;
      } else {
        cout<<"Error: unknow miscellaneous, that shouldn't happen: "<<type<<endl;;
        exit(-1);
      }
    } else if (type=="\"") {
      if (i+1<nodes->size()) {
        if (string(ts_node_type(nodes->at(i+1)))=="\"") {
          auto it=find(delimiter.begin(),delimiter.end(),"");
          size_t index=distance(delimiter.begin(),it);
          CCC_ADD_COMPOMENT(out,generate_delimiter(index));
          cout<<"double quotes mark, inserting delimiter for empty string"<<endl;
          i++;
        } else {
          CCC_ADD_COMPOMENT(out,CCC_QUOTE);
          cout<<"single quote mark"<<endl;
        }
      }
    } else if (type=="comment") {
      continue;
    } else {
      string text=code.substr(ts_node_start_byte(nodes->at(i)),ts_node_end_byte(nodes->at(i))-ts_node_start_byte(nodes->at(i)));
      cout<<"unknow node type: "<<type<<endl;
      cout<<"unknow node text: "<<text<<endl;
      exit(-1);
    }
  }
  return out;
}
int main(int argc,char **argv) {
  if (argc!=2) {
    cout<<"Usage: ccc <c file>"<<endl;
    return -1;
  }
  string filepath=string(argv[1]);
  if (!fs::exists(filepath)) {
    cout<<"Error: provided file doesn't exist."<<endl;
    return -1;
  }
  ifstream file(filepath,ios::binary);
  if (!file) {
    cout<<"Error: couldn't open provided file."<<endl;
    return -1;
  }
  string code((istreambuf_iterator<char>(file)),istreambuf_iterator<char>());
  TSParser *parser=ts_parser_new();
  ts_parser_set_language(parser,tree_sitter_c());
  TSTree *tree=ts_parser_parse_string(parser,nullptr,code.c_str(),code.size());
  TSNode root=ts_tree_root_node(tree);
  map<string,int> rec_map;
  vector<string> rec_list;
  get_all_nodes(root,code,rec_map);
  for (auto s:rec_map) {
    if (s.second>=2 and s.first.size()>=3 && s.first.size()<=256) {
      rec_list.push_back(s.first);
    }
  }
  auto payload=process_all_nodes(&all_tokens,code,rec_list);
  vector<bool> out={0,1,0,0,0,0,1,1,0,1,0,0,0,0,1,1,0,1,0,0,0,0,1,1};
  for (int i=63;i>=0;i--) {
    bool enabled=(rec_list.size()>>i)&0x01;
    out.push_back(enabled);
  }
  for (int i=0;i<rec_list.size();i++) {
    uint8_t size=(uint8_t)rec_list[i].size();
    for (int i=7;i>=0;i--) {
      bool enabled=(size>>i)&0x01;
      out.push_back(enabled);
    }
    for (auto c:rec_list[i]) {
      for (int i=7;i>=0;i--) {
        bool enabled=(c>>i)&0x01;
        out.push_back(enabled);
      }
    }
  }
  CCC_ADD_COMPOMENT(out,payload);
  vector<unsigned char> outbytes;
  unsigned char current=0;
  int bit_index=0;
  for (bool b:out) {
    current|=(b<<(7-bit_index));
    bit_index++;
    if (bit_index==8) {
      outbytes.push_back(current);
      current=0;
      bit_index=0;
    }
  }
  if (bit_index!=0) {
    outbytes.push_back(current);
  }
  ofstream fileout(filepath+".ccc",ios::binary);
  if (!fileout) {
    cout<<"Error: couldn't open output file."<<endl;
    return -1;
  }
  fileout.write(reinterpret_cast<const char*>(outbytes.data()),outbytes.size());
  fileout.close();
  cout<<"Reccurences map entry count: "<<rec_list.size()<<endl;
  size_t total_bytes=0;
  for (int i=0;i<rec_list.size();i++) {
    total_bytes++;
    total_bytes+=rec_list[i].size();
  }
  cout<<"Total spaces taken by reccurences map in bytes: "<<total_bytes<<endl;
  return 0;
}