import asyncio
import pandas as pd
from . import utils

def initialize_cache_run(pod_id, **config):
    """
    Creates a new entry in the cache run table for the current address run
    Args:
        pod_id: the Airflow pod id performing the current address run
        **config: the configuration of the current address run

    Returns: the address run id of the current address run

    """
    utils.set_env_config(**config)
    return utils.set_address_run_initial(pod_id, **config)


def end_cache_run(address_run_id, pod_id, status, exception=None, **update_metadata):
    """
    Updates an existing entry in the cache run table with the status, metrics,
    and any generated exceptions from the address run
    Args:
        address_run_id: the id of the address run to be updated
        pod_id: the Airflow pod id that performed the address run
        status: the status of the address run, either FAILED or SUCCESS
        exception: the exception generated by the address run (if there is one)
        **update_metadata: a dict of column names and values to be updated

    Returns: None

    """
    utils.set_address_run_end(
        address_run_id, pod_id, status, exception, **update_metadata
    )


def process_address(request_street_address, request_city, request_state, request_zipcode, address_run_id, is_async=True):
    """
    Gets address information for one address from the Earth Define API and saves it to the cache
    Args:
        request_street_address: the street address to be sent to Earth Define
        request_city: the city of the address to be sent
        request_state: the state of the address to be sent
        request_zipcode: the zipcode of the address to be sent
        address_run_id: the id of the current address run
        is_async: boolean, if True then the API will be accessed asynchronously, otherwise synchronously

    Returns: a DataFrame of address information from Earth Define

    """
    request_df = pd.DataFrame([[request_street_address,
                                request_city,
                                request_state,
                                request_zipcode]],
                              columns=["REQUEST_STREET_ADDRESS",
                                       "REQUEST_CITY",
                                       "REQUEST_STATE",
                                       "REQUEST_ZIPCODE"])
    return process_addresses(request_df, address_run_id, is_async)


def process_addresses(request_address_df, address_run_id=None, is_async=True):
    """
    Gets address information for multiple addresses from the Earth Define API and saves it to the cache
    Args:
        request_address_df: DataFrame of address to be sent to Earth Define
            DataFrame Columns:
                REQUEST_STREET_ADDRESS: the street address component of the requested addresses
                REQUEST_CITY: the city component of the requested addresses
                REQUEST_STATE: the 2-letter abbreviated state component of the requested addresses
                REQUEST_ZIPCODE: the 5-digit zipcode component of the requested addresses
        address_run_id: the id of the current address run
        is_async: boolean, if True then the API will be accessed asynchronously, otherwise synchronously

    Returns: a DataFrame of address information from Earth Define

    """
    if is_async:
        api_response_df = asyncio.run(
            utils.get_api_address_bulk(request_address_df, update_queue=True)
        )
    else:
        loop = asyncio.new_event_loop()
        api_response_df = loop.run_until_complete(
            utils.get_api_address_bulk(request_address_df, update_queue=False)
        )
        loop.close()
    utils.set_addresses(api_response_df, address_run_id)
    return api_response_df


def get_address(request_street_address, request_city, request_state, request_zipcode):
    """
    Gets address information for one address from the cache, if present. If not, information will be
    retrieved from the Earth Define API and added to the cache
    Args:
        request_street_address: the street address whose information will be retrieved
        request_city: the city of the address whose information will be retrieved
        request_state: the state of the address to retrieve
        request_zipcode: the zipcode of the address to retrieve

    Returns: a DataFrame of address information from either the cache or the Earth Define API

    """
    request_street_address = request_street_address.upper()
    request_city = request_city.upper()
    request_state = request_state.upper()
    request_zipcode = request_zipcode[:5]
    response_df = utils.get_address(request_street_address, request_city, request_state, request_zipcode)
    if response_df.shape[0] == 0:
        return process_address(request_street_address, request_city, request_state, request_zipcode, None, is_async=False)
    return response_df


