import os
import json
import boto3
import psycopg2
from psycopg2.extras import RealDictCursor
import logging

log = logging.getLogger()
log.setLevel(logging.INFO)
REGION = os.environ["AWS_REGION"] if "AWS_REGION" in os.environ else "eu-west-1"
S3_BUCKET_DEVEL = "frontend-api-devel-filestore"


def get_secret(secret_name):
    # Create a Secrets Manager client
    session = boto3.session.Session()
    client = session.client(
        service_name='secretsmanager',
        region_name=REGION
    )
    get_secret_value_response = json.loads(client.get_secret_value(SecretId=secret_name)['SecretString'])
    return get_secret_value_response


DB_CONN = {'fansifter-rds': None,
           'fansifter-rds-test': None,
           'fansifter-rds-live': None,
           'sandbox': None}

DB_PORT = {'fansifter-rds': 5432,
           'fansifter-rds-test': 5433,
           'fansifter-rds-live': 5434,
           'sandbox': None}


def get_rds_connection(conn_name):
    global DB_CONN
    if DB_CONN[conn_name] is None:
        try:
            secret = get_secret(conn_name)
            DB_CONN[conn_name] = psycopg2.connect(host=secret['host'], #'localhost',
                                                    port=DB_PORT[conn_name],
                                                    database=secret['dbname'],
                                                    user=secret['username'],
                                                    password=secret['password'],
                                                    cursor_factory=RealDictCursor)
        except Exception as e:
            log.exception(e)
    return DB_CONN[conn_name]


def rds_query(connection_name, query, vars=None, fetch=False):
    with get_rds_connection(connection_name) as conn:
        with conn.cursor() as cur:
            cur.execute(query, vars=vars)
            if fetch:
                return cur.fetchall()


databases = [
    'fansifter-rds',
#    'fansifter-rds-test',
#    'fansifter-rds-live',
]

other_buckets = set()
for connection in databases:
    all_chemas = rds_query(connection,
                           "select table_schema, table_name from information_schema.tables where table_name IN ('meta_data' )", fetch=True)
    print(f"Found {len(all_chemas)} schemas in database {connection}")
    for schema, table in [sch.values() for sch in all_chemas]:
        print(f'Working on: {schema}.{table}')
        statements = []
        vars = {}
        # try:
        #     for sub in rds_query(connection, f"SELECT * from {schema}.collection_airflow", fetch=True):
        #         pass
        #     print(f'All good.')
        # except Exception as e:
        #     print(f'Oops, missing fields: {schema}.{table}')
        statements.extend([
                 f"""CREATE TABLE IF NOT EXISTS {schema}.collection_airflow
(
    collection_id  integer,
    dag_id         varchar(250),
    dag_run_id     varchar(250),
    execution_date timestamp WITH TIME ZONE,
    fan_count      integer
);""",
f"""CREATE INDEX IF NOT EXISTS idx_collection_id
    ON {schema}.collection_airflow (collection_id);"""
             ])

        for statement in statements:
            try:
                rds_query(connection, statement, vars)
            except Exception as e:
                log.exception("oops", exc_info=e)