# coding=UTF8

"""
This classes are used to read a cadastrial CAT file and get several CSV
files with the lines contained on it for PostgreSQL import.
"""

import codecs, os, gzip, time, zipfile, hashlib, fiona, shutil, psycopg2
import psycopg2.extensions, datetime, numpy as np
from subprocess import call

import logging

# Log
logging.basicConfig(filename="catastro_reader_log-%s.log" % \
                    (str(datetime.datetime.now()).replace(' ','-').replace(':','_').split('.')[0]),
                    level=logging.INFO,
                    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                    datefmt="%Y/%m/%d %H:%M:%S")

logging.info("Logging enabled")


class CatastroCatProcess(object):
    """
    This class process a single compressed CAT.
    """
    
    def __init__(self):
        """Constructor."""

    def read(self, file, formatModule, encoding="iso-8859-1", delimiter="|", null="###", indexBase=1):
        """Reads a file."""
        logging.info("Starting process of CAT file %s, format module %s, encoding %s, delimiter %s, null %s, index base %s",
                     file, formatModule.__name__, encoding, delimiter, null, indexBase)
        reader = CatReader(formatModule, encoding=encoding, delimiter=delimiter, null=null, indexBase=indexBase)
        reader.read(file)
        files = dict()
        m = hashlib.md5()
        m.update(file+str(time.time()))
        folder = m.hexdigest()
        os.mkdir(folder)

        # Write files
        for i in ["11", "13", "14", "15", "16", "17"]:
            files[i] = codecs.open(folder+"/"+i+".csv", encoding="utf-8", mode="w")
        
        for i in ["11", "13", "14", "15", "16", "17"]:
            for l in reader.lines[i]:
                files[i].write(l+"\n")

        for k,f in files.iteritems():
            f.close()

    ### CONTINUE HERE writing the CAT single process
        
class CatastroCatBatchProcess(object):
    """
    This class process a batch of compressed CAT.
    """

    root = None
    """Folder root to search."""

    extensions = None
    """Tuple of extensions to search."""

    formatModule = None
    """Reference to the module that holds line formats."""

    encoding = None
    """CAT files encoding, ISO-8859-1 by default."""

    separator = None
    """CSV files separator, | by default."""

    null = None
    """CSV files null value, ### by default."""

    baseIndex = None
    """Base index for decoding CAT fixed length lines (0 or 1), default 1."""

    files = None
    """Dictionary of files to store the final CSV."""

    csvFolder = None
    """Folder to store final CSV."""
    
    def __init__(self, root, formatModule, extensions=(".CAT", ".CAT.gz"), encoding="iso-8859-1",
                 separator="|", null="###", baseIndex=1):
        """Constructor."""
        self.root = root
        self.extensions = extensions
        self.formatModule = formatModule
        self.encoding = encoding
        self.separator = separator
        self.null = null
        self.baseIndex = baseIndex
        self.files = dict()
        self.csvFolder = "CAT_CSV-"+root
        
        os.mkdir(self.csvFolder)
        logging.info("CAT batch process started at folder %s, extensions %s", root, extensions)
        logging.info("Created CAT CSV folder %s" % (self.csvFolder))
    
        for i in ["11", "13", "14", "15", "16", "17"]:
            self.files[i] = codecs.open(self.csvFolder+"/"+i+".csv", encoding="utf-8", mode="w")

    def process(self):
        """Process all CAT found in the search tree."""
        files = self.__getAllCat()
        reader = CatReader(self.formatModule, self.encoding, self.separator, self.null, self.baseIndex)
        times = []

        totalFiles = []
        for r,l in files.iteritems():
            totalFiles.extend(l)

        totalFiles = len(totalFiles)
        count = 0
        totalTime = time.time()
        success = 0
        failures = 0
                        
        for r,l in files.iteritems():
            for f in l:
                t = time.time()
                meanT = (np.mean(times))*(totalFiles-count) if len(times)>0 else "--"
                logging.info("Processing file %s, %i to go, ETA: %s",
                  r+"/"+f, totalFiles-count,
                  str(datetime.timedelta(seconds=meanT)) if not isinstance(meanT, str) else "--")
                reader.read(r+"/"+f)
                  
                for i in ["11", "13", "14", "15", "16", "17"]:
                    for l in reader.lines[i]:
                        self.files[i].write(l+"\n")

                pTime = time.time()-t
                times.append(pTime)
                count+=1
                header = reader.headerDict()
                footer = reader.footerDict()
                metadata = "Results: %s|%i|%f|%s|%s|%s|%s|%s|%i|%i|%i|%i|%i|%i" % \
                    (reader.file, reader.status, pTime, header["Entity code"], header["Entity name"], \
                    header["Date"], header["File type"], header["File description"], footer["N 11"], \
                    footer["N 13"], footer["N 14"], footer["N 15"], footer["N 16"], footer["N 17"])
                logging.info("%s", metadata)
                logging.info("Finished reading file %s", r+"/"+f)

                if reader.status==0:
                    success+=1
                else:
                    failures+=1
                    
        for k,f in self.files.iteritems():
            f.close()

        totalTime = datetime.timedelta(seconds=(time.time()-totalTime))
        logging.info("\n\nFinished processing CAT files at %s:\nTotal files: %s\nTotal time: %s\n"+ \
                     "Success: %s\nFailures: %s\n", \
                     self.root, totalFiles, totalTime, success, failures)
            
    def __getAllCat(self):
        """Finds all CAT files in the tree folder."""
        catFiles = {}

        for root, dirs, files in os.walk(self.root):
            catFile = []

            for file in files:
                if file.endswith(self.extensions):
                    catFile.append(file)

            if catFile!=[]:
                catFiles[root] = catFile

        return catFiles


