import pandas as pd

from tableau_api_lib import TableauServerConnection
from tableau_api_lib.utils import querying, flatten_dict_column


#  using personal access tokens is preferred; otherwise, comment those details out and use username / password
tableau_server_config = {
    'my_env': {
        'server': 'https://10ax.online.tableau.com',  # replace with your own server
        'api_version': '3.8',  # replace with your REST API version
        'personal_access_token_name': '<PAT NAME>',
        'personal_access_token_secret': '<PAT SECRET',
#         'username': '<USERNAME>',
#         'password': '<PASSWORD>',
        'site_name': 'your-pretty-site-name',  # if accessing your default site, set this to ''
        'site_url': 'YourSiteContentUrl'  # if accessing your default site, set this to ''
    }
}

NEW_OWNER_USERNAME = 'estam'  # replace this with your desired owner's username
USER_TO_REMOVE = 'estam2'  # replace this with the username of the user being removed


def change_content_owner(conn, content_type, content_id, new_owner_id):
    if content_type == 'project':
        response = "Cannot update project owners via REST API (yet)."
    if content_type == 'workbook':
        response = conn.update_workbook(workbook_id=content_id, new_owner_id=new_owner_id)
    if content_type == 'datasource':
        response = conn.update_data_source(datasource_id=content_id, new_owner_id=new_owner_id)
    if content_type == 'flow':
        response = conn.update_flow(flow_id=content_id, new_owner_id=new_owner_id)
    return response


conn = TableauServerConnection(tableau_server_config, env='tableau_server')
conn.sign_in()

projects_df = querying.get_projects_dataframe(conn)
projects_df['content_type'] = 'project'

workbooks_df = querying.get_workbooks_dataframe(conn)
workbooks_df['content_type'] = 'workbook'

datasources_df = querying.get_datasources_dataframe(conn)
datasources_df['content_type'] = 'datasource'

flows_df = querying.get_flows_dataframe(conn)
flows_df['content_type'] = 'flow'

content_dfs = [projects_df, workbooks_df, datasources_df, flows_df]

content_ownership_df = pd.DataFrame()

cols_to_keep = ['id', 'name', 'content_type', 'owner_id']
for index, content_df in enumerate(content_dfs):
    if not content_df.empty:
        temp_df = content_df.copy()
        temp_df = flatten_dict_column(temp_df, keys=['id'], col_name='owner')
        temp_df = temp_df[cols_to_keep]
        content_ownership_df = content_ownership_df.append(temp_df)

users_df = querying.get_users_dataframe(conn)
new_owner_id = list(users_df[users_df['name'] == NEW_OWNER_USERNAME]['id']).pop()
user_id_to_remove = list(users_df[users_df['name'] == USER_TO_REMOVE]['id']).pop()

owner_swap_df = content_ownership_df[content_ownership_df['owner_id'] == user_id_to_remove]

for index, row in owner_swap_df.iterrows():
    response = change_content_owner(conn, row['content_type'], row['id'], new_owner_id)
    print(response.json())
    
response = conn.remove_user_from_site(user_id_to_remove)

conn.sign_out()