#include <iostream>
#include <fstream>
#include <sstream>
#include <time.h>
#include <unistd.h>
#include <getopt.h>

#include <TROOT.h>
#include <TRint.h>
#include <TSystem.h>

#include "ICOMETLog.hxx"
#include "ICOMETInput.hxx"
#include "ICOMETOutputMerged.hxx"
#include "IInputFileList.hxx"

void usage(void) {
    std::cout 
    <<"Usage: oaEvent_index-file.exe [options] file\n"
      "  where file specifier is the path to an input oaEvent file\n"
      "  <filepath>/<oaeventfile>\n\n"
      "The specified file has IEntryLists added to it.  This is helpful for\n"
      "splitting a large file into smaller segments that the event loop can\n"
      "read seperately. It creates a similar file to concatenating many\n"
      "smaller oaEvent files into one large file.\n"
      "Options: \n"
      "    -N    Number of segments to split the file into.  If this is not a\n"
      "          factor of the overall number of events in the input file,\n"
      "          the last segment in the file will be smalller than the rest.\n"
      "          NOTE: default is 10\n\n"
      "    -o    Name of the output file" << std::endl;
}

int main(int argc, char **argv) {
    // Get the default options
    std::string fOutFileName = "indexed_file.root";
    int nSegments = 10;
    // Lets count how many events we need to segment
    int nEventsToSegment = 0;

    // Set the options
    int c;
    while ((c = getopt(argc, argv, "N:o:h")) != -1)
        switch (c){
            // Set the number of segments to split the file into
            case 'N':
                nSegments = atoi(optarg);
                break;
            // Choose the output name
            case 'o':
                fOutFileName= optarg;
                break;
            // Print the help menu
            case 'h':
                usage();
                return 0;
            // Let getopt deal with the rest
            case '?':
                return 1;
            default:
                abort();
        }

    // Ensure we only have on option left
    if ((argc - optind) > 1){
        COMETError("Error: Only one positional argument allowed! Found: " 
                   << argc - optind << std::endl);
        return 1;
    }
    // Prepare the input files
    const char* fInputFilePath = argv[optind];

    //Configure COMETLog levels from input config file
    TString configPath = "${PWD}/cometlog.config";
    gSystem->ExpandPathName(configPath);
    COMET::ICOMETLog::Configure(configPath.Data());
    
    // Open the output
    COMET::ICOMETOutputMerged* mergedFile = 
        new COMET::ICOMETOutputMerged(fOutFileName.c_str(), "CREATE");

    // Get the next input file
    COMET::ICOMETInput fInputFile = COMET::ICOMETInput(fInputFilePath);
    // Add this input file, but no entry list for it, i.e. clone the file
    mergedFile->AddFile(&fInputFile, false);

    // Add the number of entries in total
    nEventsToSegment = fInputFile.GetEventsInFile();
    // Figure out how many we need in each split
    int nEventsInSegment = nEventsToSegment / nSegments;
    int nEventsRemain = nEventsToSegment % nSegments;
    if (nEventsRemain != 0) ++nEventsInSegment;
    // Set the amount for each split
    int nEventsPerSegment[nSegments];
    // Fill all but the last segment with nEventsInSegment
    for (int i=0; i < nSegments-1; i++) nEventsPerSegment[i] = nEventsInSegment;
    // Fill the last segment with the remaining events
    nEventsPerSegment[nSegments-1] = nEventsToSegment;
    nEventsPerSegment[nSegments-1] -= nEventsInSegment * (nSegments-1);

    // Add all the entry lists
    std::stringstream iName;
    for (int i=0; i < nSegments; i++){
        iName << "IndexedFile_" << i; 
        mergedFile->AddEntryList(iName.str().c_str(), nEventsPerSegment[i], i);
        iName.str("");
    }

    // Close the merged file
    mergedFile->Close();
    // Recored the number of events merged
    COMETLog("Number of Events Merged : " << nEventsToSegment);
    for (int i=0; i < nSegments; i++)
        COMETLog("  Number of Events in Segment "
                 << i << " = " << nEventsPerSegment[i]);
    return 0;
}