#!/usr/bin/env python
#
# Data Quality Checker also known as DQC
#
# Read the README.md for more detailed information
#
#       Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
#       Earth Sciences Department - Computational Earth Sciences
#     
#     Authors: 
#       Joan Sala Calero (joan.sala@bsc.es)

import sys
import magic
from netCDF4 import Dataset

# Local import
from dqc_wrapper.utils.utils_files import *
from dqc_wrapper.utils.utils_grib import *

"""
Class: File format checks
---
- Checks for file format correctness
- Uses bash 'file' operation
- Uses specific libs like eccodes and netcdf4 for gridded files
"""
class FileFormatCheck:

    def __init__(self, files):
        self.testid = 'test00-file_format'
        self.logger = logging.getLogger(self.testid)      
        self.files = files

    """
    Deprecated:
    This function uses the CDO module, with command sinfon to retrieve file format
    To be used with gridded files
    """
    def cdo_return_format(self, filename):

        # This is faster than using the cdo bindings (avoids reading whole file)
        cmd0 = "cdo sinfon %s | head | grep \"File format\"" %filename
        p0 = subprocess.Popen(
            cmd0, shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT
            )
        out0, _ = p0.communicate()
        out0 = str(out0).lower()
        
        # Allowed file formats
        if 'netcdf' in out0:
            if 'netcdf4' in out0:       return 'NetCDF4'
            elif 'netcdf3' in out0:     return 'NetCDF3'
            elif 'netcdf2' in out0:     return 'NetCDF2'
            else:                       return 'NetCDF'

        return 'unkown format'

    """
    This function uses the ECCODES api, looking into the first GRIB message containing the version
    """
    def eccodes_return_format(self, grib_file):
        gf = GribFileCheck(grib_file)
        if gf.version == 2:
            return 'GRIB2'
        elif gf.version == 1:
            return 'GRIB1'
        else:
            return 'unkown format'

    """
    This function uses the ECCODES api, looking into the first GRIB message containing the version
    """
    def netcdf_return_format(self, nc_file):
        nc = Dataset(nc_file)
        return nc.data_model

    """
    This function returns the type of a file (using the bytes: equivalent to file UNIX operation)
    To be used with other files
    """
    def file_return_format(self, filename):    
        return magic.from_file(filename)

    """
    Append formats to a list
    """
    def append_formats(self, formats, f):
        if f.endswith(('.grib', '.grb')):
            formats.append(self.eccodes_return_format(f))
        elif f.endswith(('.nc', '.nc3', '.nc4')):
            formats.append(self.netcdf_return_format(f))
        elif f.endswith(('.zip', '.gzip', '.gz')):
            formats += self.zip_return_formats(f)
        else:
            formats.append(self.file_return_format(f))
        return formats

    """
    This function returns the type of the files contained in a zip file
    For that purpose creates a temporary dir
    """
    def zip_return_formats(self, filename):

        formats = []
        with tempdir() as tmpdir:
            with zipfile.ZipFile(filename, 'r') as zipf:
                zipf.extractall(tmpdir)
                for f in os.listdir(tmpdir):
                    formats = self.append_formats(formats, os.path.join(tmpdir, f))

        return formats

    def main(self):

        # Init
        res = 'ok'
        msg = ''

        # Progress init
        i = 0
        n = len(self.files)
        print()

        # Loop over files
        ffchecks = []
        for f in self.files:

            # Progress update            
            i+=1
            print_progress(i, n)

            # Formats check
            ffchecks = self.append_formats(ffchecks, f)

        # Only unique values
        ffchecks = list(set(ffchecks))
        self.logger.info('Formats: {}'.format(ffchecks))

        # File formats found (it should be only one)
        if len(ffchecks) > 1:
            res = 'err'
            msg = 'Multiple file formats found {}'.format(ffchecks)
        else:            
            if len(ffchecks) == 0:
                res = 'err'
                msg = 'Directory is empty or no accepted file formats were found'
            else:
                # One homogenous file format
                f = ffchecks[0]
                if 'grib' in f.lower():
                    if '2' in f:
                        msg = 'File format found {} is considered valid'.format(f)
                    else:
                        res = 'warn'
                        msg = 'File format found {} is considered obsolete (GRIB2 is recommended)'.format(f)
                if 'netcdf' in f.lower():
                    if f == 'NETCDF4_CLASSIC':
                        msg = 'File format found {} is considered valid'.format(f)
                    else:
                        res = 'warn'
                        msg = 'File format found {} is considered obsolete (NETCDF4_CLASSIC is recommended)'.format(f)

        # Logging
        if res == 'ok':
            self.logger.info(msg)
        else:
            self.logger.error(msg)

        # Prepare output        
        return { 
            'res': res, 
            'msg': msg
        }


if __name__ == "__main__":        
    sys.exit()
