#!/usr/bin/env python

from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

import sys
import re
import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
from rat import ROOT

def help_ratinfo():
    '''
    Function to print information about usage of this script.
    TODO:
    - use proper argument parsing, rather than using sys.argv
    '''
    print('''Usage:
  ratinfo <info> file.root [item2 item3...]

  where <info> is
    log    Print log from RAT run which generated first file.
    macro  Print RAT macro which generated first file.
    meta   Print metadata (versions, overrides) for first file.
           If more files listed, compare metadata with the first file,
           listing matches with a +, and mismatches with a -.
    dbkeys List keys in db TMap in first file.
    db     Print tables in db TMap, with regular expressions to
           match keys as further arguments.
''')

def dump_metadata(meta):
    '''
    Function to get a variety of metadata from the specified file.
    '''
    npass = meta.GetPassCount()
    if npass > 0:
        for ipass in range(0, npass):
            print('Pass {}:'.format(ipass))
            print('  Host                {} ({})'.format(meta.GetHostNames()[ipass],
                                                         meta.GetHostSystems()[ipass]))
            print('  RAT version         {}'.format(meta.GetVersions()[ipass]))
            print('  RAT revision        {}'.format(meta.GetRevisions()[ipass]))
            print('  ROOT version        {}'.format(meta.GetRootVersions()[ipass]))
            print('  GEANT4 version      {}'.format(meta.GetGeant4Versions()[ipass].strip()))
            print('  Contains MC         {}'.format(meta.GetContainsMCFlags()[ipass]))
            print('  Contains Data       {}'.format(meta.GetContainsDataFlags()[ipass]))
            print('  Events Generated    {}'.format(meta.GetEventsGeneratedCounts()[ipass]))
            print('  Events Stored       {}'.format(meta.GetEventsStoredCounts()[ipass]))

            db = meta.GetMetaDB(ipass)
            for ioc in range(0, db.GetOverrideCommandCount()):
                oc = db.GetOverrideCommand(ioc)
                print('  {:<19} {}'.format(oc.first.strip(), oc.second.strip()))
            for iof in range(0, db.GetFileCount()):
                print('  File {}'.format(db.GetFile(iof)))
    else:
        sys.exit("No metadata for any passes in ROOT file.")

def compare_metadata(meta, fname):
    '''
    Function to compare metadata.
    Compares metadata to the specified file.
    '''
    g = ROOT.TFile.Open(fname)
    if not g:
        return '0'
    if g.IsZombie():
        return '0'
    gmeta = g.Get('meta')
    if gmeta:
        npass = meta.GetPassCount()
        if npass != gmeta.GetPassCount():
            return '-'
        for ipass in range(0, npass):
            if meta.GetVersions()[ipass] != gmeta.GetVersions()[ipass]:
                return '-'
            if meta.GetRevisions()[ipass] != gmeta.GetRevisions()[ipass]:
                return '-'
            if meta.GetRootVersions()[ipass] != gmeta.GetRootVersions()[ipass]:
                return '-'
            if meta.GetGeant4Versions()[ipass] != gmeta.GetGeant4Versions()[ipass]:
                return '-'
            if meta.GetContainsMCFlags()[ipass] != gmeta.GetContainsMCFlags()[ipass]:
                return '-'
            if meta.GetContainsDataFlags()[ipass] != gmeta.GetContainsDataFlags()[ipass]:
                return '-'
            if meta.GetEventsStoredCounts()[ipass] != gmeta.GetEventsStoredCounts()[ipass]:
                return '-'
            if meta.GetEventsGeneratedCounts()[ipass] != gmeta.GetEventsGeneratedCounts()[ipass]:
                return '-'

            db = meta.GetMetaDB(ipass)
            gdb = gmeta.GetMetaDB(ipass)
            if db.GetOverrideCommandCount() != gdb.GetOverrideCommandCount():
                return '-'
            if db.GetFileCount() != gdb.GetFileCount():
                return '-'
            for ioc in range(0, db.GetOverrideCommandCount()):
                oc = db.GetOverrideCommand(ioc)
                match = False
                for i in range(0, gdb.GetOverrideCommandCount()):
                    goc = gdb.GetOverrideCommand(i)
                    if oc.first == goc.first:
                        if oc.second == goc.second:
                            match = True
                            break
                        return '-'
                if not match:
                    return '-'
            for iof in range(0, db.GetFileCount()):
                of = db.GetFile(iof)
                match = False
                for i in range(0, gdb.GetFileCount()):
                    if of == gdb.GetFile(i):
                        match = True
                        break
                if not match:
                    return '-'
    else:
        return '-'
    return '+'

if len(sys.argv) < 3:
    if len(sys.argv) == 1:
        help_ratinfo()
        sys.exit(2)
    else:
        sys.exit("Invalid number of arguments (run \"ratinfo\" for usage).")

# we will detect and report errors ourselves
ROOT.gErrorIgnoreLevel = ROOT.kFatal

f = ROOT.TFile.Open(sys.argv[2])
if not f:
    sys.exit("No such file {}.".format(sys.argv[2]))

if f.IsZombie():
    sys.exit("Failed to open file {}.".format(sys.argv[2]))

if sys.argv[1] == 'log':
    if len(sys.argv) != 3:
        help_ratinfo()
        sys.exit(2)
    log = f.Get("log")
    if log:
        print(log.GetString())
    else:
        sys.exit("No log present in ROOT file.")

elif sys.argv[1] == 'macro':
    if len(sys.argv) != 3:
        help_ratinfo()
        sys.exit(2)
    macro = f.Get('macro')
    if macro:
        print(macro.GetString())
    else:
        sys.exit("No macro present in ROOT file.")

elif sys.argv[1] == 'dbkeys':
    if len(sys.argv) != 3:
        help_ratinfo()
        sys.exit(2)
    db = f.Get('db')
    if db:
        idb = db.MakeIterator()
        while idb.Next():
            print(idb.Key())
    else:
        sys.exit("No db present in ROOT file.")

elif sys.argv[1] == 'db':
    db = f.Get('db')
    if db:
        idb = db.MakeIterator()
        for iarg in range(3, len(sys.argv)):
            p = re.compile(sys.argv[iarg])
            idb.Reset()
            while idb.Next():
                if p.match(idb.Key().GetName()):
                    print('{} = {}'.format(idb.Key(), idb.Value()))

elif sys.argv[1] == 'meta':
    meta = f.Get('meta')
    if meta:
        dump_metadata(meta)
        if len(sys.argv) > 3:
            print('+ {}'.format(sys.argv[2]))
            for iarg in range(3, len(sys.argv)):
                c = compare_metadata(meta, sys.argv[iarg])
                print('{} {}'.format(c, sys.argv[iarg]))
    else:
        sys.exit("No metadata present in ROOT file.")

else:
    sys.exit("Invalid argument for info: {} (run \"ratinfo\" for usage).".format(sys.argv[1]))