import sys
from pyspark.sql import SparkSession
from pyspark.sql import DataFrame
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from awsglue.context import GlueContext
from awsglue.dynamicframe import DynamicFrame
from awsglue.job import Job
from pyspark.sql.functions import expr, col, base64
import botocore
import boto3
from pyspark.sql import functions as F
from pyspark.sql.types import DateType, StringType, IntegerType
from datetime import datetime
from typing import Dict, List, Union, Tuple
import numpy as np
import yaml
from pyathena import connect as pyathena_connect

# spark session and other variables
spark = SparkSession.builder.appName('<app_name>').getOrCreate()
sc = spark.sparkContext
glueContext = GlueContext(sc)
job = Job(glueContext)
glue_client = boto3.client('glue')
s3_client = boto3.client('s3')
athena_cursor = pyathena_connect(
s3_staging_dir=athena_output_path, work_group=work_group).cursor()

# get config file path from job parameters
args = getResolvedOptions(sys.argv, ['config_file_path'])
config_file_path = args['config_file_path']
config_file_bucket = config_file_path.split('/')[2]
config_file_key = '/'.join(config_file_path.split('/')[3:])

def create_df_from_catalog(database: str, table: str, year: str, month: str, day: str) -> DataFrame:
    """
    Creates a Spark DataFrame from a Glue catalog using the specified parameters.

    Parameters:
    database (str): The name of the Glue catalog database.
    table (str): The name of the table in the Glue catalog.
    year (str): The year value as a string.
    month (str): The month value as a string.
    day (str): The day value as a string.

    Returns:
    DataFrame: A Spark DataFrame created from the Glue catalog.
    """
    df = glueContext.create_data_frame_from_catalog(
        database=database,
        table_name=table,
        push_down_predicate=f"(year=={year} and month=={month} and day=={day})"
    )
    return df

#DON'T NEED IT
def create_latest_partition_df_from_catalog(database: str, table: str) -> DataFrame:
    """
    Creates a Spark DataFrame from the latest partition in the Glue catalog for the specified database and table.

    Parameters:
    database (str): The name of the Glue catalog database.
    table (str): The name of the table in the Glue catalog.

    Returns:
    DataFrame: A Spark DataFrame created from the latest partition in the Glue catalog.
    """
    partitions = get_partitions(
        database=database,
        table=table
    )
    latest_partition = max(partitions)
    latest_year, latest_month, latest_day = latest_partition

    df = glueContext.create_data_frame_from_catalog(
        database=database,
        table_name=table,
        push_down_predicate=f"(year=={latest_year} and month=={latest_month} and day=={latest_day})"
    )
    return df

def create_all_partition_df_from_catalog(database: str, table: str) -> DataFrame:
    """
    Creates a Spark DataFrame from all partitions in the Glue catalog for the specified database and table.

    Parameters:
    database (str): The name of the Glue catalog database.
    table (str): The name of the table in the Glue catalog.

    Returns:
    DataFrame: A Spark DataFrame created from the latest partition in the Glue catalog.
    """
    df = glueContext.create_data_frame_from_catalog(
        database=database,
        table_name=table,
    )
    return df

def get_partitions(database: str, table: str) -> List[List[str]]:
    """
    Retrieves and returns a sorted list of partitions for the specified database and table.

    Parameters:
    database (str): The name of the Glue catalog database.
    table (str): The name of the table in the Glue catalog.

    Returns:
    List[List[str]]: A sorted list of partitions, where each partition is represented as a list of strings.
    """
    partitions = glue_client.get_partitions(
        DatabaseName=database,
        TableName=table
    )
    partitions_list = partitions['Partitions']
    partitions_sorted = sorted(
        list(map(lambda x: x['Values'], partitions_list))
    )
    return partitions_sorted

def remaining_dates_to_sync(src_database: str, dest_database: str, src_table: str, dest_table: str) -> List[List[str]]:
    """
    Determines the remaining dates to synchronize between the source and destination tables in the Glue catalog.

    Parameters:
    src_database (str): The name of the source Glue catalog database.
    dest_database (str): The name of the destination Glue catalog database.
    src_table (str): The name of the source table in the Glue catalog.
    dest_table (str): The name of the destination table in the Glue catalog.

    Returns:
    List[List[str]]: A list of remaining dates to synchronize, where each date is represented as a list of strings.
    If there are no remaining dates, a list containing the latest destination partition is returned.
    If the destination database or table does not exist, the function returns all source partitions.
    """
    src_partitions = get_partitions(
        database=src_database,
        table=src_table
    )
    try:
        # Try to get destination partitions
        dest_partitions = get_partitions(
            database=dest_database,
            table=dest_table
        )
    except:
        # If destination database and table do not exist, return all source partitions
        return src_partitions

    latest_dest_partition = dest_partitions[-1]
    remaining_dates = [date for date in src_partitions if date > latest_dest_partition]
    if not remaining_dates:
        return [latest_dest_partition]
    return remaining_dates

