#include "Database.h"

#include <regex>

#include <boost/algorithm/string/join.hpp>

#include "tools.h"
#include "Time.h"
#include "Configuration.h"

#include <TROOT.h>
#include <TSystem.h>
#include <TFile.h>
#include <TTree.h>

using namespace std;

// ------------------------------------------------------------------------

void SetupConfiguration(Configuration &conf)
{
    po::options_description control("Rootify SQL");
    control.add_options()
        ("uri,u",         var<string>()->required(),   "Database link as in\n\tuser:password@server[:port]/database.")
        ("query,q",       var<string>(""),             "MySQL query (overwrites --file)")
        ("file",          var<string>("rootify.sql"),  "An ASCII file with the MySQL query (overwrites --query)")
        ("ignore-null,i", po_switch(),                 "Do not skip rows containing any NULL field")
        ("out,o",         var<string>("rootify.root"), "Output root file name")
        ("force,f",       po_switch(),                 "Force overwriting an existing root file ('RECREATE')")
        ("update",        po_switch(),                 "Update an existing root file with the new tree ('UPDATE')")
        ("compression,c", var<uint16_t>(1),            "zlib compression level for the root file")
        ("tree,t",        var<string>("Result"),       "Name of the root tree")
        ("ignore",        vars<string>(),              "Ignore the given columns")
        ("display,d",     po_switch(),                 "Displays contents on the screen (most usefull in combination with mysql statements as SHOW or EXPLAIN)")
        ("null,n",        po_switch(),                 "Redirect the output file to /dev/null (mainly for debugging purposes, e.g. performance studies)")
        ("no-fill",       po_switch(),                 "Do not fill events into the root file (mainly for debugging purposes, e.g. performance studies)")
        ("delimiter",     var<string>(""),             "The delimiter used if contents are displayed with --display (default=\\t)")
        ("var.*",         var<string>(),               "Predefined SQL user variables (@VAR)")
        ("env.*",         vars<string>(),              "Predefined environment for substitutions in the query ($ENV)")
        ("verbose,v",     var<uint16_t>(1),            "Verbosity (0: quiet, 1: default, 2: more, 3, ...)")
        ;

    po::positional_options_description p;
    p.add("file", 1); // The 1st positional options (n=1)
    p.add("out",  1); // The 2nd positional options (n=1)

    conf.AddOptions(control);
    conf.SetArgumentPositions(p);
}

void PrintUsage()
{
    cout <<
        "rootifysql - Converts the result of a mysql query into a root file\n"
        "\n"
        "For convenience, this documentation uses the extended version of the options, "
        "refer to the output below to get the abbreviations.\n"
        "\n"
        "Writes the result of a mysql query into a root file. For each column, a branch is "
        "created of type double with the field name as name. This is usually the column name "
        "if not specified otherwise by the AS mysql directive.\n"
        "\n"
        "Columns with CHAR or VARCHAR as field type are ignored. DATETIME, DATE and TIME "
        "columns are converted to unix time (time_t). Rows containing any file which is "
        "NULL are skipped if not suppressed by the --ignore-null option. Ideally, the query "
        "is compiled in a way that no NULL field is returned. With the --display option the "
        "result of the request is printed on the screen (NULL skipping still in action). "
        "This can be useful to create an ascii file or to show results as 'SHOW DATABASES' "
        "or 'EXPLAIN table'. To redirect the contents into an ascii file, the option -v0 "
        "is useful. To suppredd writing to an output file --null can be used.\n"
        "\n"
        "The default is to read the query from a file called rootify.sql. Except if a different "
        "filename is specified by the --file option or a query is given with --query.\n"
        "\n"
        "As a trick, the rootify.sql file can be made excutable (chmod u+x rootify.sql). "
        "If the first line contains '#!rootifysql', the script can be executed directly.\n"
        "\n"
        "Columns whose name start with @ are skipped. If you want them in your output file "
        "give them a name using AS, e.g. 'SELECT @A:=5 AS A'.\n"
        "\n"
        "You can use variables in your sql query like @MyVar and define them on the "
        "command line. In this example with --var.MyVar=5\n"
        "\n"
        "You can use environment definitions for substitutions in your SQL query. "
        "For example --env.TEST=5 would replace $TEST or ${TEST} in your query by 5."
        "If you specifiy one environmentvariable more than once, a list is created. "
        "For example --env.TEST=1 --env.TEST=2 --env.TEST=3 would substitute "
        "$TEST or ${TEST} by '1, 2, 3'. This is useful for the SQL `IN` keyword.\n"
        "\n"
        "Comments in the query-file can be placed according to the SQL standard inline "
        "/*comment*/ or introduced with # (shell script style) or -- (SQL style).\n"
        "\n"
        "In case of succes, 0 is returned, a value>0 otherwise.\n"
        "\n"
        "Usage: rootifysql [rootify.sql [rootify.root]] [-u URI] [-q query|-f file] [-i] [-o out] [-f] [-cN] [-t tree] [-vN]\n"
        "\n"
        ;
    cout << endl;
}

