import argparse
import csv
import os
import subprocess
from pathlib import Path

import dotenv
import snowflake.connector
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from snowflake.connector import DictCursor

DEFAULT_LIMIT = 10
BASE_DIR = Path(__file__).parent.parent
dotenv.load_dotenv(BASE_DIR / '.env')


def _generate(table: str, file_path: str, where: str, order_by: str, limit: int):
    print('🔗Connecting to snowflake with private key...')
    passphrase = os.environ.get('SNOWFLAKE_KEY_PASSPHRASE')
    private_key_path = os.environ.get('SNOWFLAKE_PRIVATE_KEY_PATH')
    private_key = os.environ.get('SNOWFLAKE_PRIVATE_KEY')

    if private_key_path:
        with open(private_key_path, 'rb') as key:
            private_key_bytes = key.read()
    elif private_key:
        private_key_bytes = bytes(private_key, 'utf-8')

    passphrase_bytes = bytes(passphrase, 'utf-8')

    decrypted_key = serialization.load_pem_private_key(
        private_key_bytes,
        password=passphrase_bytes,
        backend=default_backend(),
    )

    decrypted_key_bytes = decrypted_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )

    connection_kwargs = dict(
        user=os.environ['SNOWFLAKE_USER'],
        private_key=decrypted_key_bytes,
        account=os.environ['SNOWFLAKE_ACCOUNT'],
        warehouse=os.environ.get('SNOWFLAKE_WAREHOUSE'),
        database=os.environ.get('SNOWFLAKE_DATABASE'),
        schema=os.environ.get('SNOWFLAKE_SCHEMA'),
        login_timeout=60,
    )
    try:
        con = snowflake.connector.connect(**connection_kwargs)
        where_statement = f'WHERE {where}' if where else ''
        order_by = f'ORDER BY {order_by}' if order_by else ''
        query = f'SELECT * FROM {table} {where_statement} {order_by} LIMIT {limit}'
        cur = con.cursor(DictCursor)
        print(f'querying {table} in snowflake...')
        cur.execute(query, timeout=75)
    except Exception as e:
        print(f'Failed connecting to snowflake. Please filter your query to make it faster. {e}')
        exit(1)

    print('✍️ Writing to csv..')
    file_object = Path(file_path)
    file_existed = file_object.exists()

    with file_object.open(mode='a') as opened_file:
        csv_writer = csv.writer(opened_file)
        first_row = True
        for rec in cur:
            if first_row and not file_existed:
                columns = rec.keys()
                print(columns)
                csv_writer.writerow(columns)
            first_row = False
            values = rec.values()
            print(values)
            csv_writer.writerow(values)

    _plant_seed(file_object.name, file_path)


def _plant_seed(file_name: str, file_path:  str):
    print(f'🌱 Planting {file_path} to /seeds directory...')
    try:
        subprocess.call(f'cp {file_path}  {BASE_DIR}/seeds/{file_name}', shell=True)
        print(f'🎉 Finished planting seed: {file_name}')
    except Exception as e:
        print(f'Error: {e}')


def _parse_arguments(args):
    parser = argparse.ArgumentParser(description='Generate DBT seed from Snowflake table.')
    parser.add_argument('-t', '--table', required=True, help='database.schema.table')
    parser.add_argument('-f', '--file',
                        help='filename with .csv extension. By default tablename.csv will be used. If file exists it will be appended with new data')
    parser.add_argument('-w', '--where', help='where clause for the query')
    parser.add_argument('-o', '--orderby', help='order by clause for the query')
    parser.add_argument('-l', '--limit', type=int, help=f'limit number of rows, default {DEFAULT_LIMIT}', default=DEFAULT_LIMIT)
    return vars(parser.parse_args(args))


def main():
    parsed = _parse_arguments(args=None)
    table = parsed['table']
    table_name = table.split('.')[-1]
    file_arg = parsed['file']
    if file_arg:
        file_path = f'tests/fixtures/expected_{file_arg}.csv'
    else:
        file_path = f'tests/fixtures/expected_{table_name}.csv'

    _generate(table, file_path, parsed['where'], parsed['orderby'], parsed['limit'])

    return parsed


if __name__ == '__main__':
    main()