#include <TFile.h>
#include <TTree.h>
#include <TString.h>
#include <TGeoManager.h>
#include <TSystem.h>

#include "ICOMETInput.hxx"
#include "ICOMETOutputMerged.hxx"
#include "ICOMETEvent.hxx"
#include "IFieldDescription.hxx"
#include "IOADatabase.hxx"
#include "IGeomIdManager.hxx"
#include "ISHAHashValue.hxx"
#include "ICOMETLog.hxx"
#include "IEntryList.hxx"
#include "IDataVector.hxx"

ClassImp(COMET::ICOMETOutputMerged);

// Define the name for the containter for the lists
const char* COMET::ICOMETOutputMerged::listContName = "event_selection_lists";

COMET::ICOMETOutputMerged::ICOMETOutputMerged(const char *fileName,
                               Option_t* opt, 
                               int compress) 
    : COMET::ICOMETOutput(fileName, opt, compress),
      fListContainer(NULL),
      fNumFiles(0),
      fGeomHash(COMET::ISHAHashValue()){}

COMET::ICOMETOutputMerged::~ICOMETOutputMerged(void) {
    if (IsOpen()) Close();
}

void COMET::ICOMETOutputMerged::AddFile(ICOMETInput* inFile, bool addEntryList){
    // Ensure the file is attached
    if (!IsAttached()) return;
    // Check that the geometry and EMField are consistent
    if (!CheckGeometry(inFile)) throw COMET::ECOMETOutputMergedDiffGeom(); 
    // TODO some exeception here
    if (!CheckEMField(inFile)) return; 
    // Get the file pointer and the tree inside it
    TFile* thisInputFile = inFile->GetFilePointer();
    TTree* fileTree = dynamic_cast<TTree*>(thisInputFile->Get("COMETEvents"));
    AddToTree(fileTree);
    if (addEntryList)
        AddEntryList(inFile->GetFilename(), fileTree->GetEntries());
    // Increment the number of files added
    fNumFiles++;
}

bool COMET::ICOMETOutputMerged::CheckGeometry(ICOMETInput* inFile){
    // Check if the geometry has not been written yet
    if (!GeometryWritten()) {
        // Write the current geometry to file
        WriteGeometry(IOADatabase::Get().UpdateGeometry());
        // Store the geometry hash of this geometry
        fGeomHash = IOADatabase::Get().GeomId().GetHash();
        return true;
    }
    // If the geometry already exists on file, ensure the geometry hash of the
    // currently loaded file matches the stored geometry hash
    return fGeomHash.Equivalent(IOADatabase::Get().GeomId().GetHash()); 
}

bool COMET::ICOMETOutputMerged::CheckEMField(ICOMETInput* inFile){
    //Check that we have not already written geometry
    if (!EMFieldWritten()) WriteEMField(IOADatabase::Get().EMFieldDescription());
    COMETLog("NO EMFIELD CHECKER IN PLACE YET");
    //TODO fill this in with checkers to ensure it is the same EMField as the
    //input file
    return true;
}

void COMET::ICOMETOutputMerged::AddToTree(TTree* inTree){
    // Merge this file in
    TList listInputTree;
    listInputTree.Add(inTree);
    fEventTree->Merge(&listInputTree);
    // TODO some exception or boolean flag here if merge does not work
}

void COMET::ICOMETOutputMerged::CreateEntryListContainer(){
    // Check if we already have this entry list
    if (!fListContainer) {
        fListContainer = dynamic_cast<COMET::IDataVector*>(fFile->Get(listContName));
        // If it does not exist, make one
        if (!fListContainer) {
            fListContainer = new COMET::IDataVector(listContName, 
                                                    "Lists of selected event indexes");
            COMETTrace("New event list container");
        }
    }
    return;
}

void COMET::ICOMETOutputMerged::AddEntryList(const char* fileName, 
                                             int nEntries,
                                             int entryID){
    // Ensure we are attached
    if (!IsAttached()) return;
    // Add the container for entry lists if it does not exist
    CreateEntryListContainer();
    // Use the base file name for naming the list
    const char* baseFileName = gSystem->BaseName(fileName);
    // Check if the list container already has the entry list by that name
    if (fListContainer->Has<COMET::IEntryList>(baseFileName)){
        COMETError("IEntryList of Name : " << baseFileName << " already exists in " 
                   "file " << GetName());
        throw COMET::ECOMETOutputAddListFailed();
    }
    // Count the entries in it, defaulting the entry ID to the number of files 
    // we've already added
    if (entryID == -1) entryID = fNumFiles;
    IEntryList* thisList = new IEntryList(baseFileName,
        TString::Format("EventFile_%d",fNumFiles).Data());
    Int_t first_entry = GetEventsWritten();
    Int_t last_entry = first_entry + nEntries;
    fEventsWritten = last_entry;
    for (int entry = first_entry; entry < last_entry; entry++){
        thisList->GetList().Enter(entry);
    }
    // Write this list to the file
    fListContainer->AddDatum(thisList);
}

void COMET::ICOMETOutputMerged::WriteEntryListContainer(){
    // Ensure we are attached
    if (!IsAttached()) return;
    // Write the list container if it was added
    if (fListContainer) fListContainer->Write();
    // Lists written
    if (fListContainer) {
        COMETLog("Lists written : ");
        fListContainer->ls();
    }
}

void COMET::ICOMETOutputMerged::Close(Option_t* opt) {
    // Make sure we write the entry list once
    WriteEntryListContainer();
    // Close as usual
    COMET::ICOMETOutput::Close(opt);
}