#!/usr/bin/python
# -*- coding: utf-8 -*-

# (c) 2017, Valeriy Solovyov <weldpua2008@gmail.com>

DOCUMENTATION = '''
---
module: gridfs_artifact
short_description: Downloads files from MongoDB GridFS to node
description:
     - Downloads files from GridFS. The remote
       server I(must) have direct access to the remote resource.
options:
    database:
        required: True
        description:
            - the name of a getent database supported by the target system (passwd, group,
              hosts, etc).
    key:
        required: False
        default: ''
        description:
            - key from which to return values from the specified database, otherwise the
              full contents are returned.
    split:
        required: False
        default: None
        description:
            - "character used to split the database values into lists/arrays such as ':' or '\t', otherwise  it will try to pick one depending on the database"
    fail_key:
        required: False
        default: True
        description:
            - If a supplied key is missing this will make the task fail if True
notes:
   - "Not all databases support enumeration, check system documentation for details"
requirements: [ pymongo ]
author: "Valeriy Solovyov"
'''

EXAMPLES = '''
# get file
-   gridfs_artifact:    url="test.txt"    dst="/tmp/test2.txt"    database="cdn_storage"    login_host="login.host"    force=yes
'''

try:
    from pymongo.errors import ConnectionFailure
    from pymongo.errors import OperationFailure
    from pymongo import version as PyMongoVersion
    from pymongo import MongoClient
    import gridfs
except ImportError:
    try:  # for older PyMongo 2.2
        from pymongo import Connection as MongoClient
        import gridfs
    except ImportError:
        pymongo_found = False
    else:
        pymongo_found = True
else:
    pymongo_found = True


# from distutils.version import LooseVersion
# import datetime
import os
import hashlib

__author__ = 'weldpua2008@gmail'

# ==============================================================
# main

class MD5Error(Exception):
    pass


class MD5CompareError(Exception):
    pass



def connect_to_db(module):
    replica_set = module.params['replica_set']
    ssl = module.params['ssl']
    db_name = module.params['database']
    # read_preferences = module.params['read_preferences']
    login_host = module.params["login_host"]
    login_password = module.params["login_password"]
    login_user = module.params["login_user"]
    port = module.params['port']
    replica_set_hosts = module.params['replica_set_hosts']

    try:
        if replica_set:
            if not isinstance(replica_set_hosts, list):
                module.fail_json(msg="replica_set_hosts %s wasn't list" % str(replica_set_hosts))
            if len(replica_set_hosts) < 1:
                module.fail_json(msg="replica_set_hosts is empty" % str(replica_set_hosts))
            _replica_set_hosts = ','.join(replica_set_hosts)

            client = MongoClient('mongodb://%s/' % _replica_set_hosts, int(port), replicaset=replica_set, ssl=ssl)
        else:
            if login_host is None:
                module.fail_json(msg="login_host is empty" % str(replica_set_hosts))
            client = MongoClient(login_host, int(port), ssl=ssl)

        if login_user is not None and login_password is not None:
            client.admin.authenticate(login_user, login_password)

    except ConnectionFailure, error:
        module.fail_json(msg="unable to connect to database: %s" % str(error))

    # db = MongoClient('mongodb://%s/?replicaSet=%s' % (replica_set_hosts, replica_set),read_preference=pymongo.ReadPreference.NEAREST).rdecdn_storage
    db = client[db_name]
    fs = gridfs.GridFS(db)
    return fs

def check_md5(store_path, md5):
    if hashlib.md5(open(store_path, 'rb').read()).hexdigest() != md5:
        raise MD5CompareError("md5 is not equal")

def get_file(fs, key, store_path, force=False, md5=None, check_md5sum=False):

    _md5 = fs.get_version(key).md5
    # _md5 = None
    # try:
    #     _md5 = fs.get_version(key).md5
    # except Exception as error:
    #     pass

    if not os.path.exists(os.path.dirname(store_path)):
        os.makedirs(os.path.dirname(store_path))
    if not os.path.exists(store_path) or force:
        with open(store_path, "wb") as f:
            f.write(fs.get_last_version(key).read())
        try:
            if md5 is not None:
                check_md5(store_path=store_path, md5=md5)
            elif check_md5sum:
                check_md5(store_path=store_path, md5=_md5)

        except MD5CompareError as error:
            os.remove(store_path)
            raise MD5CompareError(str(error))
        except Exception as error:
            os.remove(store_path)
            raise MD5Error("Can't calculate md5 %s" % error)
        return True, _md5
    return False, _md5



def save_file(fs, key, source_path, clear_versions=False):
    md5 = None
    if os.path.exists(source_path):
        if clear_versions:
            try:
                while fs.exists(filename=key):
                    fp = fs.get_last_version(key)
                    fs.delete(fp._id)
            except Exception:
                pass
        with fs.new_file(filename=key) as fp:
            with open(source_path, "rb") as data:
                fp.write(data)

        return True, md5
    else:
        raise ValueError("Not exists '%s'" % source_path)



def main():

    module = AnsibleModule(
        argument_spec=dict(
            url=dict(required=True, aliases=['key', 'file']),
            database=dict(required=True, aliases=['db']),
            store_path=dict(default=None, required=False, aliases=['dest','dst']),
            source_path=dict(default=None, required=False, aliases=['src']),
            force=dict(default=False, required=False),
            md5sum=dict(default=None, required=False, aliases=['md5']),
            check_md5sum=dict(default=True, required=False, aliases=['check_md5sum']),
            clear_versions=dict(default=True, required=False),
            port=dict(default="27017", required=False),
            replica_set=dict(required=False),
            replica_set_hosts=dict(type='list', default=[], required=False),
            login_host=dict(default=None, required=False),
            login_user=dict(default=None, required=False),
            login_password=dict(default=None, required=False),
            ssl=dict(default=False, required=False),
            # state=dict(default="present", choices=["absent", "present"]),
        ),
        supports_check_mode=True
    )
    if not pymongo_found:
        module.fail_json(msg="the python pymongo module is required")

    url = module.params['url'].strip()
    store_path = module.params['store_path']
    source_path = module.params['source_path']
    force = module.params['force']
    clear_versions = module.params['clear_versions']
    md5sum = module.params['md5sum']
    check_md5sum = module.params['check_md5sum']
    kwargs = {"changed": False}
    try:
        fs = connect_to_db(module=module)
        if store_path is not None:
            changed, md5 = get_file(fs=fs, key=url, store_path=store_path.strip(), force=force, md5=md5sum, check_md5sum=check_md5sum)
            kwargs = {"url":url, "store_path": store_path, "src": url, "changed": changed, "md5": md5, "ok": True}
        elif source_path is not None:
            changed, md5 = save_file(fs=fs, key=url, source_path=source_path.strip(), clear_versions=clear_versions)
            kwargs = {"url":url, "source_path": source_path, "clear_versions": clear_versions, "changed": changed, "md5": md5, "ok": True}

    except gridfs.errors.CorruptGridFile as error:
        module.fail_json(msg="file %s is malformed %s" % (url, str(error)) )
    except gridfs.errors.FileExists as error:
        module.fail_json(msg="Trying to create a file %s that already exists: %s" % (url, str(error)) )
    except gridfs.errors.NoFile:
        module.fail_json(msg="Tying to get a non-existent file '%s'" % url )
    except Exception as error:
        module.fail_json(msg="There was an error %s" % str(error))
    module.exit_json(**kwargs)

# import module snippets
from ansible.module_utils.basic import *
# from ansible.module_utils.urls import *
if __name__ == '__main__':
    main()