#!/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 subprocess
from netCDF4 import Dataset

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

"""
Class: Standard compliance checks
---
- Checks for file format correctness
"""
class StandardComplianceTest:

    def __init__(self, files, forms_dir, cds_dataset, cds_variable, suffix='grib'):
        self.testid = 'test01-standard_compliance_' + suffix
        self.logger = logging.getLogger(self.testid)   
        self.files = files
        self.forms_dir = forms_dir
        self.cds_variable = cds_variable
        self.cds_dataset = cds_dataset

    """
    Run CFChecker for one file and CF-XX
    """
    def run_cfchecker(self, item):
        """RUN CFCHECKER FOR ONE FILE
        """
        cmd0 = "cfchecks --version auto %s" %item
        p0 = subprocess.Popen(
            cmd0, shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT
            )
        out0, _ = p0.communicate()
        return out0

    """
    NetCDF - Check for CF compliance (netCDF)
    """
    def netcdf_test(self, ncfile):

        # Defaults (for seasonal, nothing is compliant with C3S standard yet, even if CF-compliant)
        if 'seasonal' in ncfile:
            res = 'warn'
            msg = 'Files are not compliant with C3S-0.1 netcdf convention.'
        else:
            res = 'ok'
            msg = ''

        # Run external CF-checker
        out0 = self.run_cfchecker(ncfile)
        string = str(out0).replace('\'',"").replace("\"","").split('\\n')

        # All output is written to log file
        version = "N/A"
        for e in string:
            self.logger.info(str(e))
            if 'Checking against CF Version' in e:
                version = e.replace('Checking against CF Version', '').strip()
                version = version.replace('CF-', '')

        # Detect error messages
        if any('ERROR:' in e for e in string):
            err = [item.split('ERROR:')[-1].strip() for item in string if item.startswith('ERROR:')]
            errstring = ncfile.split('/')[-1] + ' has ' + str(len(err))
            errstring += ' error messages' + '\n'
            errstring += '\n'.join(err)
            string += errstring
            self.logger.error(errstring)
            res = 'err'

        # Detect OS error messages
        if any('OSError' in e for e in string):
            err = [item.split('OSError:')[-1].strip() for item in string if item.startswith('OSError:')]
            self.logger.error(err)
            res = 'err'

        # In any case they will not be compliant with C3S-0.1 but they can be CF-compliant
        if res != 'err':
            cf_msg = 'Files are compliant with the Climate and Forecast Metadata Convention v{vv} (CF-v{vv})'.format(vv=version)
        else:
            cf_msg = 'Files are not compliant with the Climate and Forecast Metadata Convention v{vv} (CF-v{vv})'.format(vv=version)

        return res, msg + ' ' + cf_msg

    """
    Run Tigge Check for UERRA/GRIB2 files
    """
    def run_tigge_check(self, item):
        """RUN GRIB2 TIGGE_CHECK FOR ONE FILE
        """
        cmd0 = "tigge_check -u -w -z %s" %item
        p0 = subprocess.Popen(
            cmd0, shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT
            )
        out0, _ = p0.communicate()
        return out0

    """
    GRIB1 - is not supported, directly report error
    GRIB2 - checked against WMO tables (version 2)
    """
    def grib_test(self, grib_file, grib_info):
                
        # Read GRIB file
        gfc = GribFileCheck(grib_file)

        # Version check (GRIB1)
        if gfc.version == 1:
            return 'err', 'Files are not compliant with GRIB2 standard (under development). GRIB1 is obsolete according to the WMO'

        # We only check GRIB2 keys
        else:
            if 'uerra' in grib_file:
                out0 = self.run_tigge_check(grib_file).decode('UTF-8')
                string = str(out0).split('\\n')
                err = False
                for e in string:
                    if len(e):
                        if 'error' in e.lower():
                            err = True
                            self.logger.error(str(e))
                        elif 'warning' in e.lower():
                            err = True
                            self.logger.warning(str(e))
                        else:
                            self.logger.info(str(e))

                # Summary
                if err:
                    return 'err', 'Files are not compliant with GRIB2 standard for UERRA'
                else:
                    return 'ok', 'Files are compliant with GRIB2 standard for UERRA'
            else:
                # Get parameter associated to variable (ex: 2m_temperature)
                paramId = get_param_id(self.cds_dataset, self.cds_variable, self.forms_dir)

                try:
                    check_keys = grib_info[str(paramId)]['grib_codes']
                except:
                    raise Exception('GRIB ParamId {} was not found'.format(paramId))

                # Check that all keys are present (no need to proceed if error is found)
                err_keys = gfc.check_mandatory_keys(check_keys)
                if err_keys != []:
                    for e in err_keys:
                        self.logger.error(e)
                    return 'err', 'Files are not compliant with GRIB2 standard (under development)'

        return 'warn', 'Files are compliant with GRIB2 standard (under development)'

    def main(self):

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

        # Do check
        msg = ''
        res = ''

        # Load grib defintions (already generated)
        grib2_info = load_grib_metadata(version=2)

        # Error message files not found (useful for the netcdf converter, sometimes CDSAPI is not able to return them)
        if len(self.files) == 0:
            res = 'err'
            msg = 'No files found for the analysis'
        else:
            for f in self.files:
                # Progress update
                i+=1
                print_progress(i, n)

                # Grib checker or CF checker
                if f.endswith('.nc'):
                    res, msg = self.netcdf_test(f)
                elif f.endswith('.grib'):
                    res, msg = self.grib_test(f, grib2_info)
                else:
                    res = 'err'
                    msg = 'File format {} is not supported'.format(f)

                # If result is err no need to proceed
                if res == 'err':
                    break

        # Final message
        if res == 'ok':                       
            self.logger.info(msg)
        else:            
            self.logger.error(msg)
            
        # Prepare output
        return { 
            'res': res, 
            'msg': msg
        }

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