#include <iostream>
#include <sstream>
#include <fstream>
#include <vector>
#include <set>
#include <pqxx/pqxx>
//---------------------------------------------------------------------------
// (c) 2014 Thomas Neumann.
//
// This file may be used according to the terms of the
// GNU General Public License version 3.0
// as published by the Free Software Foundation.
//---------------------------------------------------------------------------
using namespace std;
//---------------------------------------------------------------------------
namespace {
//---------------------------------------------------------------------------
/// A column info
struct Column {
   std::string name,type;
   unsigned id;
};   
//---------------------------------------------------------------------------
/// Create a DB connection on demand, it might not be needed for the input
/// TODO: Implement caching, so that we can compile without DB connection
class LazyDbConnection
{
   private:
   pqxx::lazyconnection conn;
   
   public:
   /// Constructor
   LazyDbConnection(const string& server,const string& port,const string& name) : conn("dbname='"+name+"' host='"+server+"' port='"+port+"'") {}
   
   /// Extract the signaure
   vector<Column> extractSignature(const string& query);   
};
//---------------------------------------------------------------------------
vector<Column> LazyDbConnection::extractSignature(const string& query)
   // Extract the signaure
{
   pqxx::work work(conn);
   auto queryResult=work.exec("explain schema "+query);
   vector<Column> result;
   for (auto&& r:queryResult) {
      if (r[6].as<bool>())
         throw pqxx::sql_error("query parameters currently not supported",query);
      Column c;
      c.name=r[0].as<string>();
      c.id=r[5].as<int>();
      // PostgreSQL type codes
      switch (r[2].as<int>()) {
         case 23: c.type="int"; break;
         case 1700: case 701: c.type="double"; break;
         case 8: case 1042: case 1043: c.type="std::string"; break;
         case 16: c.type="bool"; break;
         case 21: c.type="short"; break;
         case 20: c.type="long long"; break;
         // TODO these types are currently only returned as strings, should pick proper C++ types
         case 17 /*bytea*/: case 1082 /*date*/: case 1083 /*time*/: case 1114 /*timestamp*/:
         case 1186 /*interval*/: case 704 /*interval*/: case 114 /*JSON*/: case 705 /*Unknown*/:
            c.type="std::string"; break;
         default:
            throw pqxx::sql_error("unsupported column type "+r[1].as<string>(),query);
      }
      result.push_back(c);
   }
   return result;
}
//---------------------------------------------------------------------------
/// Extractor
class Extractor
{
   private:
   /// The connection
   LazyDbConnection db;
   /// Next query id
   unsigned nextId;
   /// The input
   ifstream in;
   /// The file name
   string fileName;
   /// The current line
   unsigned line;
   /// The output
   stringstream code,header;

   /// Get the next character
   bool next(char& c);
   /// Report an error
   void reportError(const string& message);
   /// Process a // comment
   void processNewComment();
   /// Process a /* ... */ comment
   void processOldComment();
   /// Process a string
   void processString();   
   /// Extract the signaure
   vector<Column> extractSignature(const string& query);
   /// Translate a query
   void translateQuery(const string& query);
   /// Extract and process a query. Extracts up to the closing bracket
   void processQuery();

   public:
   /// Constructor
   Extractor(const string& db,const string& port,const string& name) : db(db,port,name),nextId(0),line(1) {}
   