class CatastroShapeBatchProcess(object):
    """
    This class process a batch of compressed shapes.
    """

    root = None
    """Folder root to search."""

    extensions = None
    """Tuple of extensions to search."""

    encoding = None
    """CAT files encoding, ISO-8859-1 by default."""

    connection = None
    """PG connection object."""

    filter = None
    """A list of lowercase strings of shapefiles to import."""

    sessionName = None
    """Bulk import session name."""
    
    def __init__(self, root, connection, sessionName, filter=None, extensions=(".ZIP", ".zip"), encoding="iso-8859-1"):
        """Constructor."""
        self.root = root
        self.extensions = extensions
        self.connection = connection
        self.encoding = encoding
        self.filter = filter
        self.sessionName = sessionName
        logging.info("Shape batch process started at folder %s, extensions %s", root, extensions)

    def process(self):
        """Process all shape packages found in the search tree."""
        files = self.__getAllShapes()
        reader = ShapeReader(self.connection, self.sessionName, ogr2ogrexe="ogr2ogr", encoding=self.encoding)
        count = 0
        times = []
        totalFiles = []
        success = 0
        failures = 0
        totalTime = time.time()
        reader.pgOpen()
                
        # Count total files
        for r,l in files.iteritems():
            totalFiles.extend(l)

        totalFiles = len(totalFiles)

        for r,l in files.iteritems():
            for f in l:
                t = time.time()
                meanT = (np.mean(times))*(totalFiles-count) if len(times)>0 else "--"
                logging.info("Processing file %s, %i to go, ETA: %s",
                  r+"/"+f, totalFiles-count,
                  str(datetime.timedelta(seconds=meanT)) if not isinstance(meanT, str) else "--")
                results = reader.read(r+"/"+f, filter=self.filter)
                success+=results[0]
                failures+=results[1]
                pTime = time.time()-t
                times.append(pTime)
                count+=1
                logging.info("Ending process of shape package %s/%s, success/failures: %s/%s", r, f, results[0], results[1])

        totalTime = datetime.timedelta(seconds=(time.time()-totalTime))
        logging.info("\n\nFinished processing shape packages at %s:\nTotal packages: %s\nTotal time: %s\n"+ \
                     "Total shapes import success: %s\nTotal shapes import failures: %s\n", \
                     self.root, totalFiles, totalTime, success, failures)
        reader.pgClose()
                      
    def __getAllShapes(self):
        """Finds all shape files in the tree folder."""
        shapeFiles = {}

        for root, dirs, files in os.walk(self.root):
            shapeFile = []

            for file in files:
                if file.endswith(self.extensions):
                    shapeFile.append(file)

            if shapeFile!=[]:
                shapeFiles[root] = shapeFile

        return shapeFiles
                

