from ewah.hooks.base import EWAHBaseHook

from airflow.utils.file import TemporaryDirectory

import os
from copy import deepcopy
from pymongo import MongoClient
from pymongo import ASCENDING as asc
from tempfile import NamedTemporaryFile


class EWAHMongoDBHook(EWAHBaseHook):

    _ATTR_RELABEL = {
        "default_database": "schema",
        "uri": "password",
        "username": "login",
    }

    conn_name_attr: str = "ewah_mongodb_conn_id"
    default_conn_name: str = "ewah_mongodb_default"
    conn_type: str = "ewah_mongodb"
    hook_name: str = "EWAH MongoDB Connection"

    @staticmethod
    def get_ui_field_behaviour() -> dict:
        return {
            "hidden_fields": ["extra"],
            "relabeling": {
                "host": "Host",
                "schema": "Default database",
                "login": "Username",
                "password": "Uri OR password",
            },
        }

    @staticmethod
    def get_connection_form_widgets() -> dict:
        """Returns connection widgets to add to connection form"""
        from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
        from wtforms import StringField, SelectField, BooleanField, PasswordField
        from ewah.ewah_utils.widgets import EWAHTextAreaWidget

        return {
            "extra__ewah_mongodb__conn_style": SelectField(
                "Connection Style",
                choices=[("uri", "URI"), ("credentials", "Credentials")],
            ),
            "extra__ewah_mongodb__tls": BooleanField("SSL / TLS?"),
            "extra__ewah_mongodb__tls_insecure": BooleanField(
                "TLS: Allow insecure connections? Aka 'tlsInsecure'"
            ),
            "extra__ewah_mongodb__ssl_cert": StringField(
                "SSL Certificate",
                widget=EWAHTextAreaWidget(rows=12),
            ),
            "extra__ewah_mongodb__ssl_private": StringField(
                "SSL Private Key",
                widget=EWAHTextAreaWidget(rows=12),
            ),
            "extra__ewah_mongodb__ssl_password": PasswordField(
                "SSL Certificate / Private Key Password",
            ),
            "extra__ewah_mongodb__auth_source": StringField(
                "Auth Source", widget=BS3TextFieldWidget()
            ),
            "extra__ewah_mongodb__auth_mechanism": StringField(
                "Auth Mechanism", widget=BS3TextFieldWidget()
            ),
            "extra__ewah_mongodb__ssh_conn_id": StringField(
                "SSH Connection ID", widget=BS3TextFieldWidget()
            ),
        }

    @property
    def mongoclient(self):
        if not hasattr(self, "_mc"):
            if self.conn.ssh_conn_id:
                if self.conn.conn_style == "uri":
                    raise Exception("Cannot have SSH tunnel with uri connection type!")
                if not hasattr(self, "_ssh_hook"):
                    self._ssh_hook = EWAHBaseHook.get_hook_from_conn_id(
                        conn_id=self.conn.ssh_conn_id
                    )
                    self.local_bind_address = self._ssh_hook.start_tunnel(
                        self.conn.host, self.conn.port
                    )
            else:
                self.local_bind_address = (self.conn.host, self.conn.port)

            conn_kwargs = {"tz_aware": True}
            if self.conn.conn_style == "uri":
                conn_kwargs["host"] = self.conn.uri
            else:
                conn_kwargs["host"] = self.local_bind_address[0]
                conn_kwargs["port"] = self.local_bind_address[1]
                if self.conn.username:
                    conn_kwargs["username"] = self.conn.username
                if self.conn.password:
                    conn_kwargs["password"] = self.conn.password

            with TemporaryDirectory() as tmp_dir:
                if self.conn.tls:
                    conn_kwargs["tls"] = True
                with NamedTemporaryFile(dir=tmp_dir) as ssl_cert:
                    with NamedTemporaryFile(dir=tmp_dir) as ssl_private:
                        if self.conn.ssl_cert:
                            ssl_cert.write(self.conn.ssl_cert.encode())
                            ssl_cert.seek(0)
                            conn_kwargs["ssl_certfile"] = os.path.abspath(ssl_cert.name)
                        if self.conn.ssl_private:
                            ssl_private.write(self.conn.ssl_private.encode())
                            ssl_private.seek(0)
                            conn_kwargs["ssl_keyfile"] = os.path.abspath(
                                ssl_private.name
                            )
                        if self.conn.ssl_password:
                            conn_kwargs[
                                "tlsCertificateKeyFilePassword"
                            ] = self.conn.ssl_password
                        if self.conn.tls_insecure:
                            conn_kwargs["tlsInsecure"] = True
                        if self.conn.auth_source:
                            conn_kwargs["authSource"] = self.conn.auth_source
                        if self.conn.auth_mechanism:
                            conn_kwargs["authMechanism"] = self.conn.auth_mechanism
                        self._mc = MongoClient(**conn_kwargs)

        return self._mc

    def get_data_in_batches(
        self,
        collection,
        database=None,
        filter_expression=None,
        batch_size: int = 100000,
    ):

        mongo_database = self.mongoclient.get_database(
            name=database or self.conn.default_database
        )
        mongo_collection = mongo_database[collection]
        self.log.info(
            "\n\nFetching data with filter expression:\n{0}\n\n".format(
                filter_expression
            )
        )
        filter_expression = deepcopy(filter_expression) or {}
        cursor = mongo_collection.find(filter_expression)

        data = []
        for document in cursor:
            data += [document]
            if len(data) >= batch_size:
                yield data
                data = []
        if data:
            yield data

    def close(self):
        if hasattr(self, "_mc"):
            if hasattr(self, "_ssh_hook"):
                self._ssh_hook.stop_tunnel()
                del self._ssh_hook
            self._mc.close()
            del self._mc
