import logging
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from zenyatta.common import get_airflow_connection
import psycopg2


class SQL:
    def __init__(self, connection_id):
        self.conn_id = connection_id

    @staticmethod
    def create_sql_driver(sql_type, connection_id=None):
        sql_type = sql_type.lower()
        if 'postgres' in sql_type:
            return PostgresSQL(connection_id)
        elif 'mysql' in sql_type or 'aurora' in sql_type:
            return MySQL(connection_id)

    def engine_string(self, host: str=None, port: str=None, login: str=None,
                      password: str=None, schema: str=None) -> str:
        raise NotImplementedError()

    def _engine_string(self, host: str=None, port: str=None, login: str=None,
                       password: str=None, schema: str=None) -> str:
        if not login:
            connection = get_airflow_connection(self.conn_id)
            login = connection.login
        if not password:
            connection = get_airflow_connection(self.conn_id)
            password = connection.get_password()
        if not host:
            connection = get_airflow_connection(self.conn_id)
            host = connection.host
        if not port:
            connection = get_airflow_connection(self.conn_id)
            port = connection.port
        if not schema:
            connection = get_airflow_connection(self.conn_id)
            schema = connection.schema

        return "://{login}:{password}@{host}:{port}/{schema}".format(login=login,
                                                                     password=password,
                                                                     host=host,
                                                                     port=port,
                                                                     schema=schema)

    def create_sql_engine(self, host: str=None, port: str=None, login: str=None, password: str=None,
                          schema: str=None, stream_results=True) -> Engine:

        return create_engine(self.engine_string(
            host=host,
            port=port,
            login=login,
            password=password,
            schema=schema)).execution_options(stream_results=stream_results)


class PostgresSQL(SQL):
    def engine_string(self, host: str=None, port: str=None, login: str=None,
                      password: str=None, schema: str=None) -> str:

        return "postgresql" + self._engine_string(host=host, port=port,
                                                  login=login, password=password, schema=schema)


class MySQL(SQL):
    def engine_string(self, host: str=None, port: str=None, login: str=None,
                      password: str=None, schema: str=None) -> str:

        return "mysql+pymysql" + self._engine_string(host=host, port=port,
                                                     login=login, password=password, schema=schema)


def get_ec2_sql_conn_id(conn_id, ts_nodash):
    return conn_id + '-' + ts_nodash


def run_table_update(host: str=None, dbname: str=None, login: str=None, password: str=None,
                     stmt: str=None) -> bool:

    db_conn = None
    try:
        db_conn = psycopg2.connect("host='{host}' dbname='{dbname}' user='{login}' password='{password}'"
                                   .format(**locals()))
        cur = db_conn.cursor()
        cur.execute(stmt)
        db_conn.commit()
        return True
    except psycopg2.DatabaseError as e:
        if db_conn:
            db_conn.rollback()
        logging.info("error run_table_update {}".format(str(e)))
    finally:
        if db_conn:
            db_conn.close()
