# 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 SettingKafkaParameters():   # 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 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_params = SettingKafkaParameters()         # Define the Kafka 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_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))

        # 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()