from pyspark import SparkConf
from pyspark import SparkContext
from pyspark import sql
from pyspark.sql.types import *

import config as conf_values
import json



def get_users( m):
    if  m['user_id'] in b_USERS.value:
        return m
    else:
        return ([])

def split_tracks(part):
    return (part[0].strip(), part[1].strip(),
            part[2], part[3].strip(),
            part[4].strip(), part[5].strip(), part[6],
            part[7], part[8], part[9], part[10], part[11], part[12],
            part[13].strip(), part[14], part[15].strip(),
            part[16], part[17].strip(), part[18])


if __name__ == "__main__":

    conf = SparkConf() \
        .setAppName(conf_values.spark_app) \
        .setMaster(conf_values.spark_master)

    sc = SparkContext(conf=conf)
    sc._jsc.hadoopConfiguration().set("fs.s3n.awsAccessKeyId", conf_values.aws_access_key)
    sc._jsc.hadoopConfiguration().set("fs.s3n.awsSecretAccessKey", conf_values.aws_secret_key)
    sql_context = sql.SQLContext(sc)

   # Broadcast variables

    b_USERS = sc.broadcast(['000000af49b341a65422ecfac2a98c92', '0000023b43b307354943ae74caff3886',
                '000004e97233ae3848ea7817d5529b9e', '00000527702c4e0e649985d7a8b99ade',
                '0000056d65e364bd13b5e56f73c284ec', '000005c0d34708aa2eecfb33bf7f8cb9',
                '0000063ed232d617a2a724b824e04172', '000006fa16ecb4db6a763d4f567a00a5'])

    track_data = sc.textFile("s3n://dev-cucumbers/spark/dim_track/slice*") \
                 .map(lambda line: line.split("\t")) \
                 .map(lambda parts: split_tracks(parts))

    track_fields = [StructField("track_id", LongType(), True),
          StructField("track_unique_id", LongType(), True),
          StructField("trackname", StringType(), True),
          StructField("upc", LongType(), True),
          StructField("cd", LongType(), True),
          StructField("track_id", LongType(), True),
          StructField("isrc", StringType(), True),
          StructField("p_line", StringType(), True),
          StructField("track_type", StringType(), True),
          StructField("version", StringType(), True),
          StructField("source", StringType(), True),
          StructField("last_updated", StringType(), True),
          StructField("date_created", StringType(), True),
          StructField("labelid", LongType(), True),
          StructField("vendor_track_identifier", StringType(), True),
          StructField("isrcid", LongType(), True),
          StructField("third_party_publisher", StringType(), True),
          StructField("resolution", StringType(), True),
          StructField("duration", IntegerType(), True)]
 track_schema = StructType(track_fields)

    # Apply the schema to the RDD.
    track_temp = sql_context.createDataFrame(track_data, track_schema)
    track_temp.registerTempTable("dim_track")


    # SQL can be run over DataFrames that have been registered as a table.
    results = sql_context.sql("SELECT isrc FROM dim_track limit 5")

    # The results of SQL queries are RDDs and support all the normal RDD operations.
    names = results.map(lambda p: "ISRC: " + p.isrc)
    for name in names.collect():
        print(name)

    stream_list = ['s3://dev-cucumbers/spark/streams_2016-07-11_NZ', 's3://dev-cucumbers/spark/streams_2016-07-11_PA']
    # Iterate because we want to get truncated files
    #file_list = sc.parallelize(key_list)
    #file_list.map(fetch_data).foreach(process_data)
    #stream_data = sc.textFile(
    for stream in stream_list:
        sc.textFile(stream) \
          .map(lambda (json_line): json.loads(json_line)) \
          .map(lambda m: (get_tracks(m))) \


    #rdd_streams = self.sc.textFile("s3n://cucumbers/SpotifyV2/archives/2016-07-11/streams_2016-07-11_MX.gz")
    user_data = sc.textFile("s3n://dev-cucumbers/spark/users_2016-07-11") \
                .map(lambda (json_line): json.loads(json_line)) \
                .map(lambda m: (get_users(m))) \
                .filter(lambda x: len(x) > 0) \
                .coalesce(1) \
                .saveAsTextFile("s3n://dev-cucumbers/spark/new_users_2016-07-11")
 #track_data = sc.textFile("s3n://dev-cucumbers/spark/tracks_2016-07-11")
    #track_parsed = 

        #filenames = rdd_files.map(lambda (filename, content): (filename))
        #filenames.collect()
        #schema_my_rdd = self.sql_context.createDataFrame(rdd_streams).collect()
        #schema_my_rdd.printSchema()
        #schema_my_rdd.registerTempTable("test_table")
        #df_pared = self.sql_context.sql("SELECT user.id_str as uid, id_str as tid, text as tweet, user.screen_name as uname, user.description as udesc from tweet WHERE lang = \"en\"" )