import os
import argparse
from datetime import datetime, timedelta

from snowflake.snowpark.session import Session
from dotenv import load_dotenv

# Initialize parser
parser = argparse.ArgumentParser(description="A script that accepts arguments.")

# Add arguments
parser.add_argument("--file", required=False, help="The file to be deployed")
parser.add_argument("--branch", required=False, default="main", help="The branch to be used")
parser.add_argument("--procedure", required=False, help="The procedure to be called")
parser.add_argument("--start_date", required=False, help="The start date to be processed")
parser.add_argument("--start_hour", required=False, help="The start hour to be processed")
parser.add_argument("--end_date", required=False, help="The end date to be processed")
parser.add_argument("--end_hour", required=False, help="The end hour to be processed")
parser.add_argument("--warehouse", required=False, default="SUNO_PROD_ENGINEER_X_SMALL", help="The warehouse to be used")

args = parser.parse_args()
deploy_file = args.file is not None
call_procedure = args.procedure is not None
procedure_no_date = args.start_date is None and args.end_date is None
procedure_single_date = args.start_date is not None and args.end_date is None
procedure_date_range = args.start_date is not None and args.end_date is not None
procedure_no_hour = args.start_hour is None and args.end_hour is None
procedure_single_hour = args.start_hour is not None and args.end_hour is None
procedure_hour_range = args.start_hour is not None and args.end_hour is not None

# Load environment variables
load_dotenv()
SNOWFLAKE_CONFIGS = {
    "account": os.getenv("SNOWFLAKE_ACCOUNT"),
    "user": os.getenv("SNOWFLAKE_ACCOUNT_USER"),
    "private_key_file": os.getenv("SNOWFLAKE_PRIVATE_KEY_FILE"),
    "role": os.getenv("SNOWFLAKE_ACCOUNT_ROLE"),
}

# Get Snowflake session
def get_snowflake_session(database: str, warehouse: str, schema: str) -> Session:
    return Session.builder.configs(
        {"warehouse": warehouse, "database": database, "schema": schema, **SNOWFLAKE_CONFIGS}
    ).create()

def run_procedure(session: Session, sql: str):
    print(sql)
    print()

    exec_start_time = datetime.now()
    print(f'Starting ({exec_start_time.strftime("%Y-%m-%d %H:%M:%S")})')
    print(session.sql(sql).collect())

    exec_end_time = datetime.now()
    print(f'Finished ({exec_end_time.strftime("%Y-%m-%d %H:%M:%S")})')
    print(f'Elapsed time: {(exec_end_time - exec_start_time).total_seconds()} seconds')
    print('----------------------------------------------')

# cd to the same directory as this file and run the following script to deploy the procedure
# make sure you push the code to main branch before deploy
# uv run deploy_to_prod.py --file <file path>
if __name__ == "__main__":
    session = get_snowflake_session(database="SUNO_PROD", warehouse=args.warehouse, schema="PROD")
    print(f'Using warehouse {args.warehouse}...')
    print('----------------------------------------------')
    try:
        # fetch the latest code from the git repository
        sql = '''
            ALTER GIT REPOSITORY snowflake FETCH;
        '''
        print (session.sql(sql).collect())
        print('==============================================')

        if deploy_file:
            print(f'Deploying file {args.file} from branch origin/{args.branch} to prod...')
            # deploy specified sql file from specified branch (default main) to prod
            sql = f'''
                EXECUTE IMMEDIATE FROM @suno_etl/branches/{args.branch}/snowflake/code/tables/{args.file};
            '''
            print (session.sql(sql).collect())
        else:
            print('No file to deploy, skipping...')
        print('==============================================')

        if call_procedure:
            # Procedures that do not require a date input
            if procedure_no_date:
                print('No start/end provided, calling procedure with no arguments...')
                print('----------------------------------------------')
                sql = f"CALL {args.procedure}();"
                run_procedure(session, sql)

            # Run procedure for a single date/datetime
            elif procedure_single_date:
                # If hour is provided, call procedure with date and hour
                if procedure_single_hour:
                    print(f'Calling procedure with date={args.start_date} and hour={args.start_hour}')
                    print('----------------------------------------------')
                    sql = f"CALL {args.procedure}('{args.start_date}', '{args.start_hour}');"
                    run_procedure(session, sql)
                # If no hour is provided, call procedure with date
                else:
                    print(f'Calling procedure with date={args.start_date} and no hour')
                    print('----------------------------------------------')
                    sql = f"CALL {args.procedure}('{args.start_date}');"
                    run_procedure(session, sql)

            # Run procedure over a date range/datetime range
            elif procedure_date_range:
                # If hour range is provided, increment by hour
                if procedure_hour_range:
                    start_datetime = datetime.strptime(f"{args.start_date} {args.start_hour}", '%Y-%m-%d %H')
                    end_datetime = datetime.strptime(f"{args.end_date} {args.end_hour}", '%Y-%m-%d %H')
                    assert start_datetime <= end_datetime, "Start datetime must be before end datetime, but got start_datetime={start_datetime} and end_datetime={end_datetime}"

                    print(f'Calling procedure from {start_datetime.strftime("%Y-%m-%d %H:%M:%S")} to {end_datetime.strftime("%Y-%m-%d %H:%M:%S")}, incrementing hourly')
                    print('----------------------------------------------')
                    current_datetime = start_datetime
                    while current_datetime <= end_datetime:
                        sql = f"CALL {args.procedure}('{current_datetime.strftime('%Y-%m-%d')}', '{current_datetime.strftime('%H')}');"
                        run_procedure(session, sql)
                        current_datetime += timedelta(hours=1)
                # If no hour range is provided, increment by day
                else:
                    start_date = datetime.strptime(args.start_date, '%Y-%m-%d')
                    end_date = datetime.strptime(args.end_date, '%Y-%m-%d')
                    assert start_date <= end_date, "Start date must be before end date, but got start_date={start_date} and end_date={end_date}"

                    print(f'Calling procedure from {start_date.strftime("%Y-%m-%d")} to {end_date.strftime("%Y-%m-%d")}, incrementing daily')
                    print('----------------------------------------------')
                    current_date = start_date
                    while current_date <= end_date:
                        sql = f"CALL {args.procedure}('{current_date.strftime('%Y-%m-%d')}');"
                        run_procedure(session, sql)
                        current_date += timedelta(days=1)

        else:
            print('No procedure to call, skipping...')
        print('==============================================')


    except Exception as e:
        print(e)
    finally:
        session.close()