#!/usr/bin/env python
'''
  Copyright European Organization for Nuclear Research (CERN) 2013

  Licensed under the Apache License, Version 2.0 (the "License");
  You may not use this file except in compliance with the License.
  You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0

  Authors:
  - Vincent Garonne, <vincent.garonne@cern.ch>, 2015

  Probe to check datatype, project and convention from AMI
'''

import sys
import traceback

from rucio.db.sqla.session import get_session
from rucio.core.scope import add_scope

# Exit statuses
OK, WARNING, CRITICAL, UNKNOWN = 0, 1, 2, 3


if __name__ == "__main__":
    try:
        session = get_session()

        # Import datatypes
        query = '''MERGE INTO atlas_rucio.did_key_map D
USING (select datatype from ATLAS_AMI_PRODUCTION_01_R.data_type@AMI_ATLR.CERN.CH where writestatus='valid') S
ON (D.key='datatype' and S.datatype = D.value)
WHEN NOT MATCHED THEN INSERT (key,value,created_at, updated_at)
VALUES ('datatype', S.datatype, sys_extract_utc(systimestamp), sys_extract_utc(systimestamp))'''
        session.execute(query)
        session.commit()

        # Import projects
        query = '''MERGE INTO atlas_rucio.did_key_map D
USING (select distinct(projecttag) as value from ATLAS_AMI_PRODUCTION_01_R.projects@AMI_ATLR.CERN.CH where writestatus='valid') S
ON (D.key='project' and S.value = D.value)
WHEN NOT MATCHED THEN INSERT (key, value,created_at, updated_at)
VALUES ('project', S.value, sys_extract_utc(systimestamp), sys_extract_utc(systimestamp))'''
        session.execute(query)
        session.commit()

        # Import scopes
        query = '''select projecttag
from ATLAS_AMI_PRODUCTION_01_R.projects@AMI_ATLR.CERN.CH
where writestatus='valid'
and not exists (select 1 from atlas_rucio.scopes where scope = projecttag)
and projecttag not in ('user', 'group', 'group10', 'user10', 'user09', 'group', 'data', 'condR2')'''
        for project, in session.execute(query):
            add_scope(scope=project, account='root')

        # import naming convention
        query = '''select projecttag, nomenclaturetemplate
from ATLAS_AMI_PRODUCTION_01_R.projects@AMI_ATLR.CERN.CH A,   ATLAS_AMI_PRODUCTION_01_R.nomenclature@AMI_ATLR.CERN.CH B
where A.writestatus='valid' and B.identifier = A.NOMENCLATUREFK and B.nomenclaturestatus = 'VALID'
and nomenclaturetemplate in ('project.runNumber.streamName.prodStep.dataType.Version', 'project.datasetNumber.physicsShort.prodStep.dataType.Version')'''
        for project, naming_convention in session.execute(query):
            try:
                # regexp = '^(?P<project>%(project)s)\.(?P<run_number>\d+)\.(?P<stream_name>[a-zA-Z0-9\_\-]+)\.(?P<prod_step>[a-zA-Z0-9\_\-]+)\.(?P<datatype>\w+)(?:\.(?P<version>[a-zA-Z\'_\'-.0-9]{1,32}))?$' % locals()
                regexp = '^(?P<project>%(project)s)\.(?P<run_number>\d+)\.(?P<stream_name>[a-zA-Z0-9\_\-]+)\.(?P<prod_step>[a-zA-Z0-9\_\-]+)\.(?P<datatype>\w+)(?:\.(?P<version>\w+))?$' % locals()
                query = '''MERGE INTO atlas_rucio.naming_conventions D
USING (SELECT '%(project)s' as scope, '%(regexp)s' as regexp FROM DUAL) S
ON (D.scope = S.scope)
WHEN NOT MATCHED THEN INSERT(scope, regexp, convention_type, updated_at, created_at)
VALUES('%(project)s', '%(regexp)s', 'DATASET', sys_extract_utc(systimestamp), sys_extract_utc(systimestamp))''' % locals()
                session.execute(query)
                session.commit()
            except Exception as error:
                print traceback.format_exc(error)
            except:
                print traceback.format_exc(error)

    except Exception as error:
        print traceback.format_exc(error)
        sys.exit(UNKNOWN)
    sys.exit(OK)
