"""Database connection and data retrieval functions."""

import psycopg2
from psycopg2.extras import DictCursor

from ioda import config
from ioda import const
from ioda import misc
from ioda import queries


def get_connection():
    """Connect to RedShift DB.

    Returns:
        Connection object
    """
    try:
        print('Connecting to RedShift...')
        con = psycopg2.connect(
            dbname=config.REDSHIFT_DB,
            host=config.REDSHIFT_HOST,
            port=config.REDSHIFT_PORT,
            user=config.REDSHIFT_USER,
            password=config.REDSHIFT_PASSWORD)
        print('Connection succeeded.')
    except Exception as ex:
        print('Failed to connect to DB. \n{}'.format(ex))
        return None
    return con


@misc.retry()
def get_batch(con, offset, limit, fields=queries.ALL_FIELDS):
    """Get a number of rows from DB for generating the CSV file.

    Args:
        con (Connection): DB connection object
        offset (int): offset from start
        limit (int): number of rows to return
        fields (str): comma-separated fields to query (* by default)

    Returns:
        list(DictRow): list of rows from DB
    """
    cur = con.cursor(cursor_factory=DictCursor)
    query = queries.IODA_DATA.format(fields)
    cur.execute(query, {
        const.OFFSET: offset,
        const.LIMIT: limit
    })
    data = cur.fetchall()
    cur.close()
    return data


@misc.retry()
def get_total_records_count(con):
    """Get total number of records to process

    Args:
        con (Connection): DB connection object

    Returns:
        int: number of records to process
    """
    cur = con.cursor(cursor_factory=DictCursor)
    cur.execute(queries.COUNT)
    result = cur.fetchone()
    count = int(result[const.ASSET_COUNT])
    cur.close()
    return count