import influxdb
from flask import current_app, _app_ctx_stack, Flask
from flask.globals import _app_ctx_err_msg
from typing import Generator

_no_influx_msg = """\
No Influx connection is present.

This means that something has overwritten _app_ctx_stack.top.influxdb_db.
"""


class InfluxDB(object):
    def __init__(self, app: Flask = None) -> None:
        """
        Class constructor
        :param app: Flask Application object
        """
        self.app = app
        if app is not None:
            self.init_app(app)

    def init_app(self, app: Flask) -> None:
        """
        Initialize extension for application
        :param app: Flask Application object
        :return:
        """
        app.config.setdefault('INFLUXDB_HOST', 'localhost')
        app.config.setdefault('INFLUXDB_PORT', '8086')
        app.config.setdefault('INFLUXDB_USER', 'root')
        app.config.setdefault('INFLUXDB_PASSWORD', 'root')
        app.config.setdefault('INFLUXDB_DATABASE', None)
        app.config.setdefault('INFLUXDB_SSL', False)
        app.config.setdefault('INFLUXDB_VERIFY_SSL', False)
        app.config.setdefault('INFLUXDB_RETRIES', 3)
        app.config.setdefault('INFLUXDB_TIMEOUT', None)
        app.config.setdefault('INFLUXDB_USE_UDP', False)
        app.config.setdefault('INFLUXDB_UDP_PORT', 4444)
        app.config.setdefault('INFLUXDB_PROXIES', None)
        app.config.setdefault('INFLUXDB_POOL_SIZE', 10)

        app.teardown_appcontext(self.teardown)

        with app.app_context():
            self.database.create(app.config['INFLUXDB_DATABASE'])

    @staticmethod
    def connect() -> influxdb.InfluxDBClient:
        """
        Connect to InfluxDB using configuration parameters
        :return: InfluxDBClient object
        """
        return influxdb.InfluxDBClient(
            host=current_app.config['INFLUXDB_HOST'],
            port=current_app.config['INFLUXDB_PORT'],
            username=current_app.config['INFLUXDB_USER'],
            password=current_app.config['INFLUXDB_PASSWORD'],
            database=current_app.config['INFLUXDB_DATABASE'],
            ssl=current_app.config['INFLUXDB_SSL'],
            verify_ssl=current_app.config['INFLUXDB_VERIFY_SSL'],
            timeout=current_app.config['INFLUXDB_TIMEOUT'],
            retries=current_app.config['INFLUXDB_RETRIES'],
            use_udp=current_app.config['INFLUXDB_USE_UDP'],
            udp_port=current_app.config['INFLUXDB_UDP_PORT'],
            proxies=current_app.config['INFLUXDB_PROXIES'],
            pool_size=current_app.config['INFLUXDB_POOL_SIZE']
        )

    @staticmethod
    def teardown(exception) -> None:
        """
        This is really a sub in case a influxdb input actually does need
        to be able to be torn down
        """
        ctx = _app_ctx_stack.top
        if hasattr(ctx, 'influxdb_db') and ctx.influxdb_db is not None:
            ctx.influxdb_db.close()

    @property
    def connection(self) -> influxdb.InfluxDBClient:
        """
        InfluxDBClient object
        :return:
        """
        ctx = _app_ctx_stack.top
        if ctx is None:
            raise RuntimeError(_app_ctx_err_msg)

        if not hasattr(ctx, 'influxdb_db'):
            ctx.influxdb_db = self.connect()

        if ctx.influxdb_db is None:
            raise RuntimeError(_no_influx_msg)

        return ctx.influxdb_db

    @property
    def close(self) -> callable:
        return self.connection.close

    @property
    def query(self) -> callable:
        return self.connection.query

    @property
    def write(self) -> callable:
        return self.connection.write

    @property
    def write_points(self) -> callable:
        return self.connection.write_points

    @property
    def database(self):
        class Database:
            switch = self.connection.switch_database
            create = self.connection.create_database
            drop = self.connection.drop_database
            all = self.connection.get_list_database

        return Database()

    @property
    def user(self):
        class User:
            switch = self.connection.switch_user
            create = self.connection.create_user
            drop = self.connection.drop_user
            all = self.connection.get_list_users
            password = self.connection.set_user_password

        return User()

    @property
    def policy(self):
        class RetentionPolicy:
            alter = self.connection.alter_retention_policy
            create = self.connection.create_retention_policy
            drop = self.connection.drop_retention_policy
            all = self.connection.get_list_retention_policies

        return RetentionPolicy()

    @property
    def measurement(self):
        class Measurement:
            tag_values = self._tag_values
            tag_keys = self._tag_keys
            drop = self.connection.drop_measurement
            all = self.connection.get_list_measurements

        return Measurement()

    def _tag_values(self,
                    measurement: str,
                    key: str,
                    *args, **kwargs) -> Generator:
        """
        Returns generator with tag's values for measurement.
        :param measurement: name of the measurement
        :param key: tag key
        :return:
        """
        conn = self.connection
        q = f"SHOW TAG VALUES ON {conn._database} FROM {measurement} WITH KEY = {key}"
        result = conn.query(query=q, *args, **kwargs)
        values = (p['value'] for p in result.get_points())
        return values

    def _tag_keys(self,
                  measurement: str,
                  *args, **kwargs) -> Generator:
        """
        Returns generator with tag's keys for measurement.
        :param measurement: name of the measurement
        :return:
        """
        conn = self.connection
        q = f'SHOW TAG KEYS ON {conn._database} FROM {measurement}'
        result = conn.query(query=q, *args, **kwargs)
        keys = (p['tagKey'] for p in result.get_points())
        return keys
