"""Read Kafka Streams into Spark, perform simple filtering/aggregation"""

import argparse
import sys
from time import sleep

try:
    from pyspark.sql import SparkSession
    from pyspark.sql.functions import col, from_json, udf
    from pyspark.sql.types import BooleanType, IntegerType, StringType, StructType
except ImportError:
    print("\npip install pyspark")
    sys.exit(1)

try:
    import tldextract
except ImportError:
    print("\nThis example needs tldextract. Please do a $pip install tldextract and rerun this example")
    sys.exit(1)


def exit_program():
    """Exit on Signal"""
    print("Exiting Program...")
    sys.exit()


def compute_domain(query):
    # Pull out the domain
    if query.endswith(".local"):
        return "local"
    return tldextract.extract(query).registered_domain if query else None


if __name__ == "__main__":
    """Read Kafka Streams into Spark, perform simple filtering/aggregation"""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--server", type=str, default="localhost:9092", help="Specify the Kafka Server (default: localhost:9092)"
    )
    args, commands = parser.parse_known_args()

    # Check for unknown args
    if commands:
        print("Unrecognized args: %s" % commands)
        sys.exit(1)

    # Grab the Kafka server
    kserver = args.server

    # Spin up a local Spark Session (with 4 executors)
    spark = (
        SparkSession.builder.master("local[4]")
        .appName("my_awesome")
        .config("spark.jars.packages", "org.apache.spark:spark-sql-kafka-0-10_2.11:2.4.4")
        .getOrCreate()
    )
    spark.sparkContext.setLogLevel("ERROR")

    # Optimize the conversion to Spark
    spark.conf.set("spark.sql.execution.arrow.enable", "true")

    # SUBSCRIBE: Setup connection to Kafka Stream
    raw_data = (
        spark.readStream.format("kafka")
        .option("kafka.bootstrap.servers", kserver)
        .option("subscribe", "dns")
        .option("startingOffsets", "earliest")
        .load()
    )

    # Define the schema for the DNS message (do this better)
    dns_schema = (
        StructType()
        .add("ts", StringType())
        .add("uid", StringType())
        .add("id.orig_h", StringType())
        .add("id.orig_p", IntegerType())
        .add("id.resp_h", StringType())
        .add("id.resp_p", IntegerType())
        .add("proto", StringType())
        .add("trans_id", IntegerType())
        .add("query", StringType())
        .add("qclass", IntegerType())
        .add("qclass_name", StringType())
        .add("qtype", IntegerType())
        .add("qtype_name", StringType())
        .add("rcode", IntegerType())
        .add("rcode_name", StringType())
        .add("AA", BooleanType())
        .add("TC", BooleanType())
        .add("RD", BooleanType())
        .add("RA", BooleanType())
        .add("Z", IntegerType())
        .add("answers", StringType())
        .add("TTLs", StringType())
        .add("rejected", BooleanType())
    )

    # ETL: Convert raw data into parsed and proper typed data
    parsed_data = raw_data.select(from_json(col("value").cast("string"), dns_schema).alias("data")).select("data.*")

    # FILTER: Only get DNS records that have 'query' field filled out
    filtered_data = parsed_data.filter(parsed_data.query.isNotNull() & (parsed_data.query != ""))

    # FILTER 2: Remove Local/mDNS queries
    filtered_data = filtered_data.filter(~filtered_data.query.like("%.local"))  # Note: using the '~' negation operator

    # COMPUTE: A new column with the 2nd level domain extracted from the query
    udf_compute_domain = udf(compute_domain, StringType())
    computed_data = filtered_data.withColumn("domain", udf_compute_domain("query"))

    # AGGREGATE: In this case a simple groupby operation
    group_data = computed_data.groupBy("`id.orig_h`", "domain", "qtype_name").count()

    # At any point in the pipeline you can see what you're getting out
    group_data.printSchema()

    # Take the end of our pipeline and pull it into memory
    dns_count_memory_table = (
        group_data.writeStream.format("memory").queryName("dns_counts").outputMode("complete").start()
    )

    # Let the pipeline pull some data
    print("Pulling pipline...Please wait...")

    # Create a Pandas Dataframe by querying the in memory table and converting
    # Loop around every 5 seconds to update output
    for _ in range(10):
        sleep(5)
        dns_counts_df = spark.sql("select * from dns_counts").toPandas()
        print("\nDNS Query Total Counts = {:d}".format(dns_counts_df["count"].sum()))
        print(dns_counts_df.sort_values(ascending=False, by="count"))

    # Stop the stream
    dns_count_memory_table.stop()
    sleep(1)