   /// Process the file and replace embedded SQL of the form " ... : SQL ... )"
   bool processFile(const string& inFile,const string& outFile);
};
//---------------------------------------------------------------------------
static inline bool isWS(char c) { return (c==' ')||(c=='\t')||(c=='\n')||(c=='\r'); }
//---------------------------------------------------------------------------
bool Extractor::next(char& c)
   // Get the next character
{
   c=in.get();
   if (!in) return false;
   if (c=='\n') ++line;
   return true;
}
//---------------------------------------------------------------------------
void Extractor::reportError(const string& message)
   // Report an error
{
   cerr << fileName << ":" << line << ": " << message << endl;
   exit(1);
}
//---------------------------------------------------------------------------
void Extractor::processNewComment()
   // Process a // comment
{
   char c;
   while (next(c)) {
      code << c;
      if (c=='\n') break;
   }
}
//---------------------------------------------------------------------------
void Extractor::processOldComment()
   // Process an /* ... */ comment
{
   char c;
   while (next(c)) {
      checkChar: code << c;
      if (c=='*') {
         if (!next(c)) break;
         if (c=='/') { code << c; break; }
         goto checkChar;
      }
   }
}
//---------------------------------------------------------------------------
void Extractor::processString()
   // Process a string
{
   char c;
   while (next(c)) {
      code << c;
      if (c=='"') break;
      if (c=='\\') {
         if (!next(c)) break;
         if (c=='\n') ++line;
      }
   }
}
//---------------------------------------------------------------------------
vector<Column> Extractor::extractSignature(const string& query)
   // Extract the signaure
{
   try {
      return db.extractSignature(query);
   } catch (const pqxx::pqxx_exception& e) {
      reportError(e.base().what());
      throw; // unreachable
   }
}
//---------------------------------------------------------------------------
static bool isValidName(const string& name)
   // Is avalid C++ name?
{
   if (name.empty()) return false;
   char c=name[0];
   if (c=='_') {
      if ((name.length()>1)&&(name[1]=='_')) return false;
   } else {
      if (!(((c>='A')&&(c<='Z'))||((c>='a')&&(c<='z')))) return false;
   }
   for (char c:name) {
      if (!(((c>='A')&&(c<='Z'))||((c>='a')&&(c<='z'))||((c>='0')&&(c<='9'))||(c=='_'))) return false;
   }
   static const char* const keywords[]={"alignas","alignof","and","and_eq","asm","auto","bitand","bitor","bool","break","case","catch","char","char16_t","char32_t","class","compl","const","constexpr","const_cast","continue","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","not","not_eq","nullptr","operator","or","or_eq","private","protected","public","register","reinterpret_cast","return","short","signed","sizeof","static","static_assert","static_cast","struct","switch","template","this","thread_local","throw","true","try","typedef","typeid","typename","union","unsigned","using(1)","virtual","void","volatile","wchar_t","while","xor","xor_eq"};
   for (auto k:keywords)
      if (k==name)
         return false;
   return true;
}
//---------------------------------------------------------------------------
static string getColumnName(set<string>& names,const string& originalName,unsigned id)
   // Construct a valid name
{
   string name;
   if (isValidName(originalName))
      name=originalName; else
      name="col"+to_string(id);
   if (!names.count(name)) {
      names.insert(name);
      return name;
   }
   for (unsigned index=2;;++index) {
      string n=name+"_"+to_string(index);
      if (!names.count(n)) {
         names.insert(n);
         return n;
      }
   }
}
//---------------------------------------------------------------------------
void Extractor::translateQuery(const string& query)
   // Translate the SQL statement
{
   vector<Column> columns=extractSignature(query);
   if (!nextId) {
      header << "#include <pqxx/pqxx>" << endl
             << "namespace {" << endl;
   }
   header << "class __sql_QueryProvider" << nextId << " {" << endl
          << "   pqxx::result result;" << endl
          << "   public: __sql_QueryProvider" << nextId << "(pqxx::work& work) : result(work.exec(\"";
   for (auto c:query) {
      if (c=='\\') header << "\\\\"; else
      if (c=='\"') header << "\\\""; else
      if (c<' ') {
         unsigned v=c&0xFF;
         header << "\\" << static_cast<char>('0'+(v/64)) << static_cast<char>('0'+((v/8)%8)) << static_cast<char>('0'+(v%8));
      } else header << c;
   }
   header << "\")) {}" << endl
          << "   struct const_iterator;" << endl
          << "   struct __sql_proxy {" << endl
          << "      private: friend struct const_iterator;" << endl
          << "      pqxx::result::const_iterator __sql_iterator;" << endl
          << "      __sql_proxy(pqxx::result::const_iterator iter) : __sql_iterator(iter) {}" << endl
          << "      public:" << endl;
   set<string> columnNames;
   for (auto& c:columns) {
      string name=getColumnName(columnNames,c.name,c.id);
      header << "      " << c.type << " " << name << "() const { return (*__sql_iterator)[" << c.id << "].as<" << c.type << ">(); }\n";
      header << "      bool " << name << "_isnull() const { return (*__sql_iterator)[" << c.id << "].is_null(); }\n";      
   }
   header << "   };" << endl
          << "   struct const_iterator {" << endl
          << "      pqxx::result::const_iterator iter;" << endl
          << "      const_iterator(pqxx::result::const_iterator iter) : iter(iter) {}" << endl
          << "      __sql_proxy operator*() const { return __sql_proxy(iter); }" << endl
          << "      bool operator==(const const_iterator& o) const { return iter==o.iter; }" << endl
          << "      bool operator!=(const const_iterator& o) const { return iter!=o.iter; }" << endl
          << "      const_iterator& operator++() { ++iter; return *this; }" << endl
          << "   };" << endl
          << "   const_iterator begin() const { return const_iterator(result.begin()); }" << endl
          << "   const_iterator end() const { return const_iterator(result.end()); }" << endl
          << "};" << endl;
   code << "__sql_QueryProvider" << nextId << "(work)";
   nextId++;
}
//---------------------------------------------------------------------------
void Extractor::processQuery()
   // Extract and process a query. Extracts up to the closing bracket
{
   // Extract the query
   stringstream query;
   unsigned level=0,start=line;
   char c;
   while (next(c)) {
      checkChar: switch (c) {
         case '-': query << c; if (!next(c)) break; if (c=='-') { while (next(c)) { if (c=='\n') goto checkChar; query << c; } break;} else goto checkChar;
         case '(': query << c; ++level; break;
         case ')': if (!level) { translateQuery(query.str()); while (start<line) { cout << "\n"; ++start; } code << ")"; return; } query << c; --level; break;
         case '\'': while (next(c)) { if (c=='\'') { query << c; if (!next(c)) break; if (c!='\'') goto checkChar; } query << c; } break; 
         case '\"': while (next(c)) { if (c=='\"') { query << c; if (!next(c)) break; if (c!='\"') goto checkChar; } query << c; } break;
         default: query << c; break;
      }
   }
   
   // We did not find the closing bracket
   stringstream s;
   s << "unterminated SQL query starting from line " << start;
   reportError(s.str());
}
//---------------------------------------------------------------------------
bool Extractor::processFile(const string& inFile,const string& outFile)
   // Process the file and replace embedded SQL of the form " ... : SQL ... )"
   // It is meant to be used as for (auto& t:SQL select * from foo)
{
   if (outFile==inFile) {
      cerr << "refusing to overwrite input file" << endl;
      return false;
   }
   in.open(inFile);
   if (!in.is_open()) {
      cerr << "unable to open " << inFile << endl;
      return false;
   }
   fileName=inFile;
   // Scan the file for embedded SQL commands, keeping the rest as it is
   char c;
   while (next(c)) {
      checkChar: switch (c) {
         case '/': code << c; if (!next(c)) break; if ((c=='/')||(c=='*')) { code << c; if (c=='/') processNewComment(); else processOldComment(); break; } else goto checkChar;
         case '"': code << c; processString(); break;
         case ':': {
            // currently a bit crude. Check for :<ws>*SQL<ws>
            code << c;
            if (!next(c)) break;
            const char* hadWS="";
            while (isWS(c)) { hadWS=" "; if (!next(c)) break; }
            if (c!='S') { code << hadWS; goto checkChar; } if (!next(c)) break;
            if (c!='Q') { code << hadWS << "S"; goto checkChar; } if (!next(c)) break;
            if (c!='L') { code << hadWS << "SQ"; goto checkChar; } if (!next(c)) break;
            if (!isWS(c)) { code << hadWS << "SQL"; goto checkChar; }
            processQuery();
            break;
         }
         default: code << c; break;
      }
   }
   
   // Write out result
   ofstream out(outFile);
   if (!out.is_open()) {
      cerr << "unable to write " << outFile << endl;
      return false;
   }
   if (!nextId) {
      out << code.str();
   } else {
      out << header.str() 
          << "}" << endl
          << "#line 1 \"" << inFile << "\"" << endl;
      out << code.str();
   }
   return true;
}
//---------------------------------------------------------------------------
}   
//---------------------------------------------------------------------------
int main(int argc,char* argv[])
{
   if (argc!=6) {
      cerr << "usage: " << argv[0] << " dbhost dbport dbname infile outfile" << endl;
      return 1;
   }
   Extractor extractor(argv[1],argv[2],argv[3]);
   return !extractor.processFile(argv[4],argv[5]);
}
//---------------------------------------------------------------------------