class ShapeReader(object):
    """
    This class reads catastro shapefiles packages.
    """

    ogr2ogrexe = None
    """Path to the ogr2ogr executable."""

    encoding = None
    """PGClient encoding."""

    pgConnection = None
    """PG connection details."""

    pgConn = None
    """PG connection."""

    sessionName = None
    """Bulk loader session. TODO: should not belong here."""
    
    def __init__(self, pgConnection, sessionName, ogr2ogrexe="ogr2ogr", encoding="iso-8859-1"):
        self.ogr2ogrexe = ogr2ogrexe
        self.encoding = encoding
        self.pgConnection = pgConnection
        self.sessionName = sessionName
        logging.info("ShapeReader initialized")

    def pgOpen(self):
        """Open connection to PG to create metadata."""
        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
        psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)

        try:
            self.pgConn = psycopg2.connect(database=self.pgConnection["db"], user=self.pgConnection["user"],
                                            password=self.pgConnection["pass"], host=self.pgConnection["host"],
                                            port=self.pgConnection["port"])
            logging.info("Connected to PG database %s on %s", self.pgConnection["db"], self.pgConnection["host"])
        except:
            logging.critical("Unable to connect to PG database %s on %s", self.pgConnection["db"], self.pgConnection["host"])
            exit(1)

    def pgClose(self):
        """Closes connection to PG."""
        self.pgConn.close()
        
    def read(self, file, filter=None, resume=True):
        """Imports the ZIP."""

        # TODO: This object should not write logs to PostgreSQL and should
        # interrupt execution by launching exceptions
        # TODO: decompose functionality to other modules, like unzipping
        # TODO: handle only the shapefile
        
        # Check if file exists
        if not os.path.isfile(file):
            logging.error("File %s does not exists", file)
            return(0,1)

        # Check if file has been already processed
        if resume:
            cur = self.pgConn.cursor()
            cur.execute("select gid from "+self.pgConnection["schema"]+".metashape where packagepath=%s;", (file,))
                
            if cur.fetchone() is not None:
                logging.warning("Package %s has been already processed", file)
                cur.close()
                return(0,1)

            cur.close()
                
        # Guess catastro branch
        if "urbano" in file.lower():
            catastroType="URBANO"
        elif "rustico" in file.lower():
            catastroType="RUSTICO"
        else:
            logging.error("Could not guess catastro branch for file %s", file)
            return(0,1)

        success = 0
        failures = 0
                
        # Unzip the shapes
        try:
            m = hashlib.md5()
            m.update(file+str(time.time()))
            folder = m.hexdigest()
            os.mkdir(folder)
            z = zipfile.ZipFile(file)
            z.extractall(folder)
            z.close()

            logging.info("Coding file %s with hash %s", file, folder)            
            logging.info("Created temporary folder %s", folder)
             
            for dirpath, dirs, files in os.walk(folder):
                for f in files:
                    z = zipfile.ZipFile(folder+"/"+f)
                    z.extractall(folder)
                    z.close()

            logging.info("Unzipped shape package %s", file)
        except:
            logging.info("Error unzipping shape package %s", file)

        # Find shapes and DBF
        for dirpath, dirs, files in os.walk(folder):
            shapes = []

            for f in files:
                if f.endswith((".dbf", ".DBF")):
                    shapes.append(f)

        # Process each shape
        tables = []
        
        for shape in shapes:
            if (filter is not None and shape.split(".")[0].lower() in filter) or (filter is None):
                src = fiona.open(folder+"/"+shape)
                logging.info("Processing: %s, features: %i", shape, len(src))

                # Determine CRS
                if str(src.crs)=="{u'units': u'm', u'no_defs': True, u'ellps': u'intl', u'proj': u'utm', u'zone': 30}":
                    epsg='-a_srs "epsg:23030"'
                elif str(src.crs)=="{u'units': u'm', u'no_defs': True, u'ellps': u'GRS80', u'proj': u'utm', u'zone': 30}":
                    epsg='-a_srs "epsg:25830"'
                elif str(src.crs)=="{u'units': u'm', u'no_defs': True, u'ellps': u'GRS80', u'proj': u'utm', u'zone': 29}":
                    epsg='-a_srs "epsg:25829"'                    
                else:
                    epsg = "" 
                
                # # Determine CRS
                # if src.crs<>{}:
                #     if src.crs[u"ellps"]==u"GRS80":
                #         if src.crs[u"proj"]==u"utm":
                #             if src.crs[u"zone"]==30:
                #                 epsg='-a_srs "epsg:25830"'
                #             elif src.crs[u"zone"]==29:
                #                 epsg='-a_srs "epsg:25829"'
                #             else:
                #                 # raise ShapeImporterExceptionUnknownSRS(shape=folder+"/"+shape, crs=str(src.crs))
                #                 epsg = ""
                #         else:
                #             # raise ShapeImporterExceptionUnknownSRS(shape=folder+"/"+shape, crs=str(src.crs))
                #             epsg = ""
                #     else:
                #         # raise ShapeImporterExceptionUnknownSRS(shape=folder+"/"+shape, crs=str(src.crs))
                #         epsg = ""
                # else:
                #     epsg = ""

                if epsg<>"":
                    logging.info("Identified CRS: %s", epsg)
                    tableName = shape.lower().rstrip(".dbf")+"___"+folder
                    command = ('SHAPE_ENCODING=UTF-8 PGCLIENTENCODING=%s %s -lco OVERWRITE=YES -progress -gt 65000 '+ \
                            '-f "PostgreSQL" PG:"host=%s user=%s dbname=%s password=%s port=%s" %s '+ \
                            '-lco SCHEMA=%s -nln %s -lco LAUNDER=YES -lco GEOMETRY_NAME=geom -nlt %s %s') % \
                            (self.encoding, self.ogr2ogrexe, self.pgConnection["host"], \
                            self.pgConnection["user"], self.pgConnection["db"], \
                            self.pgConnection["pass"], self.pgConnection["port"], epsg, self.pgConnection["schema"], \
                            tableName, src.schema["geometry"], folder+"/"+shape)
                    tables.append(shape.lower().rstrip(".dbf"))
                    out = call(command, shell=True)

                    if out==0:
                        logging.info("Successfully imported file %s, features %s, hash %s, table %s", \
                                    shape, len(src), folder, tableName)
                        success+=1
                    else:
                        logging.error("Import failed for file %s, hash %s, table %s___%s", \
                                    file, folder, shape.lower().rstrip(".dbf"), folder)
                        failures+=1
                
                    src.close()

                    ts = time.time()
                    st = datetime.datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")
                    
                    # Write meta to database
                    cur = self.pgConn.cursor()
                    cur.execute("insert into "+self.pgConnection["schema"]+".metashape(importsession, time, "+\
                                "pgconnection, catastrotype, packagepath, ogr2ogrexe, srcencoding, "+\
                                "hash, unpackedpath, crs, ogrcommand, schemaname, tablename, geomtype, "+\
                                "features, status) "+\
                                "values(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)", \
                                (self.sessionName, st, str(self.pgConnection), catastroType, \
                                file, self.ogr2ogrexe, self.encoding, folder, folder+"/"+shape, epsg, \
                                command, self.pgConnection["schema"], tableName, src.schema["geometry"], \
                                len(src), out))
                    self.pgConn.commit()
                    cur.close()
                else:
                    logging.error("Shapefile %s has an unknown CRS: %s", shape, str(src.crs))
            else:
                logging.info("Shapefile %s filtered", shape)
            
        shutil.rmtree(folder)
        logging.info("Cleaned temporary folder %s", folder)
        return (success, failures)
    
    
