import logging
import sys
from typing import Any, Dict, List
from sqlalchemy import not_, and_, inspect
from airflow.models import Connection
from airflow import settings
from zenyatta.db.sql import SQL
from zenyatta.tasks import create_rds_resource
from zenyatta.aws.rds import RDSResource, SingleInstanceDriver
from zenyatta.common.util import update_aws_resource_variable, dump_to_json, get_aws_resource_variable
from zenyatta.common import get_zenyatta_base_connections


def fetch_connections_partition_info(connections: List[Connection],
                                     ts_nodash: str) -> List:
    """
    One time fetch table partition information for spark apps
    """
    sesh = settings.Session()
    try:
        for conn in connections:
            var_key = conn.conn_id + '-pitr-table-partition'
            partitions = {"dummy": None}
            update_aws_resource_variable(var_key, 'partitions', partitions, is_append=False)
            update_aws_resource_variable(var_key, 'ts_nodash', ts_nodash, is_append=False)
            logging.info("inserting {var_key} partition as {partitions}".format(**locals()))
    except Exception as e:
        logging.info("caught exception: {}".format(e))


def init_db_tables_variables(connections: List[Connection], ts_nodash: str):
    """
    Init db tables and store it as Variables
    """
    sesh = settings.Session()
    try:
        for conn in connections:
            # db_engine = SQL.create_sql_driver(conn.conn_type, conn.conn_id).create_sql_engine()
            # tables = db_engine.table_names()
            tables = ['dummy_only_know_when_pitr']
            var_key = conn.conn_id + '-pitr-schema'
            if get_aws_resource_variable(var_key, 'db_tables') is None:
                update_aws_resource_variable(var_key, 'db_tables', tables, is_append=False)
            if get_aws_resource_variable(var_key, 'ts_nodash') is None:
                update_aws_resource_variable(var_key, 'ts_nodash', ts_nodash, is_append=False)
            logging.info("inserting {var_key} schema as {tables}".format(**locals()))
    except Exception as e:
        logging.info("caught exception: {}".format(e))


def set_table_priority():
    """
    This key-value map represents how many instances each table needs when they are submitted to Yarn.
    Only when a table needs more than 1 instnace do we need to set it up, otherwise it defaults to 1,
    e.g. {'justintv_prod.users': 4, 'cohesion-following.associations': 10}
    """
    priority_map = {'conn_id.table': 1}
    var_key = 'priority_map'
    update_aws_resource_variable(var_key, var_key, priority_map, False)


if __name__ == '__main__':
    connections = get_zenyatta_base_connections()
    if len(connections) == 0:
        logging.info("could not find valid connections in airflow. this is unusual."
                     "please add valid connections.")
        sys.exit(0)
    logging.info("populating db schema for dags.")
    ts_nodash = "20170101T000000"
    init_db_tables_variables(connections, ts_nodash)
    logging.info("creating spark table partition for connections.")
    output = settings.conf.get('spark', 'spark_home') + '/spark_table_partition.json'
    fetch_connections_partition_info(connections, ts_nodash)
    set_table_priority()
