#!/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 scratchdb import ScratchDb
from gpc1db import Gpc1Db
from ipptopspsdb import IppToPspsDb
from fits import Fits

from java.lang import *
from java.sql import *
from xml.etree.ElementTree import ElementTree, Element, tostring

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

    '''
    Constructor

    >>> batch = Batch(1,2,3,4,5,6,7)
    >>> print batch.pspsVoTableFilePath
    "../config/2/tables.vot"
    '''
    def __init__(self, 
                 logger, 
                 configPath,
                 doc,
                 gpc1Db,
                 ippToPspsDb,
                 id,
                 batchID,
                 batchType, 
                 fits): 

        self.everythingOK = False
        self.readHeader = False
        self.configPath = configPath
        self.doc = doc
        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.batchType = batchType;
        self.pspsVoTableFilePath = "../config/" + batchType + "/tables.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 for id", "%d" % id)
                return

        # if no fits file, and this is not an IN batch, then get out of here
        else:
           if batchType != "IN": 
               logger.errorPair("Could not read FITS for id", "%d" % id)
               return

        # get info from config
        self.survey = self.doc.find("options/survey").text
        self.pspsSurvey = self.doc.find("options/pspsSurvey").text
        self.dvoGpc1Label = self.doc.find("dvo/gpc1Label").text
        self.dvoLocation = self.doc.find("dvo/location").text
        self.useFullTables = int(self.doc.find("dvo/useFullTables").text)
        self.scratchDb = ScratchDb(logger, self.doc, self.useFullTables)

        if not self.scratchDb.everythingOK: return

        # TODO
        self.tablesToExport = []

        if self.survey != "":
            self.surveyID = self.scratchDb.getSurveyID(self.survey)
        else:
            self.surveyID = -1;
       
        # get some options from the config
        self.testMode = int(self.doc.find("options/testMode").text)
        self.dataRelease = int(self.doc.find("metadata/dataRelease").text)

        # create datastore object
        self.datastore = Datastore(self.logger, self.doc, self.ippToPspsDb)

        # get local storage location from config
        self.batchName = Batch.getNameFromID(self.batchID)
        self.basePath = self.doc.find("localOutPath").text
        self.subDir = Batch.getSubDir(
                self.basePath,
                self.batchType, 
                self.dvoGpc1Label)

        self.localOutPath = Batch.getOutputPath(
                self.basePath,
                self.batchType, 
                self.dvoGpc1Label,
                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 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", self.survey)
        self.logger.infoPair("Survey ID", "%d" % self.surveyID)
        self.logger.infoPair("Publishing to PSPS as survey", self.pspsSurvey)
        self.logger.infoPair("DVO location", self.dvoLocation)
        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.logger.infoPair("Output path", self.localOutPath)

        self.everythingOK = True
    
    '''
    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.testMode: 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") 
        if self.batchType != "IN": root.attrib['survey'] = self.pspsSurvey
        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
        p = Popen("xmllint --format " + tmpPath + " > " + outPath, shell=True, stdout=PIPE)
        p.wait()
        os.remove(tmpPath)


    '''
    Publishes this batch to the datastore
    '''
    def tarAndZip(self):
      
        # set up filenams and paths
        tarFile = Batch.getTarFile(self.batchID)
        tarPath = Batch.getTarPath(self.basePath, self.batchType, self.dvoGpc1Label, self.batchID)
        tarballFile = Batch.getTarballFile(self.batchID)
        tarballPath = Batch.getTarballPath(self.basePath, self.batchType, self.dvoGpc1Label, 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

        # only now can we report that the batch has successfully processed
        self.ippToPspsDb.updateProcessed(self.batchID, 1)

        # 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)
      tables = stilts.treads(self.fits.getPath())

      count = 0
      for table in tables:

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

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

    '''
    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("dvoMeta") < 1:
            self.logger.error("No DVO IDs found in dvoMeta")
            return False

        # TODO path to DVO prog hardcoded temporarily
        cmd = "../src/dvograbber " + self.configPath + " " + self.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 methods call below should throw exceptions on failure
    '''
    def run(self):

        if not self.everythingOK:
            self.logger.error("Aborting this batch due to (hopefully) previously reported errors")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            return

        if not self.createEmptyPspsTables():
            self.logger.error("Aborting this batch due to (hopefully) previously reported errors")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
            return

        self.importIppTables()
        if not self.populatePspsTables():
            self.logger.errorPair("Aborting this batch", "unable to populate PSPS tables")
            self.ippToPspsDb.updateProcessed(self.batchID, -1)
        else:
            if not self.exportPspsTablesToFits():
                self.logger.errorPair("Aborting this batch", "unable to export tables to FITS file")
            else:
                self.writeBatchManifest()
                if int(self.doc.find("options/publishToDatastore").text): 
                    # tar and zip ready for publication to datastore
                    if self.tarAndZip():
                        # now publish to the datastore
                        tarballFile = Batch.getTarballFile(self.batchID)
                        Batch.publishToDatastore(self.datastore, self.batchID, self.batchName, self.subDir, tarballFile)

                if int(self.doc.find("options/reportNulls").text): self.reportNullsInAllPspsTables(False)

from datastore import Datastore

