#!/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 datetime
import sqlite3
import re, os

import datetime as dt
from sqlite3 import Error
from cdo import Cdo

"""
Create SQLite DB connection to file
"""
def create_connection(db_file):

    conn = None
    try:
        conn = sqlite3.connect(db_file)

    except Error as e:
        print(e)

    return conn

"""
-- Main database table
CREATE TABLE dataset (
	"_datasetId" TEXT PRIMARY KEY,
	"datasetGroup" TEXT,
	"lastUpdated" TEXT
);
"""
def insert_dataset(conn, _datasetId, datasetGroup):

    lastUpdated = datetime.datetime.now().isoformat()
    sql = '''INSERT INTO dataset(_datasetId, datasetGroup, lastUpdated)
             VALUES("{}","{}","{}") '''.format(_datasetId, datasetGroup, lastUpdated)
    cur = conn.cursor()
    cur.execute(sql)
    conn.commit()

    return cur.lastrowid

def select_datasets(conn):

    sql = '''SELECT _datasetId,datasetGroup FROM dataset'''
    cur = conn.cursor()
    cur.execute(sql)
    rows = cur.fetchall()

    return rows

"""
-- Stores all files for a dataset X (normally a file per month)
CREATE TABLE file (
	"_fname" TEXT,
	"downRequired" INTEGER,
	"downPerformed" INTEGER,
	-- Foreign keys
	"_ds" TEXT,
	FOREIGN KEY (_ds) REFERENCES dataset(_datasetId),
	PRIMARY KEY (_fname, _ds)
);
"""
def insert_files(conn, datasetId, cds_forms_dates, cds_variable, origin, system):

    # Perform download (2 months per year, those months change per variable(m))
    startM = len(cds_variable) % 6 + 1
    endM = (startM + 6)
    count = 0
    files = []

    dates_to_download = []
    for yr in range(cds_forms_dates[0].year, cds_forms_dates[-1].year + 1):
        if yr < 2020:
            mstart = 1 + (startM + count) % 12
            mend = 1 + (endM + count) % 12
            dates_to_download.append('{}{:02}01'.format(yr, mstart))
            dates_to_download.append('{}{:02}01'.format(yr, mend))
            count += 1

    # Go through all cds-forms dates, extract month and mark for download
    for dt in cds_forms_dates:

        # We store all data per month YYYYMMDD
        fnameDate = dt.strftime('%Y%m01')
        year = int(dt.strftime('%Y'))
        if 'seasonal' in datasetId:
            fname = "{o}-{s}-{dt}".format(dt=fnameDate, o=origin, s=system)
        else:
            if 'uerra' in datasetId:
                fname = "{o}-{dt}".format(dt=fnameDate, o=origin)
            else:
                fname = "{v}-{dt}".format(dt=fnameDate, v=cds_variable)

        # 2 months per year strategy for hourly products - before Jan2020
        downRequired = 1
        downPerformed = 0
        if not('monthly' in datasetId) and not(fnameDate in dates_to_download) and year<2020:
            downRequired = 0

        # Insert or replace (avoid double entries)
        sql = '''INSERT OR REPLACE INTO file(_fname, downRequired, downPerformed, _ds)
                 VALUES("{}", {}, {}, "{}")'''.format(fname, downRequired, downPerformed, datasetId)
        files.append(fname)

        cur = conn.cursor()
        cur.execute(sql)
        conn.commit()

    return list(set(files))

def update_downperformed(conn, val, _ds, _fn):

    sql = '''UPDATE file
             SET downPerformed={}
             WHERE _ds="{}" AND _fname="{}";'''.format(val, _ds, _fn)
    cur = conn.cursor()
    cur.execute(sql)
    conn.commit()

    return cur.lastrowid

"""
Get all timesteps given a dataset
"""
def select_files(conn, _ds):

    sql = '''SELECT _fname FROM file WHERE _ds="{}"'''.format(_ds)
    cur = conn.cursor()
    cur.execute(sql)
    rows = cur.fetchall()

    return rows

"""
Get all files needed for download
"""
def select_required_downloads(conn, _ds):

    sql = '''SELECT _fname FROM file WHERE downRequired=1 AND downPerformed=0 AND _ds="{}"'''.format(_ds)
    cur = conn.cursor()
    cur.execute(sql)
    rows = cur.fetchall()

    return rows

