import argparse
import os
from datetime import datetime, timedelta

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

# 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(
    "--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_local_changes.py --file <file path>
if __name__ == "__main__":
    session = get_snowflake_session(database="SUNO_STAGING", warehouse=args.warehouse, schema="STAGING")
    print(f"Using warehouse {args.warehouse}...")
    print("----------------------------------------------")

    try:
        if deploy_file:
            print(f"Deploying local file {args.file} to staging...")
            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:
            # 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()