#!/usr/bin/env python
# coding: utf8

# Script to create superset users

import pandas as pd
from sqlalchemy import create_engine
import datetime
import pexpect
import os

DB = create_engine('mysql+pymysql://superset:superset@localhost/superset')
CSV_USER = 'ext_source/superset_users.csv'
CSV_PASSWORD = 'secrets/email_password.csv'
DOCKER_PATH = '/usr/local/bin'  # Useful for Mac


def main():
    create_missing_users()
    # reset_all_passwords()


def create_missing_users():
    """ Create users present in csv file but missing in Superset"""
    print("== Creating users from csv file, if they do not exist in Superset")
    df_existing_users = pd.read_sql_table('ab_user', DB)
    df_password_users = pd.read_csv(CSV_PASSWORD, sep=';')

    df_csv_users = pd.read_csv(CSV_USER, sep=';')
    df_csv_users['username'] = concat_username(df_csv_users['first_name'], df_csv_users['last_name'])
    now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    df_csv_users['now'] = "TIMESTAMP '{}'".format(now)

    for user_row in df_csv_users.iterrows():
        user_dict = user_row[1].to_dict()

        # Skip existing users
        to_insert = True
        for key in ['username', 'email']:
            if user_dict[key] in df_existing_users[key].values:
                print("{} {} already exists in Superset.".format(key, user_dict[key]))
                to_insert = False
                break

        # Insert others
        if to_insert:
            print("Creating user : {first_name}, {last_name}, {username}, {email}".format(**user_dict))
            query = """
            INSERT INTO ab_user (first_name, last_name, username, active, email, created_on, changed_on, created_by_fk, changed_by_fk, login_count, fail_login_count)
            VALUES ('{first_name}', '{last_name}', '{username}', 1, '{email}', {now}, {now}, 1, 1, 0, 0);
            """.format(**user_dict)
            connection = DB.connect()
            connection.execute(query)
            connection.close()

            password = get_corresponding_value(df_password_users, 'email', user_dict['email'], 'password')
            if password is None:
                print("No password defined for user with email '{email}'. You will have to reset it manually".format(**user_dict))
            else:
                print("Reset password of user '{username}' with email '{email}'".format(**user_dict))
                reset_password(username=user_dict['username'], password=password)


def concat_username(first_name, last_name):
    return first_name + '.' + last_name


def reset_all_passwords():
    """ Reset all passwords from password file"""
    print("== Reset all users passwords from secret file")
    df_csv_users = pd.read_csv(CSV_USER, sep=';')
    df_password_users = pd.read_csv(CSV_PASSWORD, sep=';')

    for user_row in df_password_users.iterrows():
        user_dict = user_row[1].to_dict()
        email = user_dict['email']
        password = user_dict['password']

        first_name = get_corresponding_value(df_csv_users, 'email', email, 'first_name')
        last_name = get_corresponding_value(df_csv_users, 'email', email, 'last_name')
        if first_name is None or last_name is None:
            print("No user found for email {} in {}\n".format(email, CSV_USER))
            continue
        username = concat_username(first_name, last_name)
        reset_password(username, password)


def reset_password(username, password):
    """ Reset password of user, by executing fabmanager script in docker
    """
    env = os.environ.copy()
    env['PATH'] = "{}{}{}".format(DOCKER_PATH, os.pathsep, env['PATH'])
    child = pexpect.spawnu(
        """docker exec -ti superset sh -c \"fabmanager reset-password --username {} --app superset\"""".format(username),
        env=env,
        timeout=3)
    child.expect('Password:')
    child.sendline(password)
    child.expect('Repeat for confirmation:')
    child.sendline(password)


def get_corresponding_value(df, column_to_equal, value, column_to_get='id'):
    """ Get corresponding value in column_to_get, for column_to_equal == value

    Suppose that only one row match column_to_equal == value
    """
    ids = list(df[df[column_to_equal] == value][column_to_get])
    if ids:
        if len(ids) > 1:
            raise ValueError("Value {} is present multiple times in columns {}".format(value, column_to_equal))
        return ids[0]
    else:
        return None

if __name__ == "__main__":
    main()