import sys
import abc
from datetime import datetime, timedelta
import time

from pyspark.context import SparkContext  # type: ignore
from awsglue.utils import getResolvedOptions  # type: ignore
from awsglue.context import GlueContext  # type: ignore
from awsglue.job import Job  # type: ignore
from utils.snowflake.snowflake_client import get_snowflake_connection  # type: ignore

import boto3  # type: ignore

DYNAMODB_DEFAULT_EXPORT_TIMEOUT = 600


class DynamoDBtoS3Job(abc.ABC):
    """
    Abstract base class representing an AWS Glue Job that gets data from DynamoDB,
    exports it into an S3 bucket, and then calls a procedure to upsert it into Snowflake.
    """

    def __init__(
        self,
        dynamodb_table_name: str,
        procedure_name: str,
        task_monitor_name: str,
        region_name: str = "us-east-2",
        dynamodb_export_timeout: int = DYNAMODB_DEFAULT_EXPORT_TIMEOUT,
        prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL",
        staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL",
        print_debug_logs: bool = False,
        execute_snowflake_procedure: bool = True,
    ):
        self.dynamodb_table_name = dynamodb_table_name
        self.procedure_name = procedure_name
        self.task_monitor_name = task_monitor_name
        self.print_debug_logs = print_debug_logs
        self.prod_warehouse_name = prod_warehouse
        self.staging_warehouse_name = staging_warehouse
        (
            self.account_id,
            self.environment,
            self.s3_bucket,
            self.warehouse,
            self.db_name,
        ) = self._get_running_info()
        self.region = region_name
        self.dynamodb_export_timeout = dynamodb_export_timeout
        self.start_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        self.execute_snowflake_procedure = execute_snowflake_procedure
        
    def debug_print(self, message: str):
        if self.print_debug_logs:
            print(message)

    def _get_running_info(self):
        """
        Gets important information about the job's running environment, depending on whether
        the job is run in staging or production. Returns the environment (staging or prod) and
        the S3 bucket, Snowflake warehouse, and Snowflake database to write to.
        """
        PROD_ACCOUNT_ID = "734185074900"
        STAGING_ACCOUNT_ID = "590183763515"

        # uses the account id to determine whether we're in prod or staging
        sts = boto3.client("sts")
        account_id = sts.get_caller_identity()["Account"]
        if account_id == PROD_ACCOUNT_ID:
            environment = "PROD"
            s3_bucket = "analytics-database-data"
            warehouse = self.prod_warehouse_name

        elif account_id == STAGING_ACCOUNT_ID:
            environment = "STAGING"
            s3_bucket = "analytics-database-data-staging"
            warehouse = self.staging_warehouse_name

        else:
            raise Exception("Invalid account ID.")

        # sets the Snowflake database to write to depending on whether we're in prod or staging
        db_name = f"SUNO_{environment}"

        return account_id, environment, s3_bucket, warehouse, db_name

    def _initialize_glue_context(self):
        """
        Initializes some context necessary for the Glue Job to be run.
        """
        args = getResolvedOptions(sys.argv, ["JOB_NAME"])
        sc = SparkContext()
        glueContext = GlueContext(sc)
        spark = glueContext.spark_session
        job = Job(glueContext)
        job.init(args["JOB_NAME"], args)
        return glueContext, spark, job

    def get_call_snowflake_procedure_sql(
        self,
        procedure_name: str,
        params: list,
    ):
        """
        Returns a SQL query that calls a Snowflake procedure with the given parameters.
        """
        return f"""call {procedure_name}({", ".join([f"'{param}'" for param in params])})"""

    @abc.abstractmethod
    def get_s3_save_location_path(self) -> str:
        """
        Returns the path in S3 where the data from DynamoDB will be saved.
        """
        pass

    @abc.abstractmethod
    def get_dynamodb_export_time_interval(self) -> tuple[datetime, datetime]:
        """
        Returns the "from" and "to" times specifying a range of time, such that data
        updated or created within that time range will be exported from DynamoDB.
        """
        pass

    @abc.abstractmethod
    def get_snowflake_procedure_sql_parameters(self) -> list[str]:
        """
        Returns the parameters that should be passed to the Snowflake procedure that
        performs the upsert.
        """
        pass

    @abc.abstractmethod
    def get_job_frequency(self) -> str:
        """
        Returns the frequency with which this job is called, e.g. "DAILY" or "HOURLY".
        """
        pass

    def _start_incremental_dynamodb_export_to_s3(self, client):
        """
        Triggers an incremental export of the DynamoDB table to S3 using the given boto3
        client.
        """
        from_time, to_time = self.get_dynamodb_export_time_interval()
        export_description = client.export_table_to_point_in_time(
            TableArn=f"arn:aws:dynamodb:{self.region}:{self.account_id}:table/{self.dynamodb_table_name}",
            S3Bucket=f"{self.s3_bucket}",
            S3Prefix=f"{self.get_s3_save_location_path()}",
            ExportType="INCREMENTAL_EXPORT",
            IncrementalExportSpecification={
                "ExportFromTime": from_time,
                "ExportToTime": to_time,
            },
        )["ExportDescription"]

        return export_description["ExportArn"], export_description["ExportStatus"]

    def _poll_for_export_completion(
        self,
        client,
        export_arn,
        export_status,
        timeout=DYNAMODB_DEFAULT_EXPORT_TIMEOUT,
        frequency=30,
    ):
        """
        Checks repeatedly until the DynamoDB to S3 export is complete using the given
        boto3 client. Throws an exception if the export takes longer than the timeout
        period, or if the export terminates with any status other than "COMPLETED".
        """
        timeout_period = timedelta(seconds=timeout)
        export_start_time = datetime.now()
        timed_out = False

        while export_status == "IN_PROGRESS" and not timed_out:
            time.sleep(frequency)

            self.debug_print("Checking export status...")

            export_description = client.describe_export(ExportArn=export_arn)[
                "ExportDescription"
            ]
            export_status = export_description["ExportStatus"]

            self.debug_print(f"Export status: {export_status}")

            if datetime.now() - export_start_time > timeout_period:
                timed_out = True
                break

        if timed_out:
            raise Exception(f"Export took longer than {timeout} seconds and timed out.")

        elif export_status == "COMPLETED":
            item_count = export_description["ItemCount"]
            self.debug_print(f"Exported table to S3. Item count: {item_count}")

        else:
            raise Exception(f"Export failed to complete. Status: {export_status}.")
        return item_count

    def run_job(self):
        """
        Runs this Glue Job. Gets data from the DynamoDB table with the table name
        associated with this job, exports it to S3 using the boto3 client, and then
        calls a Snowflake procedure to upsert it. Does not handle any errors.
        """
        # Glue parameters
        _, _, job = self._initialize_glue_context()

        self.debug_print(f"Environment: {self.environment}")
        self.debug_print(f"S3 Bucket: {self.s3_bucket}")
        self.debug_print(f"Warehouse: {self.warehouse}")
        self.debug_print(f"DB Name: {self.db_name}")

        client = boto3.client("dynamodb")

        self.debug_print("Exporting table to S3...")

        export_arn, export_status = self._start_incremental_dynamodb_export_to_s3(
            client
        )

        item_count = self._poll_for_export_completion(
            client, export_arn, export_status, timeout=self.dynamodb_export_timeout
        )

        if self.execute_snowflake_procedure:
            task_monitor_sql_params = [
                self.task_monitor_name,
                self.start_time,
                item_count,
                "SUCCESS",
                self.get_job_frequency().upper(),
            ]

            snowflake_procedure_sql = self.get_call_snowflake_procedure_sql(
                self.procedure_name, self.get_snowflake_procedure_sql_parameters()
            )

            self.debug_print(f"Calling SQL: {snowflake_procedure_sql}")
            task_monitor_sql = self.get_call_snowflake_procedure_sql(
                "TASK_MONITOR_INSERT_PROC", task_monitor_sql_params
            )

            connection = get_snowflake_connection(
                self.warehouse, self.db_name, self.environment
            )
            cursor = connection.cursor()

            self.debug_print(f"Calling SQL: {snowflake_procedure_sql}")
            result = cursor.execute(snowflake_procedure_sql)
            self.debug_print(f"Calling SQL: {task_monitor_sql}")
            result = cursor.execute(task_monitor_sql)

        self.debug_print("Query executed successfully.")

        job.commit()