"""
-- Each file has T timestamps (this table is filled and used by the temporal_completeness and temporal_consistency checks)
CREATE TABLE filedatetime (
	"_timestamp" TEXT,
	-- Foreign keys
	"_ds" TEXT,
	"_fn" TEXT,
	FOREIGN KEY (_ds) REFERENCES dataset(_datasetId),
	FOREIGN KEY (_fn) REFERENCES file(_fname),
	PRIMARY KEY (_ds, _fn, _timestamp)
);
"""
def insert_filedatetime(conn, _timestamp, _fn, _ds):

    try:
        sql = '''INSERT INTO filedatetime(_timestamp, _fn, _ds)
                 VALUES("{}","{}","{}") '''.format(_timestamp, _fn, _ds)
        cur = conn.cursor()
        cur.execute(sql)
        conn.commit()
    except sqlite3.IntegrityError:
        pass # existing timestamp

    return cur.lastrowid

"""
Get all timesteps given a dataset
"""
def select_filedatetimes(conn, _ds):

    sql = '''SELECT _timestamp FROM filedatetime WHERE _ds="{}"'''.format(_ds)
    cur = conn.cursor()
    cur.execute(sql)
    rows = cur.fetchall()

    return rows

"""
Insert into the table filedatetime all the timestamps collected from the GRIB files (initialization dummy db)
"""
def update_filedatetimes(conn, dataset, files):

    # Scan for available dates
    cdo = Cdo()
    for f in files:
        info = cdo.showtimestamp(input=f)

        # Collect all timestamps including HHMMSS
        dates = []
        for inf in info:
            for t in inf.split():
                match = re.search(r'\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}', t)
                if match:  # YYYY-MM-DDTHH:MM:SS minimum requirement
                    dates.append(dt.datetime.strptime(match.group(), '%Y-%m-%dT%H:%M:%S'))

        # Store unique dates
        for d in list(set(dates)):
            insert_filedatetime(conn, d.isoformat(), os.path.basename(f).replace('.grib', ''), dataset)

"""
-- Each file goes through D datachecks (standard_compliance, spatial_completeness, ... , valid_ranges)
CREATE TABLE filedatacheck (
	"_checkId" TEXT,
	"checkResult" INTEGER DEFAULT 0,
	"checkMsg"	INTEGER,
	-- Foreign keys
	"_ds" TEXT,
	"_fn" TEXT,
	FOREIGN KEY (_ds) REFERENCES dataset(_datasetId),
	FOREIGN KEY (_fn) REFERENCES file(_fname),
	PRIMARY KEY (_ds, _fn, _checkId)
);
"""
def insert_filedatacheck(conn, _ds, _fn, _checkId, checkResult, checkMsg):

    sql = '''INSERT INTO filedatacheck(_ds, _fn, _checkId, checkResult, checkMsg)
             VALUES("{}","{}","{}","{}","{}") '''.format(_ds, _fn, _checkId, checkResult, checkMsg)
    cur = conn.cursor()
    cur.execute(sql)
    conn.commit()

    return cur.lastrowid

"""
Update a given file+check, if file is None, all files from dataset
"""
def update_filedatacheck(conn, _ds, _fn, _checkId, checkResult, checkMsg):

    # CheckResult can be either a string or a number, but is encoded as a number
    if isinstance(checkResult, str):
        if checkResult == 'ok':
            checkResult = 1
        elif checkResult == 'err':
            checkResult = -1
        elif checkResult == 'warn':
            checkResult = 2
        else:
            checkResult = 0

    if _fn is None:
        sql = '''UPDATE filedatacheck
                 SET checkResult={}, checkMsg="{}"
                 WHERE _ds="{}" AND _checkId="{}";'''.format(checkResult, checkMsg, _ds, _checkId)
    else:
        sql = '''UPDATE filedatacheck
                 SET checkResult={}, checkMsg="{}"
                 WHERE _ds="{}" AND _fn="{}" AND _checkId="{}";'''.format(checkResult, checkMsg, _ds, _fn, _checkId)

    cur = conn.cursor()
    cur.execute(sql)
    conn.commit()

    return cur.lastrowid

"""
Get all the files/checks to be performed (NEW)
"""
def select_filedatachecks_todo(conn, _ds):

    sql = '''SELECT _fname, _checkId
             FROM file
             INNER JOIN filedatacheck ON file._ds = filedatacheck._ds
             WHERE downRequired=1 AND downPerformed=1 AND checkResult=0 AND file._ds="{}"'''.format(_ds)
    cur = conn.cursor()
    cur.execute(sql)
    rows = cur.fetchall()

    return rows

"""
Get all file/checks done already
"""
def select_filedatachecks(conn, _ds, checkId):

    sql = '''SELECT DISTINCT _fname, checkResult, checkMsg
             FROM file
             INNER JOIN filedatacheck ON file._ds = filedatacheck._ds
             WHERE _checkId="{}" AND checkResult!=0 AND file._ds="{}"
             ORDER BY _fname'''.format(checkId, _ds)
    cur = conn.cursor()
    cur.execute(sql)
    rows = cur.fetchall()

    return rows