class CatReader(object):
    """
    This class reads a CAT file and returns several CSV for PostgreSQL with
    the line data.
    """

    lines = None
    """Dictionary of lists to hold different kinds of lines."""
    
    formatModule = None
    """Reference to the module that holds line formats."""

    encoding = None
    """CAT file encoding."""

    delimiter = None
    """Delimiter unicode sequence for CSV."""

    null = None
    """None unicode sequence for CSV."""

    indexBase = None
    """Represents if the first element of the fixed length string is 0 or 1."""

    header = None
    """01 header info."""

    footer = None
    """90 footer info."""

    file = None
    """Processed file."""

    status = None
    """Last read status code."""
    
    def __init__(self, formatModule, encoding='iso-8859-1', delimiter='|', null='###', indexBase=1):
        """Constructor."""
        self.formatModule = formatModule
        self.encoding = encoding
        self.delimiter = delimiter
        self.null = null
        self.indexBase = indexBase
        self.lines = dict()

    def read(self, file):
        """Reads a CAT file."""
        # Check for gzipped files
        try:
            f = gzip.open(file, mode="r")
            contents = f.read()
            f.close()
            u_str = contents.decode(self.encoding).split("\n")[:-1]
            logging.info("Readed gzipped file %s", file)
        except IOError:
            f = codecs.open(file, encoding=self.encoding, mode="r")
            u_str = f.read().split("\n")[:-1]
            f.close()
            logging.info("Readed plain file %s", file)

        lineTypes = ["01", "11", "13", "14", "15", "16", "17", "90"]
        fixedLineReader = dict()
        
        for i in lineTypes:
            fixedLineReader[i] = FixedLineReader(getattr(self.formatModule, "format_"+i), self.indexBase)
        
        for i in lineTypes[1:-1]:
            self.lines[i] = []

        for line in u_str:
            # Check line type
            idLine = line[0:2]

            if idLine=="01":
                self.header = fixedLineReader["01"].readLine(line)
                logging.info("Found header for file %s", file)
            elif idLine=="90":
                self.footer = fixedLineReader["90"].readLine(line)
                logging.info("Found footer for file %s", file)
            else: 
                self.lines[idLine].append(self.__listToCsv(fixedLineReader[idLine].readLine(line)))
        
        # Check for integrity
        if self.header is None:
            self.status = -1
            raise CatImporterExceptionNoHeader(fichero=file)

        if self.footer is None:
            self.status = -2
            raise CatImporterExceptionNoFooter(fichero=file)

        for i in range(1, 7):
            self.status = -3
            if len(self.lines[lineTypes[i]])<>int(self.footer[i]):
                raise CatImporterException(fichero=file, registro=lineTypes[i],
                                           realCount=len(self.lines[lineTypes[i]]),
                                           expectedCount=int(self.footer[i]))

        self.status = 0
        logging.info("Finished processing file %s, status %i", file, self.status)

    def headerDict(self):
        """Returns a dictionary with interpreted header data."""
        d = dict()
        d["Entity code"] = self.header[2]
        d["Entity name"] = self.header[3]
        d["Date"] = self.header[4]
        d["File type"] = self.header[6]
        d["File description"] = self.header[7]

        return d
    
    def footerDict(self):
        """Returns a dictionary with interpreted footer data."""
        d = dict()
        d["N 11"] = int(self.footer[1])
        d["N 13"] = int(self.footer[2])
        d["N 14"] = int(self.footer[3])
        d["N 15"] = int(self.footer[4])
        d["N 16"] = int(self.footer[5])
        d["N 17"] = int(self.footer[6])

        return d
                
    def __listToCsv(self, valList):
        """Converts to CSV a list of strings."""
        l = u""

        for i in valList:
            if i==u'':
                v = self.null
            else:
                v = i
                
            l += v+self.delimiter

        return l.rstrip(self.delimiter)