class HourlyDynamoDBtoS3Job(DynamoDBtoS3Job):
    """
    A Glue Job that exports data from DynamoDB to S3 hourly.
    """

    def __init__(
        self,
        dynamodb_table_name: str,
        procedure_name: str,
        task_monitor_name: str,
        region_name: str = "us-east-2",
        dynamodb_export_timeout: int = DYNAMODB_DEFAULT_EXPORT_TIMEOUT,
        prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL",
        staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL",
        print_debug_logs: bool = False,
        execute_snowflake_procedure: bool = True,
    ):
        DynamoDBtoS3Job.__init__(
            self,
            dynamodb_table_name,
            procedure_name,
            task_monitor_name,
            region_name=region_name,
            dynamodb_export_timeout=dynamodb_export_timeout,
            prod_warehouse=prod_warehouse,
            staging_warehouse=staging_warehouse,
            print_debug_logs=print_debug_logs,
            execute_snowflake_procedure=execute_snowflake_procedure,
        )

        current_date = datetime.now()
        one_hour_ago = current_date - timedelta(hours=1)

        self.start_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        self.hour_start = one_hour_ago.replace(minute=0, second=0, microsecond=0)
        self.hour_end = self.hour_start + timedelta(hours=1)
        self.date = one_hour_ago.strftime("%Y-%m-%d")
        self.hour = one_hour_ago.strftime("%H")

    def get_s3_save_location_path(self) -> str:
        return f"{self.dynamodb_table_name}/pdate={self.date}/phour={self.hour}"

    def get_dynamodb_export_time_interval(self) -> tuple[datetime, datetime]:
        return self.hour_start, self.hour_end

    def get_snowflake_procedure_sql_parameters(self) -> list[str]:
        stage_path = (
            f"@SUNO_DYNAMODB_EVENTS/{self.get_s3_save_location_path()}/AWSDynamoDB/data"
        )
        return [self.date, self.hour, stage_path]

    def get_job_frequency(self) -> str:
        return "HOURLY"