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("--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("--increment", required=False, help="The increment to be used", choices=["hour", "day", "week", "month"])
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

# 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_local_changes.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:
        if deploy_file:
            print(f'Deploying local file {args.file} to prod...')
            with open(args.file, "r") as file:
                sql_script = file.read()
                print(sql_script)
                print(session.sql(sql_script).collect())
        else:
            print('No file to deploy, skipping...')
        print('==============================================')

        if call_procedure:
            if args.increment == "hour":
                increment = timedelta(hours=0)
            elif args.increment == "day":
                increment = timedelta(days=1) - timedelta(hours=1)
            elif args.increment == "week":
                increment = timedelta(weeks=1) - timedelta(hours=1)
            elif args.increment == "month":
                increment = timedelta(days=31) - timedelta(hours=1)
            else:
                raise ValueError(f"Invalid increment: {args.increment}; options are [hour, day, week, month]")

            # All arguments are required (start/end date and hour)
            if not all([args.start_date, args.end_date, args.start_hour, args.end_hour]):
                raise ValueError("All arguments are required (start/end date and hour)")

            start_ts = datetime.strptime(args.start_date + ' ' + args.start_hour, '%Y-%m-%d %H')
            end_ts = datetime.strptime(args.end_date + ' ' + args.end_hour, '%Y-%m-%d %H')
            assert start_ts <= end_ts, f"Start timestamp must be equal to or before end timestamp, but got start = {start_ts} and end = {end_ts}"

            # Call procedure, incrementing by the specified increment
            current_start_ts = start_ts
            while current_start_ts <= end_ts:
                current_end_ts = current_start_ts + increment
                current_start_date = current_start_ts.strftime("%Y-%m-%d")
                current_start_hour = current_start_ts.strftime("%H")
                current_end_date = current_end_ts.strftime("%Y-%m-%d")
                current_end_hour = current_end_ts.strftime("%H")

                print(f'Calling procedure from {current_start_ts} to {current_end_ts}, incrementing by {args.increment}')
                print('----------------------------------------------')
                sql = f"CALL {args.procedure}('{current_start_date}', '{current_start_hour}', '{current_end_date}', '{current_end_hour}');"
                run_procedure(session, sql)
                current_start_ts = current_end_ts + timedelta(hours=1)

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

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