# Copyright 2022-2024 ETSI OSG/SDG TeraFlowSDN (TFS) (https://tfs.etsi.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import logging
from pyspark.sql                  import SparkSession
from pyspark.sql.types            import StructType, StructField, StringType, DoubleType
from pyspark.sql.functions        import from_json, col
from common.tools.kafka.Variables import KafkaConfig, KafkaTopic

LOGGER = logging.getLogger(__name__)

def DefiningSparkSession():
    # Create a Spark session with specific spark verions (3.5.0)
    return SparkSession.builder \
            .appName("Analytics") \
            .config("spark.jars.packages", "org.apache.spark:spark-sql-kafka-0-10_2.12:3.5.0") \
            .getOrCreate()

def SettingKafkaConsumerParams():   # TODO:  create get_kafka_consumer() in common with inputs (bootstrap server, subscribe, startingOffset and failOnDataLoss with default values)
    return {
            # "kafka.bootstrap.servers": '127.0.0.1:9092',
            "kafka.bootstrap.servers": KafkaConfig.get_kafka_address(),
            "subscribe"              : KafkaTopic.VALUE.value,
            "startingOffsets"        : 'latest',
            "failOnDataLoss"         : 'false'              # Optional: Set to "true" to fail the query on data loss
        }

def DefiningRequestSchema():
    return StructType([
            StructField("time_stamp" ,  StringType()  , True),
            StructField("kpi_id"     ,  StringType()  , True),
            StructField("kpi_value"  ,  DoubleType()  , True)
        ])

def SettingKafkaProducerParams():
    return {
            "kafka.bootstrap.servers" : KafkaConfig.get_kafka_address(),
            "topic"                   : KafkaTopic.ANALYTICS_RESPONSE.value
    }

def SparkStreamer(kpi_list):
    """
    Method to perform Spark operation Kafka stream.
    NOTE: Kafka topic to be processesd should have atleast one row before initiating the spark session. 
    """
    kafka_producer_params = SettingKafkaConsumerParams()         # Define the Kafka producer parameters
    kafka_consumer_params = SettingKafkaConsumerParams()         # Define the Kafka consumer parameters
    schema                = DefiningRequestSchema()              # Define the schema for the incoming JSON data
    spark                 = DefiningSparkSession()               # Define the spark session with app name and spark version

    try:
        # Read data from Kafka
        raw_stream_data = spark \
            .readStream \
            .format("kafka") \
            .options(**kafka_consumer_params) \
            .load()

        # Convert the value column from Kafka to a string
        stream_data          = raw_stream_data.selectExpr("CAST(value AS STRING)")
        # Parse the JSON string into a DataFrame with the defined schema
        parsed_stream_data   = stream_data.withColumn("parsed_value", from_json(col("value"), schema))
        # Select the parsed fields
        final_stream_data    = parsed_stream_data.select("parsed_value.*")
        # Filter the stream to only include rows where the kpi_id is in the kpi_list
        filtered_stream_data = final_stream_data.filter(col("kpi_id").isin(kpi_list))

        query = filtered_stream_data \
            .selectExpr("CAST(kpi_id AS STRING) AS key", "to_json(struct(*)) AS value") \
            .writeStream \
            .format("kafka") \
            .option("kafka.bootstrap.servers", KafkaConfig.get_kafka_address()) \
            .option("topic",                   KafkaTopic.ANALYTICS_RESPONSE.value) \
            .option("checkpointLocation",      "/home/tfs/sparkcheckpoint") \
            .outputMode("append")

        # Start the Spark streaming query and write the output to the Kafka topic
        # query = filtered_stream_data \
        #     .selectExpr("CAST(kpi_id AS STRING) AS key", "to_json(struct(*)) AS value") \
        #     .writeStream \
        #     .format("kafka") \
        #     .option(**kafka_producer_params) \
        #     .option("checkpointLocation", "sparkcheckpoint") \
        #     .outputMode("append") \
        #     .start()

        # Start the Spark streaming query
        # query = filtered_stream_data \
        #     .writeStream \
        #     .outputMode("append") \
        #     .format("console")              # You can change this to other output modes or sinks

        # Start the query execution
        query.start().awaitTermination()
    except Exception as e:
        print("Error in Spark streaming process: {:}".format(e))
        LOGGER.debug("Error in Spark streaming process: {:}".format(e))
    finally:
        spark.stop()