int main(int argc, const char* argv[])
{
    Time start;

    gROOT->SetBatch();

    Configuration conf(argv[0]);
    conf.SetPrintUsage(PrintUsage);
    SetupConfiguration(conf);

    if (!conf.DoParse(argc, argv))
        return 127;

    // ----------------------------- Evaluate options --------------------------
    const string   uri         = conf.Get<string>("uri");
    const string   out         = conf.Get<string>("out");
    const string   file        = conf.Get<string>("file");
    const string   tree        = conf.Get<string>("tree");
    const bool     force       = conf.Get<bool>("force");
    const bool     ignorenull  = conf.Get<bool>("ignore-null");
    const bool     update      = conf.Get<bool>("update");
    const bool     display     = conf.Get<bool>("display");
    const bool     noout       = conf.Get<bool>("null");
    const bool     nofill      = conf.Get<bool>("no-fill");
    const uint16_t verbose     = conf.Get<uint16_t>("verbose");
    const uint16_t compression = conf.Get<uint16_t>("compression");
    const string   delimiter   = conf.Get<string>("delimiter");
    const vector<string> _ignore = conf.Vec<string>("ignore");
    //const vector<Map> mymap    = conf.Vec<Map>("map");

    // -------------------------------------------------------------------------

    const auto vars = conf.GetWildcardOptions("var.*");

    vector<string> variables;
    for (const auto &var : vars)
        variables.emplace_back('@'+var.substr(4)+":="+Tools::Trim(conf.Get<string>(var)));

    // -------------------------------------------------------------------------

    if (verbose>0)
        cout << "\n--------------------- Rootify SQL ----------------------" << endl;

    string query  = conf.Get<string>("query");
    if (query.empty())
    {
        if (verbose>0)
            cout << "Reading query from file '" << file << "'." << endl;

        ifstream fin(file);
        if (!fin)
        {
            cerr << "Could not open '" << file << "': " << strerror(errno) << endl;
            return 1;
        }
        getline(fin, query, (char)fin.eof());
    }

    if (query.empty())
    {
        cerr << "No query specified." << endl;
        return 2;
    }

    // -------------------------------------------------------------------------

    const auto envs = conf.GetWildcardOptions("env.*");

    for (const auto &env : envs)
    {
        regex r1("\\$\\{"+env.substr(4)+"\\}");
        regex r2("\\$"+env.substr(4)+"\\b");

        const string repl = boost::join(conf.Vec<string>(env), ", ");

        query = regex_replace(query, r1, repl);
        query = regex_replace(query, r2, repl);
    }

    // -------------------------- Check for file permssion ---------------------
    // Strictly speaking, checking for write permission and existance is not necessary,
    // but it is convenient that the user does not find out that it failed after
    // waiting long for the query result
    //
    // I am using root here instead of boost to be
    // consistent with the access pattern by TFile
    TString path(noout?"/dev/null":out.c_str());
    gSystem->ExpandPathName(path);

    if (!noout)
    {
        FileStat_t stat;
        const Int_t  exist = !gSystem->GetPathInfo(path, stat);
        const Bool_t write = !gSystem->AccessPathName(path,  kWritePermission) && R_ISREG(stat.fMode);

        if ((update && !exist) || (update && exist && !write) || (force && exist && !write))
        {
            cerr << "File '" << path << "' is not writable." << endl;
            return 3;
        }

        if (!update && !force && exist)
        {
            cerr << "File '" << path << "' already exists." << endl;
            return 4;
        }
    }

    Time start2;

    // --------------------------- Connect to database -------------------------------------------------

    if (query.back()!='\n')
        query += '\n';

    if (verbose>0)
        cout << "Connecting to database..." << endl;


    Database connection(uri); // Keep alive while fetching rows

    // -------------------------- Set user defined variables -------------------
    if (variables.size()>0)
    {
        if (verbose>0)
            cout << "Setting user defined variables..." << endl;

        const string varset =
            "SET\n   "+boost::algorithm::join(variables, ",\n   ");

        try
        {
            connection.query(varset).execute();
        }
        catch (const exception &e)
        {
            cerr << varset << "\n\n";
            cerr << "SQL query failed:\n" << e.what() << endl;
            return 5;
        }

        if (verbose>2)
            cout << '\n' << varset << '\n' << endl;
    }

    // -------------------------- Request data from database -------------------
    if (verbose>0)
        cout << "Requesting data..." << endl;

    if (verbose>2)
        cout << '\n' << query << endl;

    const mysqlpp::UseQueryResult res =
        connection.query(query).use();

    // -------------------------------------------------------------------------

    if (verbose>0)
        cout << "Opening file '" << path << "' [compression=" << compression << "]..." << endl;

    // ----------------------------- Open output file --------------------------
    TFile tfile(path, update?"UPDATE":(force?"RECREATE":"CREATE"), "Rootify SQL", compression);
    if (tfile.IsZombie())
        return 6;

    // -------------------------------------------------------------------------

    // get the first row to get the field description
    mysqlpp::Row row = res.fetch_row();
    if (!row)
    {
        cerr << "Empty set returned... nothing to write." << endl;
        return 7;
    }

    if (verbose>0)
        cout << "Trying to setup " << row.size() << " branches..." << endl;

    if (verbose>1)
        cout << endl;

    const mysqlpp::FieldNames &l = *row.field_list().list;

    vector<double>  buf(l.size());
    vector<uint8_t> typ(l.size(),'n'); // n=number [double], d is used for DateTime

    UInt_t cols = 0;


    // -------------------- Configure branches of TTree ------------------------
    TTree *ttree = new TTree(tree.c_str(), query.c_str());

    size_t skipat  = 0;
    size_t skipreg = 0;
    for (size_t i=0; i<l.size(); i++)
    {
        const string t = row[i].type().sql_name();

        if (t.find("DATETIME")!=string::npos)
            typ[i] = 'd';
        else
            if (t.find("DATE")!=string::npos)
                typ[i] = 'D';
            else
                if (t.find("TIME")!=string::npos)
                    typ[i] = 'T';
                else
                    if (t.find("VARCHAR")!=string::npos)
                        typ[i] = 'V';
                    else
                        if (t.find("CHAR")!=string::npos)
                            typ[i] = 'C';

        bool found = false;
        for (const auto &pattern: _ignore)
        {
            if (regex_match(l[i], regex(pattern)))
            {
                found = true;
                typ[i] = '-';
                skipreg++;
                break;
            }
        }

        if (l[i][0]=='@')
        {
            typ[i] = '@';
            skipat++;
        }

        const bool use = l[i][0]!='@' && typ[i]!='V' && typ[i]!='C' && !found;

        if (verbose>1)
            cout << (use?" + ":" - ") << l[i].c_str() << " [" << t << "] {" << typ[i] << "}\n";

        if (use)
        {
            // string name = l[i];
            // for (const auto &m: mymap)
            //     name = boost::regex_replace(l[i], boost::regex(m.first), m.second);

            ttree->Branch(l[i].c_str(), buf.data()+i);
            cols++;
        }
    }
    // -------------------------------------------------------------------------

    if (verbose>1)
        cout << endl;
    if (verbose>0)
    {
        if (skipreg)
            cout << skipreg << " branches skipped due to ignore list." << endl;
        if (skipat)
            cout << skipat  << " branches skipped due to name starting with @." << endl;
        cout << "Configured " << cols << " branches.\nFilling branches..." << endl;
    }

    if (display)
    {
        cout << endl;
        cout << "#";
        for (size_t i=0; i<l.size(); i++)
            cout << ' ' << l[i].c_str();
        cout << endl;
    }

    // ---------------------- Fill TTree with DB data --------------------------
    size_t count = 0;
    size_t skip  = 0;
    do
    {
        count++;

        ostringstream sout;

        size_t idx=0;
        for (auto col=row.begin(); col!=row.end(); col++, idx++)
        {
            if (display)
            {
                if (idx>0)
                    sout << (delimiter.empty()?"\t":delimiter);
                sout << col->c_str();
            }

            if (!ignorenull && col->is_null())
            {
                skip++;
                break;
            }

            switch (typ[idx])
            {
            case 'd':
                buf[idx] = time_t((mysqlpp::DateTime)(*col));
                break;

            case 'D':
                buf[idx] = time_t((mysqlpp::Date)(*col));
                break;

            case 'T':
                buf[idx] = time_t((mysqlpp::Time)(*col));
                break;

            case 'V':
            case 'C':
            case '-':
            case '@':
                break;

            default:
                buf[idx] = atof(col->c_str());
            }
        }

        if (idx==row.size())
        {
            if (!nofill)
                ttree->Fill();

            if (display)
                cout << sout.str() << endl;
        }

        row = res.fetch_row();


    } while (row);

    // -------------------------------------------------------------------------

    if (display)
        cout << '\n' << endl;

    if (verbose>0)
    {
        cout << count << " rows fetched." << endl;
        if (skip>0)
            cout << skip << " rows skipped due to NULL field." << endl;

        cout << ttree->GetEntries() << " rows filled into tree." << endl;
    }

    ttree->Write();
    tfile.Close();

    if (verbose>0)
    {
        cout << "File closed.\n";
        cout << "Execution time: " << Time().UnixTime()-start.UnixTime() << "s\n";
        cout << "--------------------------------------------------------" << endl;
    }

    return 0;
}