#!/usr/bin/env python

import sys
import os
import string
import urllib
import datetime

#constants
Str_script = "get_files.py : "

#Data sites
Str_dataset_link = "/dbs_discovery/getData"
Str_dbs_runs     = "https://cmsweb.cern.ch/dbs_discovery/getRunsFromRange"
Str_dbs_lfn      = "https://cmsweb.cern.ch/dbs_discovery/getLFN_txt"
Str_dbs_search   = "https://cmsweb.cern.ch/dbs_discovery/aSearchShowAll"

#parameters
str_run       = "0"
str_dataset   = "/Cosmics/Commissioning08-v1/RAW"
str_extention = "py"
str_output    = ""

#variables
int_allrun = 0
int_run = 0

#Definition of printing formats
def head(ext,file):
    if ext == "cff" :
        file.write("replace PoolSource.fileNames = {\n")
        return
    if ext == "py" :
        file.write("import FWCore.ParameterSet.Config as cms\n\n")
        file.write("maxEvents = cms.untracked.PSet( input = cms.untracked.int32(-1) )\n")
        file.write("readFiles = cms.untracked.vstring()\n")
        file.write("secFiles = cms.untracked.vstring()\n") 
        file.write("source = cms.Source (\"PoolSource\",fileNames = readFiles, secondaryFileNames = secFiles)\n")
        file.write("readFiles.extend( [\n")
        return
    return
    
def line (ext,rootfile,file):
    if ext == "cff" :
        file.write("\t'"+rootfile+"',\n")
        return
    if ext == "py" :
        file.write("\t'"+rootfile+"',\n")
        return
    if ext == "txt" :
        file.write(rootfile+"\n")
        return
    return

def tail (ext,file):
    if ext == "cff" :
        file.write("}\n}")
        return
    if ext == "py" :
        file.write("])")
        return
    return


#Getting parameters
for int_count in range(len(sys.argv)):
    #-r <Run> if none, use last for given dataset
    if sys.argv[int_count] == "-r" :
        if int_count < len(sys.argv)-1 : 
            str_run  = sys.argv[int_count+1]
            continue
    #-d <Dataset>
    if sys.argv[int_count] == "-d" :
        if int_count < len(sys.argv)-1 : 
            str_dataset  = sys.argv[int_count+1]
            continue
    #-D <Dataset> force all runs with dataset
    if sys.argv[int_count] == "-D" :
        if int_count < len(sys.argv)-1 : 
            str_dataset  = sys.argv[int_count+1]
            int_allrun = 1
            continue
    #-f <output file name> if none <dataset>_Run_<Run>_cff.py
    if sys.argv[int_count] == "-f" :
        if int_count < len(sys.argv)-1 : 
            str_output  = sys.argv[int_count+1]
            continue
    #-cff
    #-py
    #-txt
    if sys.argv[int_count] == "-cff":
       str_extention = "cff"
    if sys.argv[int_count] == "-py":
       str_extention = "py"
    if sys.argv[int_count] == "-txt":
       str_extention = "txt"
    
try:
    int_run = int(str_run)
except:
    print "%s The run number %s is not an integer" %(Str_script, str_run)
    sys.exit(-1)

user_input = 'find run.number,dataset where dataset=' + str_dataset

if int_run > 0 : 
    user_input = user_input+" and run=" + str_run
    print "%sSearching in dbs for %s, %s" %(Str_script,str_run,str_dataset)
else :
    print "%sSearching in dbs for %s" %(Str_script,str_dataset)

str_dbs_params   = urllib.urlencode({'case':'on','cff':'0','caseSensitive':'on','grid':'0','fromRow':'0','xml':'0','sortName':'Runs_RunRumber','dbsInst':'cms_dbs_prod_global','html':'1','limit':'-1','sortOrder':'desc','output':'run.number,dataset','tabCol':'','userMode':'user','method':'dbsapi','userInput':user_input})
file_dbsRuns     = urllib.urlopen(Str_dbs_search, str_dbs_params)
#getting the last run
int_lastRun = 0
int_totRuns = 0
for str_Runs in file_dbsRuns.readlines():
    if str_Runs.find(str_dataset) >= 0 :
        run_list = str_Runs.split(', '+ str_dataset)
        for run in run_list :
            if run.find("<br/>") >= 0 :  
                run=run.split("<br/>")[1]
            try:
                int(run)
            except:
                continue
            int_totRuns = int_totRuns + 1
            if int(run) > int_lastRun :
                int_lastRun = int(run)
if int_totRuns == 0 :
    print "%sNo entries found" %(Str_script)
    sys.exit(-1)
print "%sFound %i entries" %(Str_script, int_totRuns)
if int_run == 0 and int_allrun == 0:
    print "%sFound last run: %i" %(Str_script, int_lastRun)
    int_run = int_lastRun
    file_run = open("_tmp_last_run_", 'w')
    file_run.write(str(int_lastRun))
    
str_run = str(int_run)

if str_output == "" :
    str_output = str_dataset
    if int_allrun == 1:
        str_output = str_output+"_AllRuns"
    else:
        str_output = str_output+"_Run_"+str(int_run)

    if str_extention == "cff":
        str_output = str_output+".cff"
    elif str_extention == "py":
        str_output = str_output+"_cff.py"
    elif str_extention == "txt":
        str_output = str_output+".txt"
    str_output = str_output.replace('/','_')
    str_output = str_output.replace('-','_')
    str_output = str_output.replace('__','_')


#Getting .cff (.py) file list
user_input = 'find file,lumi where dataset=' + str_dataset
if int_run > 0 : 
    user_input = user_input+" and run=" + str_run
    print "%sSearching in dbs for %s, %s" %(Str_script,str_run,str_dataset)
else :
    print "%sSearching in dbs for %s" %(Str_script,str_dataset)
user_input = user_input+" order by lumi asc"

str_lfn_params = urllib.urlencode({'case':'on','cff':'0','caseSensitive':'on','userInput':user_input,'grid':'0','fromRow':'0','xml':'0','sortName':'','dbsInst':'cms_dbs_prod_global','html':'1','limit':'-1','sortOrder':'asc','output':'file,lumi','tabCol':'','userMode':'user','method':'dbsapi'})
file_lfnRuns   = urllib.urlopen(Str_dbs_search, str_lfn_params)

f = open(str_output, 'w')
head(str_extention,f)

str_dataset_line = ""
for str_dbsLFN in file_lfnRuns.readlines():
    if str_dbsLFN.find(".root") >=0 :
        str_dataset_line = str_dbsLFN.split("<br/>")
        for str_lfn in str_dataset_line :
            if str_lfn.find(".root")>=0:
                line(str_extention,str_lfn.split(',')[0],f)

if  str_dataset_line == "" :
    print "%sNo .root files found. Maybe dbs is down." %(Str_script)
    tail(str_extention,f)
    sys.exit(-1)
    
tail(str_extention,f)

sys.exit(0)
