#!/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,
                 useFullTables): 

        # print "starting the batch"

        # self.printline = 0
        # self.testprint()

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

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

        # self.testprint()

        # set up class variables
        self.id = id
        self.batchID = batchID
        self.gpc1Db = gpc1Db
        self.ippToPspsDb = ippToPspsDb
        self.scratchDb = scratchDb
        self.batchType = batchType;
        self.pspsVoTableFilePath = self.config.configDir + "tables." + batchType + ".vot"

        # self.testprint()

        # 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

        self.scratchDb.setUseFullTables(self.useFullTables)

        # 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)

        # self.testprint()

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

        # create DVO tables if accessing DVO directly
        if not self.useFullTables: self.scratchDb.createDvoTables()

        # 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)
        self.logger.infoBool("Use full DVO tables?", self.useFullTables)

        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.testprint()

        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"

    '''
    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()
        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)

    '''
    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)
             table.write(self.scratchDb.url + '#' + table.name)
             self.tablesToExport.append(table.name)

         return self.alterPspsTables();

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


    '''
    Alter PSPS tables
    '''   
    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=""):

      self.logger.infoPair("Importing tables 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

      count = 0
      for table in tables:

          # print "import smf table ", table
          match = re.match(filter, table.name)
          if not match: continue
          self.logger.debugPair("Reading IPP table", table.name)
          table = stilts.tpipe(table, cmd='explodeall')

          # print "read smf table ", table

          # 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 *')
          # print "cleaned up values ", table

          try:
              table.write(self.scratchDb.url + '#' + table.name)
              # print "wrote to mysql ", table
              self.scratchDb.killLastConnectionID()
              count = count + 1
          except:
              self.logger.exception("Problem writing table '" + table.name + "' to the database")

      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
           _table = stilts.tpipe(_table, cmd='replaceval "" -999 *')

           match = re.match(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()

        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):

        # this is badly named : it creates the tables and then fills in
        # the ImageMeta tables for each ota (alterPspsTables)
        if not self.createEmptyPspsTables():
            self.logger.errorPair("Aborting this batch", "could not create empty PSPS tables")
            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():
            self.logger.errorPair("Aborting this batch", "could not import IPP tables")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            raise

        if not self.populatePspsTables():
            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)

from datastore import Datastore

