import sys
import re
from pyspark import SparkContext
from pyspark.streaming import StreamingContext
from pyspark.streaming.kafka import KafkaUtils


def main():
    if len(sys.argv) != 3:
        print("Usage: recoverable_network_wordcount.py <brokers> <topics>")
        exit(-1)
    brokers, topics = sys.argv[1:]

    # Create context with 1 second batch interval
    sc = SparkContext(appName="StreamingDemo",master="local[*]")
    ssc = StreamingContext(sc, 1)

    topicsSet = topics.split(",")
    kafkaParams = {"metadata.broker.list": brokers}
    kvs = KafkaUtils.createDirectStream(ssc, topicsSet, kafkaParams)
    values = kvs.map(lambda kv: kv[1])
    count = values.map(lambda v: re.search("key", v)).filter(lambda x: x != None).count
    print("find %d" % (count))

    ssc.start()
    ssc.awaitTermination()


if __name__ == "__main__":
    main()