import configparser
from datetime import datetime
import os
from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.functions import year, month, dayofmonth, hour, weekofyear, date_format
import pandas as pd

# read config file for access keys
config = configparser.ConfigParser()
config.read('dl.cfg')

os.environ['AWS_ACCESS_KEY_ID']=config.get('AWS','AWS_ACCESS_KEY_ID')
os.environ['AWS_SECRET_ACCESS_KEY']=config.get('AWS','AWS_SECRET_ACCESS_KEY')


def create_spark_session():
    '''
    Desc: Create spark session
    '''
    spark = SparkSession \
        .builder \
        .config("spark.jars.packages", "org.apache.hadoop:hadoop-aws:2.7.0") \
        .getOrCreate()
    return spark


def process_song_data(spark, input_data, output_data):
    '''
    Desc: This function loads song data from S3 then extract required columns to load songs and artists
    table and then loads back to S3.
    
    Parameters:
    spark: Spark session varaiable
    input_data: the location of song_data from where the file is load to process
    output_data: the location where results will be stored
    
    '''
    # get filepath to song data file
    song_data = input_data + 'song_data/*/*/*/*.json'
    
    # read song data file
    df = spark.read.json(song_data)

    # extract columns to create songs table
    df.createOrReplaceTempView('song_data_table')
    songs_table = spark.sql(
    '''
    select song_id,
    title,
    artist_id,
    year,
    duration
    from song_data_table
    WHERE song_id IS NOT NULL
    '''
    )
    
    
    # write songs table to parquet files partitioned by year and artist
    songs_table.write.mode('overwrite').partitionBy("year", "artist_id").option("encoding", "UTF8").parquet(output_data+'songs_table/')

    # extract columns to create artists table
    artists_table = spark.sql(
    '''
    select artist_id,
    artist_name,
    artist_location,
    artist_latitude,
    artist_longitude
    from song_data_table
    WHERE artist_id IS NOT NULL
    '''
    )
    
    # write artists table to parquet files
    artists_table.write.mode('overwrite').parquet(output_data+'artist_table/')


def process_log_data(spark, input_data, output_data):
    '''
    Desc: This function loads log data from S3 then extract required columns to load user ,time and songplay
    table and then loads back to S3.
    
    Parameters:
    spark: Spark session varaiable
    input_data: the location of log_data from where the file is load to process
    output_data: the location where results will be stored
    
    '''
    # get filepath to log data file
    log_data = input_data + 'log_data/*.json'

    # read log data file
    df = spark.read.json(log_data)
    
    # filter by actions for song plays
    df = df.filter(df.page == 'NextSong')

    # extract columns for users table
    df.createOrReplaceTempView('log_data_table')
    users_table = spark.sql(
    '''
    select userId,
    firstName,
    lastName,
    gender,
    level
    from log_data_table
    WHERE userId IS NOT NULL
    '''
    ) 
    
    # write users table to parquet files
    users_table.write.mode('overwrite').parquet(output_data+'users_table/')

    
    # extract columns to create time table

    time_table = spark.sql('''
                            SELECT 
                            t.start_time_ts as start_time,
                            hour(t.start_time_ts) as hour,
                            dayofmonth(t.start_time_ts) as day,
                            weekofyear(t.start_time_ts) as week,
                            month(t.start_time_ts) as month,
                            year(t.start_time_ts) as year,
                            dayofweek(t.start_time_ts) as weekday
                            FROM
                            (SELECT to_timestamp(timeSt.ts/1000) as start_time_ts
                            FROM log_data_table timeSt
                            WHERE timeSt.ts IS NOT NULL
                            ) t
                        ''')
    
    # write time table to parquet files partitioned by year and month
    time_table.write.mode('overwrite').partitionBy("year","month").parquet(output_data+'time_table/')

    # read in song data to use for songplays table
    song_df = spark.read.parquet(output_data+'songs_table/')

    # extract columns from joined song and log datasets to create songplays table 
    songplays_table = spark.sql('''
                                SELECT monotonically_increasing_id() as songplay_id,
                                to_timestamp(logT.ts/1000) as start_time,
                                month(to_timestamp(logT.ts/1000)) as month,
                                year(to_timestamp(logT.ts/1000)) as year,
                                logT.userId as user_id,
                                logT.level as level,
                                songT.song_id as song_id,
                                songT.artist_id as artist_id,
                                logT.sessionId as session_id,
                                logT.location as location,
                                logT.userAgent as user_agent
                                FROM log_data_table logT
                                JOIN song_data_table songT 
                                on logT.artist = songT.artist_name and logT.song = songT.title
                                ''')

    # write songplays table to parquet files partitioned by year and month
    songplays_table.write.mode('overwrite').partitionBy("year","month").parquet(output_data+'songplays_table/')


def main():
    spark = create_spark_session()
    input_data = "s3a://udacity-dend/"
    output_data = "s3a://udacity-dend-proj/project-4-full/"
    
    process_song_data(spark, input_data, output_data)    
    process_log_data(spark, input_data, output_data)


if __name__ == "__main__":
    main()
