import boto3
import json
import time
import os
import subprocess
from botocore.config import Config as BotoCoreConfig
from dotenv import load_dotenv

# Load env file if it exists
load_dotenv(verbose=True)

ENVIRONMENT = os.environ.get('Environment', 'dev')
MYSQL_CMD_USER = os.environ.get('MYSQL_CMD_USER')
OUTFILE_DIR = os.environ.get('OUTFILE_DIR')

region = os.environ.get('AWS_REGION')
boto_config = BotoCoreConfig(read_timeout=70, region_name=region)
sf_client = boto3.client('stepfunctions', config=boto_config)
activity_arn = os.environ.get('SELECT_INTO_OUTFILE_ARN')

sm_client = boto3.client('secretsmanager', region_name=region)
mysql_password = sm_client.get_secret_value(SecretId=ENVIRONMENT+'/statementdb-to-snowflake/MYSQL_PASSWORD')[
    'SecretString'
]
MYSQL_CMD_PASSWORD = '-p'+str(mysql_password)

SQL = '''select * from accountingflat.dig_sales_detail LIMIT 500000
into outfile 'OUTFILE_DIROUTPUT_FILE_REPLACE/raw/outfile.txt' 
FIELDS ESCAPED BY '\\\\' 
TERMINATED BY '\\t' 
OPTIONALLY ENCLOSED BY '"' 
LINES TERMINATED BY '\\n';'''

while True:
    response = sf_client.get_activity_task(activityArn=activity_arn,
                                    workerName=ENVIRONMENT+'_select_into_outfile')

    if 'taskToken' not in response:
        print('No Task Token')
        time.sleep(5)
    else:
        print(response['taskToken']) 
        print("===================")
        activity_token = response['taskToken']
        activity_input = response['input']
        generate_uuid = activity_input.strip('\"')
        SQL = SQL.replace("OUTFILE_DIR", OUTFILE_DIR)
        SQL = SQL.replace("OUTPUT_FILE_REPLACE", str(generate_uuid))
        
        try:
            completed_process = subprocess.run(['mysql', MYSQL_CMD_USER, MYSQL_CMD_PASSWORD, '-e', SQL], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        except subprocess.CalledProcessError as exc:
            sf_client.send_task_failure(taskToken=activity_token, error=str(exc.stderr))
        else:            
            sf_client.send_task_success(taskToken=activity_token, output=json.dumps(str(generate_uuid)))