#!/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

    """
    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 NetCDF python bindings to return the version
    """
    def netcdf_return_format(self, nc_file):
        try:
            nc = Dataset(nc_file)
        except:
            return 'Not a valid netcdf'
        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)
        res = 'ok'

        # No valid formats found
        if len(ffchecks) == 0:
            res = 'err'
            msg = 'Directory is empty or no accepted file formats were found'

        # At least one valid format found
        else:
            msg = 'File formats found: {}'.format(ffchecks)

            # For every format found
            for f in ffchecks:
                if 'grib' in f.lower():
                    if '2' in f:
                        msg += '. {} is considered valid'.format(f)
                    else:
                        res = 'warn'
                        msg += '. {} is considered obsolete (GRIB2 is recommended)'.format(f)
                if 'netcdf' in f.lower():
                    if f == 'NETCDF4_CLASSIC':
                        msg += '. {} is considered valid'.format(f)
                    elif 'valid' in f.lower():
                        res = 'error'
                        msg += '. Most likely a corrupt netcdf file'
                    else:
                        res = 'warn'
                        msg += '. {} 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()
