#!/usr/bin/env jython

import stilts
import datetime
import re
import sys
import os
import md5
import shutil
import logging
from subprocess import call, PIPE, Popen
from java.lang import *
from java.sql import *
from xml.etree.ElementTree import ElementTree, Element, tostring

from scratchdb import ScratchDb
from gpc1db import Gpc1Db
from ipptopspsdb import IppToPspsDb
from fits import Fits

'''
Abstract base class of all batch types.
'''
class Batch(object):

    '''
    Constructor
    '''
    def __init__(self, 
                 logger, 
                 config,
                 skychunk,
                 gpc1Db,
                 ippToPspsDb,
                 scratchDb,
                 id,
                 batchID,
                 batchType, 
                 fits): 

        self.readHeader = False
        self.config = config
        self.skychunk = skychunk
        self.fits = fits

        # set up logging
        self.logger = logger
        self.logger.infoSeparator()
        self.logger.debug("Batch class constructor")

        # set up class variables
        self.id = id
        self.batchID = batchID
        self.gpc1Db = gpc1Db
        self.ippToPspsDb = ippToPspsDb
        self.scratchDb = scratchDb
        self.batchType = batchType;
        self.skipBatch = False; # set to true if there is a non-fatal error
        self.pspsVoTableFilePath = self.config.configDir + "tables." + batchType + ".vot"

        # check FITS file is ok
        if self.fits: 
            self.header = self.fits.getPrimaryHeader()

            # now check that the fits header is readable
            if not self.header: 
                logger.errorPair("Could not read FITS PHU for id", "%d" % id)
                raise
        else:
            logger.errorPair("missing fits item?", str(id))

        # define the dvo table names
        self.scratchDb.setDvoTableNames()

        # TODO
        self.tablesToExport = []

        if self.skychunk.survey != "":
            self.surveyID = self.scratchDb.getSurveyID(self.skychunk.survey)
        else:
            self.surveyID = -1;
       
        # create datastore object
        self.datastore = Datastore(self.logger, self.skychunk, self.ippToPspsDb)

        # get local storage location from skychunk
        self.batchName = Batch.getNameFromID(self.batchID)
        self.subDir = Batch.getSubDir(
                self.skychunk.basePath,
                self.batchType, 
                self.skychunk.dvoLabel)

        self.localOutPath = Batch.getOutputPath(
                self.skychunk.basePath,
                self.batchType, 
                self.skychunk.dvoLabel,
                self.batchID)

        if not os.path.exists(self.localOutPath): os.makedirs(self.localOutPath)

        # store today's date
        now = datetime.datetime.now();
        self.dateStr = now.strftime("%Y-%m-%d")

        # create DVO tables (scratch database)
        ### XXX self.scratchDb.createDvoTables() <-- remove this

        # dump stuff to the log
        self.logger.infoTitle("New " + self.batchType + " batch")
        self.logger.infoPair("Batch name", self.batchName)
        self.logger.infoPair("Survey ID", "%d" % self.surveyID)

        if self.fits:
            self.logger.infoPair("Input FITS file", self.fits.getOriginalPath())
            self.logger.infoPair("Input FITS primary header", "%s cards found" % self.fits.getPrimaryHeaderCardCount())

        self.logger.infoPair("Output path", self.localOutPath)

    
    def testprint(self):
      print "here ", self.printline
      self.printline += 1

    '''
    Static method to generated batch name from batch ID
    '''
    @staticmethod
    def getNameFromID(id):
       return "B%08d" % id

    '''
    Static method to generated sub directory for batch on local disk
    '''
    @staticmethod
    def getSubDir(basePath, batchType, dvoLabel):
       return basePath + "/" + batchType + "/" + dvoLabel

    '''
    Static method to generated output path for batch
    '''
    @staticmethod
    def getOutputPath(basePath, batchType, dvoLabel, id):
       return Batch.getSubDir(basePath, batchType, dvoLabel) + "/" + Batch.getNameFromID(id)

    '''
    Static method to get name of tar file
    '''
    @staticmethod
    def getTarFile(id):
        return Batch.getNameFromID(id) + ".tar"

    '''
    Static method to get path of tar file
    '''
    @staticmethod
    def getTarPath(basePath, batchType, dvoLabel, id):
        return Batch.getSubDir(basePath, batchType, dvoLabel) + "/" + Batch.getTarFile(id)

    '''
    Static method to get name of tarball file
    '''
    @staticmethod
    def getTarballFile(id):
        return Batch.getTarFile(id) + ".gz"

    '''
    Static method to get path of tarball file
    '''
    @staticmethod
    def getTarballPath(basePath, batchType, dvoLabel, id):
        return Batch.getSubDir(basePath, batchType, dvoLabel) + "/" + Batch.getTarballFile(id)

    '''  
    Static method to delete this batch from local disk
    '''
    @staticmethod
    def deleteFromDisk(logger, basePath, batchType, dvoLabel, id):
       
        tarballPath = Batch.getTarballPath(basePath, batchType, dvoLabel, id)
        dirPath = Batch.getOutputPath(basePath, batchType, dvoLabel, id)

        if not os.path.exists(tarballPath): 
            logger.errorPair("Could not find", tarballPath)
            return 1

        try:
            os.remove(tarballPath)
            logger.infoPair("Deleted", tarballPath)
        except:
            logger.errorPair("Could not delete", tarballPath)
            return 0
   
        return 1

    '''
    Destructor
    '''
    def __del__(self):

        self.logger.debug("Batch destructor")

    '''
    Looks for a dictionary entry. If it is not there, then replaces it with the supplied default
    if, and only if, in test mode
    '''
    def safeDictionaryAccessWithDefault(self, header, key, default):

         value = self.safeDictionaryAccess(header, key)

         if value != "NULL": return value
         else:
             if not self.config.test: return "NULL"
             header[key] = default
             self.logger.infoPair("Hardcoding " + key + " to", header[key])
             return header[key]

    '''
    Returns the string value from this dictionary or else "NULL"
    '''
    def safeDictionaryAccess(self, header, key):

         if key in header: return header[key]
         else:
             self.logger.errorPair("Missing header field", key)
             return "NULL"

    '''
    Returns the string keyword value from this header or else "NULL"
    '''
    def getKeyValue(self, header, key):

         if key in header: return header[key]
         else:
             self.logger.errorPair("Missing header field", key)
             return "NULL"

    '''
    Returns the float keyword value from this header or else "NULL"
    '''
    def getKeyFloat(self, header, format, key):

         if key in header: 
             value = float(format % float(header[key]))
         else:
             self.logger.errorPair("Missing header field", key)
             value = float(format % -999.9)

         return value

    '''
    Returns the string keyword value from this header or else "NULL"
    '''
    def getKeyInt(self, header, defValue, key):

         if key in header: 
             value = str(int(header[key]))
         else:
             self.logger.errorPair("Missing header field", key)
             value = str(defValue)

         return value

    '''
    Dumps out the list of keys in a header
    '''
    def dumpKeys(self, header):
        for key in header:
            self.logger.infoPair("", str(key))
        return 0

    '''
    Writes the batch manifest file
    '''
    def writeBatchManifest(self):

        outPath = self.localOutPath + "/BatchManifest.xml"
        tmpPath = self.localOutPath + "/tmp.xml"
        self.logger.infoPair("Creating manifest", outPath)
        root = Element('manifest')

        # batch information
        root.attrib['name'] = self.batchName
        root.attrib['type'] = self.batchType
        root.attrib['timestamp'] = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") 
        root.attrib['survey'] = self.skychunk.pspsSurvey

        # min/max object IDs
        try: self.minObjID
        except: pass
        else: root.attrib['minObjId'] = str(self.minObjID)
        try: self.maxObjID
        except: pass
        else: root.attrib['maxObjId'] = str(self.maxObjID)

        # file information
        child = Element('file')
        root.append(child)
        child.attrib['name'] = self.outputFitsFile

        # now create doc and write to file
        file = open(tmpPath, 'w')
        ElementTree(root).write(file)
        file.close()

        # clunky way to prettify XML
        try:
            cmd = "xmllint --format " + tmpPath + " > " + outPath
            p = Popen(cmd, shell=True, stdout=PIPE)
            p.wait()
            if p.returncode != 0:
                self.logger.errorPair("error running xmllint for ", tmpPath)
                raise
        except:
            self.logger.errorPair("Popen failed", cmd)
            return False
        finally: 
            os.remove(tmpPath)

        return True


    '''
    Publishes this batch to the datastore
    '''
    def tarAndZip(self):
      
        # set up filenams and paths
        tarFile = Batch.getTarFile(self.batchID)
        tarPath = Batch.getTarPath(self.skychunk.basePath, self.batchType, self.skychunk.dvoLabel, self.batchID)
        tarballFile = Batch.getTarballFile(self.batchID)
        tarballPath = Batch.getTarballPath(self.skychunk.basePath, self.batchType, self.skychunk.dvoLabel, self.batchID)

        # tar directory
        cmd = "tar -cvf " + tarPath + " -C " + self.subDir + " " + self.batchName
        self.logger.infoPair("Creating tar archive", cmd)
        p = Popen(cmd, shell=True, stdout=PIPE)
        p.wait()
        if p.returncode != 0:
            self.logger.errorPair("tar command", "failed")
            return False

        # zip tar archive
        cmd = "gzip -c " + tarPath + " > " + tarballPath
        self.logger.infoPair("Compressing tar archive", cmd)
        p = Popen(cmd, shell=True, stdout=PIPE)
        p.wait()
        if p.returncode != 0:
            self.logger.errorPair("gzip command", "failed")
            return False

        # delete tar file and original directory
        os.remove(tarPath)
        shutil.rmtree(self.localOutPath)

        return True

    '''
    Static method to publish a batch to the datastore
    '''
    @staticmethod
    def publishToDatastore(datastore, batchID, batchName, subDir, tarballFile):
        datastore.publish(batchName, subDir, tarballFile, "tgz")

    '''
    Sets min and max obj ID using the provided table, or list of tables
    '''
    def setMinMaxObjID(self, tables):

        first = True

        self.totalDetections = 0
        for table in tables:

            sql = "SELECT MIN(objID), MAX(objID), COUNT(objID) FROM " + table
            rs = self.scratchDb.executeQuery(sql)
            rs.first()

            self.totalDetections = self.totalDetections + rs.getLong(3)

            if first:
                self.minObjID = rs.getLong(1)
                self.maxObjID = rs.getLong(2)
            else:
                if rs.getLong(1) < self.minObjID: self.minObjID = rs.getLong(1)
                if rs.getLong(2) > self.maxObjID: self.maxObjID = rs.getLong(2)

            first = False
            rs.close()

        self.ippToPspsDb.updateDetectionStats(self.batchID, self.minObjID, self.maxObjID, self.totalDetections)
        self.logger.infoPair("Total detections", "%ld" % self.totalDetections)
        self.logger.infoPair("Min objID", "%ld" % self.minObjID)
        self.logger.infoPair("Max objID", "%ld" % self.maxObjID)


    '''                                                                                            
    Sets min and max obj ID using the provided table, or list of tables                            
    '''
    def setMinMaxDiffObjID(self, tables):
        first = True
        self.totalDetections = 0
        for table in tables:
            sql = "SELECT MIN(diffObjID), MAX(diffObjID), COUNT(diffObjID) FROM " + table
            rs = self.scratchDb.executeQuery(sql)
            rs.first()
            self.totalDetections = self.totalDetections + rs.getLong(3)
            if first:
                self.minObjID = rs.getLong(1)
                self.maxObjID = rs.getLong(2)
            else:
                if rs.getLong(1) < self.minObjID: self.minObjID = rs.getLong(1)
                if rs.getLong(2) > self.maxObjID: self.maxObjID = rs.getLong(2)
            first = False
            rs.close()
        self.ippToPspsDb.updateDetectionStats(self.batchID, self.minObjID, self.maxObjID, self.totalDetections)
        self.logger.infoPair("Total detections", "%ld" % self.totalDetections)
        self.logger.infoPair("Min objID", "%ld" % self.minObjID)
        self.logger.infoPair("Max objID", "%ld" % self.maxObjID)



    '''
    Creates PSPS tables from VOTable descriptions
    '''   
    def createEmptyPspsTables(self):
        self.pspsTables = stilts.treads(self.pspsVoTableFilePath)
        for table in self.pspsTables:
            self.logger.debug("Creating PSPS table: " + table.name)
            self.logger.infoPair("creating psps table ",table.name)
            table.write(self.scratchDb.url + '#' + table.name)
            self.tablesToExport.append(table.name)

        return True

    '''
    Subclass should implement this to index IPP tables
    '''   
    def indexIppTables(self):
        self.logger.warn("indexIppTables not implemented")


    '''
    Alter PSPS tables (Subclass implements this)
    '''   
    def alterPspsTables(self):
        self.logger.warn("alterPspsTables not implemented")


    '''
    Imports IPP tables from FITS file

    Accepts a regular expression filter so not all tables need to be imported
    '''
    def importIppTables(self, columns="*", filter=""):
        if self.config.useOhanaMYSQLtoFITS:
            if not self.importIppTablesOhana():
                self.logger.errorPair("problem running", "importIppTablesOhana")
                return False
        else:
            if not self.importIppTablesStilts(columns, filter):
                self.logger.errorPair("problem running", "importIppTablesStilts")
                return False

        return True

    '''
    Imports IPP tables from FITS file
    '''
    def importIppTablesOhana(self):
        
        self.logger.infoPair("Importing IPP tables", "using fits_to_mysql")

        count = 0
        self.tablesLoaded = []

        fileName = self.fits.getPath()

        # this option is equivalent to the stilts "addcol table_index $0" call
        fitsOptions  = " -sequence-column table_index"
        fitsOptions += " -sequence-column-start 1"

        ## *** actually write the FITS data to mysql
        self.scratchDb.loadFITStoMYSQL(fileName, fitsOptions)

        ## use stilts to get a list of the tables from the header
        try:
            tables = stilts.treads(fileName)
        except:
            self.logger.errorPair("STILTS could not import from", fileName)
            return False

        for table in tables:
            # example of listing all columns in a table:
            ## print "-----table: " + table.name
            ## columns = table.columns()
            ## for column in columns:
            ##     print column
            ## 
            ## print "-----"

            self.logger.debugPair("Read IPP table", table.name)

            self.tablesLoaded.append(table.name)
            count += 1

        self.logger.infoPair("Done. Imported", "%d tables" % count)
        self.indexIppTables()

        return True

    '''
    Imports IPP tables from FITS file

    Accepts a regular expression filter so not all tables need to be imported
    '''
    def importIppTablesStilts(self, columns="*", filter=""):
        
        self.logger.infoPair("Importing tables (stilts) with filter", filter)
        
        print "trying to read ", self.fits.getPath()

        try:
            tables = stilts.treads(self.fits.getPath())
        except:
            self.logger.errorPair("STILTS could not import from", self.fits.getPath())
            return False

        self.tablesLoaded = []

        print "read tables from ", self.fits.getPath()

        count = 0
        for table in tables:

            # example of listing all columns in a table:
            ## print "-----table: " + table.name
            ## myColumns = table.columns()
            ## for column in myColumns:
            ##     print column
            ## print "-----"

            print "import smf table ", table.name

            match = re.match(filter, table.name)
            if not match: continue
            self.logger.debugPair("Reading IPP table", table.name)
            table = stilts.tpipe(table, cmd='addcol table_index $0')

            table = stilts.tpipe(table, cmd='explodeall')

            # print "read smf table ", table
            self.logger.infoPair("read smf table", table.name)

            # drop any previous tables before import
            self.scratchDb.dropTable(table.name)

            # IPP FITS files are littered with infinities, so remove these
            self.logger.debug("Removing Infinity values from all columns")
            table = stilts.tpipe(table, cmd='keepcols "' + columns + '"')
            table = stilts.tpipe(table, cmd='replaceval -Infinity null *')
            table = stilts.tpipe(table, cmd='replaceval Infinity null *')

            try:
                table.write(self.scratchDb.url + '#' + table.name)
                # XXX this one is not needed
                # self.scratchDb.killLastConnectionID()
                count = count + 1
            except:
                self.logger.exception("Problem writing table '" + table.name + "' to the database")

            self.tablesLoaded.append(table.name)

        self.logger.infoPair("Done. Imported", "%d tables" % count)
        self.indexIppTables()

        return True

    '''
    Exports PSPS tables from the database to FITS format. Optional regex if you want to alter table names prior to export
    '''    
    def exportPspsTablesToFits(self, regex="(.*)"):

        self.logger.infoPair("Replacing NULLs with", "-999")
        self.logger.infoPair("Changing table names with regex", regex)
        _tables = []

        self.logger.debug("Selecting database tables for export")
        for table in self.tablesToExport:

           # check for an empty table
           if self.scratchDb.getRowCount(table) < 1: continue

           # get everything from table
           try:
               _table = stilts.tread(self.scratchDb.url + '#SELECT * FROM ' + table)
               self.scratchDb.killLastConnectionID()
           except:
               self.logger.exception("Could not read from DB table: " + table)
               return False
               
           # replace nulls and empty fields with weird PSPS -999 pseudo-null
           # this should be sufficient to handle all NANs (which go in as NULL)
           _table = stilts.tpipe(_table, cmd='replaceval "" -999 *')
           self.logger.infoPair('regex', regex)
           self.logger.infoPair('table', table)
           match = re.search(regex, table)
           newTableName = match.group(1)

           # change table names
           _table = stilts.tpipe(_table, cmd='tablename ' + newTableName)
           _tables.append(_table)

        self.logger.infoPair("Writing to FITS", self.outputFitsPath)
        try:
            stilts.twrites(_tables, self.outputFitsPath, fmt='fits')
        except:
            self.logger.exception("Could not write to FITS")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            return False

        return True

    '''
    Searches all PSPS tables and reports the columns that are either partially or completely populated with NULLs
    '''
    def reportNullsInAllPspsTables(self, showPartials):

        for table in self.pspsTables:
            self.scratchDb.reportNulls(table.name, showPartials)

    '''
    Searches all PSPS tables and replaces all NULLs with the provided substitute
    '''
    def replaceAllPspsNulls(self, sub):

        self.logger.infoPair("Replacing all NULL values with", sub)
        for table in self.pspsTables:
            self.scratchDb.replaceNulls(table.name, sub)
        self.logger.infoPair("NULL replacement", "done")

    '''
    Main processing method - must be implemented by subclasses
    '''
    def populatePspsTables(self):
        self.logger.warn("populatePspsTables not implemented yet")

    '''
    Calls DVO program to 'query' DVO database and populate results to local MySQL Db table
    '''
    def getIDsFromDVO(self):

        if self.scratchDb.getRowCount(self.scratchDb.dvoImagesTable) < 1:
            self.logger.errorInfo("No DVO images found in", self.scratchDb.dvoImagesTable)
            return False

        # TODO path to DVO prog hardcoded temporarily
        cmd = "../src/dvograbber " + self.config.settingsPath + " " + self.scratchDb.dbName + " " + self.skychunk.dvoLocation
        self.logger.infoPair("Running DVO", cmd)
        p = Popen(cmd, shell=True, stdout=PIPE)
        p.wait()
        if p.returncode != 0:
            self.logger.errorPair("dvograbber command", "failed")
            return False

        rowCount = self.scratchDb.getRowCount("dvoDetection")
        if rowCount < 1:
            self.logger.errorPair("DVO error", "No IDs found")
            return False

        self.logger.infoPair("DVO access complete. Found", "%d detections" % rowCount)
            
        return True

    '''
    Creates and publishes a batch
    TODO all method calls below should throw exceptions on failure
    '''
    def run(self):

        # creates the empty PSPS tables
        self.logger.infoPair("creating","empty psps tables")
        if not self.createEmptyPspsTables():
            self.logger.errorPair("Aborting this batch", "could not create empty PSPS tables")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            raise

        # fill in basic, not cmf/smf data (eg, ImageMeta for detectionbatch)
        if not self.alterPspsTables():
            self.logger.errorPair("Aborting this batch", "could not alter PSPS tables for batch type")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            raise

        # for P2/ST, this reads the detection tables from the CMF/SMF file(s)
        # for OB, this imports object data from DVO
        if not self.importIppTables():
            if (self.batchType == "OB" or self.batchType =="DO" or self.batchType =="FO" or self.batchType =="FG"):
                self.logger.errorPair("skipping this batch", "no dvo tables for this region")
                self.ippToPspsDb.updateProcessed(self.batchID, 1)
                return True
            self.logger.errorPair("Aborting this batch", "could not import IPP tables")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            raise

        self.logger.infoPair("populatePspsTables","ok")
        if not self.populatePspsTables():
            if self.skipBatch:
                self.ippToPspsDb.updateProcessed(self.batchID, -1)
                return True

            self.logger.errorPair("Aborting this batch", "unable to populate PSPS tables")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            raise 

        if not self.exportPspsTablesToFits():
            self.logger.errorPair("Aborting this batch", "unable to export tables to FITS file")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            raise 
   
        if self.writeBatchManifest():
            if self.skychunk.datastorePublishing: 

                # tar and zip ready for publication to datastore
                if self.tarAndZip():

                    # only now can we report that the batch has successfully processed
                    self.ippToPspsDb.updateProcessed(self.batchID, 1)
                    
                    # now publish to the datastore
                    tarballFile = Batch.getTarballFile(self.batchID)
                    Batch.publishToDatastore(
                            self.datastore, 
                            self.batchID, 
                            self.batchName, 
                            self.subDir, 
                            tarballFile)

        if self.config.onebatch:
            self.logger.infoPair("onebatch requested", "exiting")
            sys.exit(0)

# XXX why is this down here??
from datastore import Datastore
