import argparse
import os
import sys
from datetime import datetime, timedelta

from cryptography.hazmat.primitives import serialization
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()


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__":
    # Load and format the private key
    raw_private_key = os.getenv("SNOWFLAKE_PRIVATE_KEY")

    if not raw_private_key:
        raise ValueError("Environment variable SNOWFLAKE_PRIVATE_KEY is not set or empty")

    # ✅ Ensure the private key is properly formatted
    raw_private_key = raw_private_key.strip().encode()  # Convert string to bytes

    # ✅ Load the private key into the correct format
    private_key = serialization.load_pem_private_key(
        raw_private_key,
        password=None,  # If encrypted, replace with the passphrase
    )

    # ✅ Convert the private key into DER format for Snowflake
    private_key_der = private_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
    CONNECTION_PARAMETERS = {
        "account": os.getenv("SNOWFLAKE_ACCOUNT"),
        "user": os.getenv("SNOWFLAKE_ACCOUNT_USER"),
        "private_key": private_key_der,
        "role": os.getenv("SNOWFLAKE_ACCOUNT_ROLE"),
        "database": "SUNO_PROD",
        "warehouse": args.warehouse,
        "schema": "PROD",
    }

    session = Session.builder.configs(CONNECTION_PARAMETERS).create()

    print(f"Using warehouse {args.warehouse}...")
    print("----------------------------------------------")

    error_message = None

    try:
        print("Fetching latest code from Github into SUNO_ETL...")
        session.sql("alter git repository SUNO_ETL fetch").collect()

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

        # formats the error message to be displayed in Github Actions output
        error_message_lines = str(e).split("\n")
        if len(error_message_lines) > 1 and "SQL compilation error" in error_message_lines[0]:
            del error_message_lines[0]

        error_message = " ".join(error_message_lines)

        # this particular error does not need to be indicated to the user
        if "SQL compilation error" in str(e) and "already exists" in error_message:
            error_message = None

    finally:
        session.close()

    # fails with the error message so it can be printed in Github Actions output
    if error_message:
        sys.exit(error_message)