def get_address_bulk(request_address_df):
    """
    Gets address information for multiple addresses from the cache, if present. If not, information will be
    retrieved from the Earth Define API and added to the cache
    Args:
        request_address_df: a DataFrame of addresses whose information will be retrieved
            DataFrame Columns:
                REQUEST_STREET_ADDRESS: the street address component of the requested addresses
                REQUEST_CITY: the city component of the requested addresses
                REQUEST_STATE: the 2-letter abbreviated state component of the requested addresses
                REQUEST_ZIPCODE: the 5-digit zipcode component of the requested addresses

    Returns: a DataFrame of address information from either the cache or the Earth Define API

    """
    request_address_df["REQUEST_STREET_ADDRESS"] = request_address_df[
        "REQUEST_STREET_ADDRESS"
    ].str.upper()
    request_address_df["REQUEST_CITY"] = request_address_df[
        "REQUEST_CITY"
    ].str.upper()
    request_address_df["REQUEST_STATE"] = request_address_df[
        "REQUEST_STATE"
    ].str.upper()
    request_address_df['REQUEST_ZIPCODE'] = request_address_df['REQUEST_ZIPCODE'].str[:5]

    cache_df = utils.get_address_bulk(request_address_df)
    not_in_cache_df = pd.merge(
        request_address_df,
        cache_df,
        on=["REQUEST_STREET_ADDRESS", "REQUEST_CITY", "REQUEST_STATE", "REQUEST_ZIPCODE"],
        how="outer",
        indicator=True,
    ).query('_merge=="left_only"')
    if not_in_cache_df.shape[0] > 0:
        earth_define_df = process_addresses(not_in_cache_df, is_async=False)
        return pd.concat([cache_df, earth_define_df])
    return cache_df


def insert_queue_bulk(request_address_df):
    """
    Adds multiple addresses to the address queue (to be processed in the future)
    Args:
        request_address_df: a DataFrame of addresses to be added to the queue
            DataFrame Columns:
                REQUEST_STREET_ADDRESS: the street address component of the requested addresses
                REQUEST_CITY: the city component of the requested addresses
                REQUEST_STATE: the 2-letter abbreviated state component of the requested addresses
                REQUEST_ZIPCODE: the 5-digit zipcode component of the requested addresses

    Returns: None

    """
    utils.insert_queue_bulk(request_address_df)


def insert_queue(request_street_address, request_city, request_state, request_zipcode):
    """
    Adds one address to the address queue (to be processed in the future)
    Args:
        request_street_address: the street address to be added to the queue
        request_city: the city of the address to be added
        request_state: the state of the address to be added
        request_zipcode: the zipcode of the address to be added

    Returns: None

    """
    request_df = pd.DataFrame([[request_street_address, request_city, request_state, request_zipcode]],
                              columns=["REQUEST_STREET_ADDRESS", "REQUEST_CITY", "REQUEST_STATE", "REQUEST_ZIPCODE"])
    insert_queue_bulk(request_df)


def get_address_queue_count():
    """
    Gets the number of addresses currently in the queue
    Returns: the number of addresses in the queue

    """
    return utils.get_address_queue_count()


def get_address_queue_records(batch_size, pod_id):
    """
    Retrieves multiple records from the front of the address queue
    Args:
        batch_size: the maximum number of records to be retrieved
        pod_id: the Airflow pod id

    Returns: DataFrame of address records from the address queue

    """
    return utils.get_address_queue_records(batch_size, pod_id)


def delete_queue(request_street_address, request_city, request_state, request_zipcode):
    """
    Removes one address from the queue, if present
    Args:
        request_street_address: the street address to be deleted from the queue
        request_city: the city of the address to be deleted
        request_state: the state of the address to be deleted
        request_zipcode: the zipcode of the address to be deleted

    Returns: None

    """
    utils.delete_queue_single(request_street_address, request_city, request_state, request_zipcode)


def delete_queue_bulk(address_run_id):
    """
    Removes addresses from the queue that have been processed during a given address run
    Args:
        address_run_id: the address run whose queued records will be deleted

    Returns: None

    """
    utils.delete_queue_bulk(address_run_id)


def get_ttl_addresses_for_processing(batch_size, pod_id):
    """
    Retrieves multiple records from the cache that are considered stale
    Args:
        batch_size: the maximum number of records to retrieve
        pod_id: the Airflow pod id of the current address run

    Returns: a DataFrame of stale records from the cache

    """
    return utils.get_ttl_addresses_for_processing(batch_size, pod_id)


def reset_queue():
    """
    Sets the number of process attempts to 0 for any non-failed records in the address queue
    Returns: None

    """
    utils.reset_queue()


def reset_cache_pod_ids():
    """
    Unsets the POD_ID column for all records in the address cache
    Returns: None

    """
    utils.reset_cache_pod_ids()