'''
Connect to Amazon Redshift Databases.

Copyright (C) 2023 Nat Hawkins & Mike Chappelow

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License (version 3 of the License)
as published by the Free Software Foundation.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
'''


# Imports 
import os
import pyodbc
from plinko.configurations import configure_username
from plinko.zoltar import zoltar_ask


def plinko_redshift(environment, 
                    database = None, server = None):
    '''
    Connect to AWS Redshift databases in Kortex

    TODO: Add support for additional databases and servers, currently
    default to REDSHIFT_default_database when connecting (limited service account
    information available at this time)

    Parameters
    ----------
    environment: str
        Kortex environment to connect to (DEV, PROD, QA)
    database: str (optional)
        Specify database to connect to in Kortex environment. If None, REDSHIFT_default_database
        information is pulled from zoltar
    server: str (optional)
        Specify custom server to connect to in Kortex environment. If None, server is determined
        based on environment

    Returns
    -------
    pyodbc.connect
        If connection can be established using pyodbc.connect(), a pyodbc.connect() object
        will be returned. This can be passed to pandas (e.g., pd.read_sql) or similar libraries
        for reading remote data. If a connection cannot be established, None will be returned
    '''
    # Check for environment variable configurations
    if os.getenv("CONNECT_API_KEY") is None: configure_username()

    # DEV
    if environment.upper() == "DEV":
        server_id   = "REDSHIFT_DEV_server"
        uid_id      = "REDSHIFT_DEV_uid"
        pwd_id      = "REDSHIFT_DEV_pwd"

    # PROD
    elif environment.upper() == "PROD":
        server_id   = "REDSHIFT_PROD_server"
        uid_id      = "REDSHIFT_PROD_uid"
        pwd_id      = "REDSHIFT_PROD_pwd"

    # QA
    elif environment.upper() == "QA":
        server_id   = "REDSHIFT_QA_server"
        uid_id      = "REDSHIFT_QA_uid"
        pwd_id      = "REDSHIFT_QA_pwd"

    else:
        print("Please specify PROD, DEV, or QA as your environment.")
        return None

    # Establish connection
    cnxn = pyodbc.connect(Driver = 'redshift',
                          servername = zoltar_ask(server_id),
                          database = zoltar_ask("REDSHIFT_default_database") if database is None else database,
                          UID = zoltar_ask(uid_id),
                          PWD = zoltar_ask(pwd_id),
                          Port = zoltar_ask("REDSHIFT_default_port"))

    return cnxn