def get_config_file(bucket: str, key: str) -> dict:
    """
    Retrieves and returns the contents of a YAML config file stored in an S3 bucket.

    Parameters:
    bucket (str): The name of the S3 bucket.
    key (str): The key or path of the config file in the S3 bucket.

    Returns:
    dict: The contents of the YAML config file, represented as a dictionary.

    Raises:
    botocore.exceptions.ClientError: If there is an error reading the config file from S3.
    """
    try:
        response = s3_client.get_object(Bucket=bucket, Key=key)
        contents = response['Body'].read().decode('utf-8')
        yaml_data = yaml.safe_load(contents)
        return yaml_data
    except botocore.exceptions.ClientError as e:
        print(
            f"Could not read config file s3://{bucket}/{key}. Exception: {e}.")
        raise

# bronze_to_silver_glue_job_config = {
#     ""source_catalog_database_name"": config['business']['ops']['glue_job_config']['bronze_to_silver']['source_catalog_database'],
#     ""destination_catalog_database_name"": config['business']['ops']['glue_job_config']['bronze_to_silver']['destination_catalog_database'],
#     ""source_test_catalog_table_name"": config['business']['ops']['glue_job_config']['bronze_to_silver']['source_test_catalog_table'],
#     ""destination_test_catalog_table_name"": config['business']['ops']['glue_job_config']['bronze_to_silver']['destination_test_catalog_table'],
#     ""destination_s3_path"" : config['business']['ops']['glue_job_config']['bronze_to_silver']['destination_s3_path']
# }

def create_df_with_partition(df: DataFrame, year: str, month: str, day: str) -> DataFrame:
    """
    Creates a new Spark DataFrame with additional partition columns.

    Parameters:
    df (DataFrame): The input Spark DataFrame.
    year (str): The year value for the partition column.
    month (str): The month value for the partition column.
    day (str): The day value for the partition column.

    Returns:
    DataFrame: A new Spark DataFrame with the additional partition columns.
    """
    return df.withColumn("year", F.lit(year)) \
        .withColumn("month", F.lit(month)) \
        .withColumn("day", F.lit(day))

def write_df_to_s3(df, dest_bucket: str, dest_prefix: str, dest_database: str, dest_table: str):
    """
    Writes a DataFrame to S3 and updates the data catalog with partition information.

    Parameters:
    df (DataFrame): The input DataFrame to write.
    dest_bucket (str): The destination S3 bucket.
    dest_prefix (str): The prefix within the S3 bucket where the DataFrame will be written.
    dest_database (str): The destination Glue catalog database.
    dest_table (str): The destination Glue catalog table.

    Returns:
    None
    """
    # Get the latest source partition
    lastest_src_partition = df.select('year', 'month', 'day').first().asDict()
    year = lastest_src_partition["year"]
    month = lastest_src_partition["month"]
    day = lastest_src_partition["day"]

    # Destination table path to write to
    dest_write_path = f's3://{dest_bucket}/{dest_prefix}{dest_table}'

    # Create DynamicFrame from the DataFrame
    dynamic_df = DynamicFrame.fromDF(df.coalesce(1), glueContext, dest_table)

    # Delete existing folder for current partition
    glueContext.purge_s3_path(
        s3_path=f'{dest_write_path}/year={year}/month={month}/day={day}',
        options={"retentionPeriod": 0}
    )

    sink = glueContext.getSink(
        connection_type="s3",
        updateBehavior="UPDATE_IN_DATABASE",
        path=dest_write_path,
        enableUpdateCatalog=True,
        partitionKeys=["year", "month", "day"]
    )
    sink.setFormat("glueparquet")
    sink.setCatalogInfo(
        catalogDatabase=dest_database,
        catalogTableName=dest_table
    )

    # Write DynamicFrame to S3 and data catalog
    sink.writeFrame(dynamic_df)