class FixedLineReader(object):
    """This class uses a format dictionary definition to decode a line of plain
    text data of fixed length.
    """

    format = None
    """List format to parse upon."""

    indexBase = None
    """Index base (depends on if first character in format is 0 or 1)."""
        
    def __init__(self, format, indexBase):
        """Constructor."""
        self.format = format
        
        if indexBase==1:
            self.indexBase = -1
        elif indexBase==0:
            self.indexBase = 0

    def readLine(self, line):
        """Reads an UNICODE line and parse based on list format."""
        lineData = []

        # Loop line information bits
        for i in range(0, len(self.format), 2):
            data = line[self.__getSlice(self.format[i], self.format[i+1])]
            lineData.append(unicode(data.strip()))

        return lineData

    def __getSlice(self, start, length):
        """Returns a string slice."""
        return slice(start+self.indexBase,start+length+self.indexBase)


# Exceptions

class CatImporterExceptionNoHeader(Exception):
    def __init__(self, **args):
        self.fichero = unicode(args["fichero"])

    def __str__(self):
        return "Excepcion importando fichero %s: sin cabecera 01." % (self.fichero)
    
class CatImporterExceptionNoFooter(Exception):
    def __init__(self, **args):
        self.fichero = unicode(args["fichero"])

    def __str__(self):
        return "Excepcion importando fichero %s: sin resumen 90." % (self.fichero)
    
class CatImporterExceptionMismatch(Exception):
    def __init__(self, **args):
        self.fichero = unicode(args["fichero"])
        self.registro = unicode(args["registro"])
        self.realCount = unicode(args["realCount"])
        self.expectedCount = unicode(args["expectedCount"])

    def __str__(self):
        return "Excepcion importando fichero %s en registros de tipo %s: encontrados %s, esperados %s" % \
          (self.fichero, self.registro, self.realCount, self.expectedCount)

class ShapeImporterExceptionUnknownSRS(Exception):
    def __init__(self, **args):
        self.shape = unicode(args["shape"])
        self.crs = unicode(args["crs"])

    def __str__(self):
        return "Excepcion importando shapefile %s: CRS desconocido %s" % \
          (self.shape, self.crs)

class ShapeImporterExceptionNoConnection(Exception):
    def __init__(self, **args):
        self.host = unicode(args["host"])
        self.user = unicode(args["user"])
        self.port = unicode(args["port"])    

    def __str__(self):
        return "Excepción conectando a %s, usuario %s, puerto %s" % \
          (self.host, self.user, self.port)

class ShapeImporterExceptionBadCatastroType(Exception):
    def __init__(self, **args):
        self.catastroType = unicode(args["catastroType"])

    def __str__(self):
        return "Critico: %s no es un catastro valido (RUSTICO o URBANO)" % (self.catastroType)
          
