import os, sys

def produce_message():
  producer = KafkaProducer(bootstrap_servers=bootstrap_servers,linger_ms=20)
  meesage = os.system('source examplemsg.sh')
  producer.send('devices', message.encode())
  producer.flush

def consume_message():
SparkSubmitOperator
from airflow.contrib.operators.spark_submit_operator import SparkSubmitOperator
.
.
.

exploded_df.printSchema()
exploded_df.show

from pyspark.sql import SparkSession
from pyspark.sql.functions import from_json
from pyspark.sql.types import StringType, StructField, StructType, ArrayType, LongType

spark = SparkSession \
    .builder \
    .appName("Streaming from Kafka") \
    .config("spark.streaming.stopGracefullyOnShutdown", True) \
    .config('spark.jars.packages', 'org.apache.spark:spark-sql-kafka-0-10_2.12:3.3.0') \
    .config("spark.sql.shuffle.partitions", 4) \
    .master("local[*]") \
    .getOrCreate()
# Define Kafka connection properties
kafka_params = {
    "kafka.bootstrap.servers": "125.6.40.186:19092",
    "subscribe": "devices",
    "startingOffsets": "earliest"
}

# Define JSON Schema

json_schema = StructType([
    StructField('customerId', StringType(), True),
    StructField('data', StructType([
        StructField('devices', ArrayType(StructType([
            StructField('deviceId', StringType(), True),
            StructField('measure', StringType(), True),
            StructField('status', StringType(), True),
            StructField('temperature', LongType(), True)
        ]), True), True)
    ]), True),
    StructField('eventId', StringType(), True),
    StructField('eventOffset', LongType(), True),
    StructField('eventPublisher', StringType(), True),
    StructField('eventTime', StringType(), True)
])

# Read Kafka messages
streaming_df = spark \
    .readStream \
    .format("kafka") \
    .options(**kafka_params) \
    .load()

# Parse JSON messages
json_df = streaming_df.selectExpr("CAST(value AS STRING) AS value") \

json_expanded_df = json_df.withColumn("value", from_json(json_df["value"], json_schema)).select("value.*")

#print(json_schema)
#json_expanded_df.printSchema()

###############################################

from pyspark.sql.functions import explode, col


exploded_df = json_expanded_df \
    .select("customerId", "eventId", "eventOffset", "eventPublisher", "eventTime", "data") \
    .withColumn("devices", explode("data.devices")) \
    .drop("data")

#exploded_df.printSchema()
#exploded_df.show

flattened_df = exploded_df \
    .selectExpr("customerId", "eventId", "eventOffset", "eventPublisher", "cast(eventTime as timestamp) as eventTime",
                "devices.deviceId as deviceId", "devices.measure as measure",
                "devices.status as status", "devices.temperature as temperature")

#flattened_df.printSchema()


# Aggregate the dataframes to find the average temparature
# per Customer per device throughout the day for SUCCESS events
from pyspark.sql.functions import to_date, avg

agg_df = flattened_df.where("STATUS = 'SUCCESS'") \
    .withColumn("eventDate", to_date("eventTime", "yyyy-MM-dd")) \
    .groupBy("customerId","deviceId","eventDate") \
    .agg(avg("temperature").alias("avg_temp"))

# Write the output to console sink to check the output
writing_df = agg_df.writeStream \
    .format("console") \
    .option("checkpointLocation","checkpoint_dir") \
    .outputMode("complete") \
    .start()

# Start the streaming application to run until the following happens
# 1. Exception in the running program
# 2. Manual Interruption
writing_df.awaitTermination()


#query = flattened_df \
#    .writeStream \
#    .outputMode("append") \
#    .format("console") \
#    .start()

# Wait for the termination of the query
#query.awaitTermination()

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from kafka import KafkaProducer, KafkaConsumer
from kafka.admin import KafkaAdminClient, NewTopic
from datetime import datetime
from airflow.models import TaskInstance
from kafka.errors import TopicAlreadyExistsError
from pyspark.sql import SparkSession
from pyspark.sql.functions import from_json
from pyspark.sql.types import StringType, StructField, StructType, ArrayType, LongType

# Define Kafka broker
bootstrap_servers = '125.6.40.186:19092'
spark_master = 'bami-cluster2:7077'
topic_name = 'devices'


# Check with Kafka
def check_connected_kafka():
    try:
        producer = KafkaProducer(bootstrap_servers=bootstrap_servers)
        consumer = KafkaConsumer(topic_name, bootstrap_servers=bootstrap_servers)
        print(f"#####SUCCESS CONNECTING WITH KAFKA#######")
        return True
    except Exception as e:
        print(f"Error connecting to Kafka: {str(e)}")
        return False

# Check with Spark
def check_connedted_spark():
    try:
    spark = SparkSession \
        .builder \
        .appName("Streaming from Kafka") \
        .config("spark.streaming.stopGracefullyOnShutdown", True) \
        .config('spark.jars.packages', 'org.apache.spark:spark-sql-kafka-0-10_2.12:3.3.0') \
        .config("spark.sql.shuffle.partitions", 4) \
        .master(spark_master) \
        .getOrCreate()
        return True
    except Exception as e:
        print(f"FAIL..............FAIL SPARK........: {str{e}}")
        return False
# Define the DAG
default_args = {
    'owner': 'admin',
    'depends_on_past': False,
    'start_date': datetime(2024, 2, 27),
    'retries': 1,
}
dag = DAG(
    'airflow_kafka_spark_interworking',
    default_args=default_args,
    description='DAG for testing interworking of Airflow and Kafka and Spark',
    schedule_interval=None,
)

# Define tasks
check_connected_kafka_task = PythonOperator(
    task_id='check_connected_kafka',
    python_callable=check_connected_kafka,
    dag=dag,
)

check_connedted_spark_task = PythonOperator(
    task_id='check_connedted_spark',
    python_callable=check_connedted_spark,
    dag=dags,
)

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from kafka import KafkaProducer, KafkaConsumer
from kafka.admin import KafkaAdminClient, NewTopic
from datetime import datetime
from airflow.models import TaskInstance
from kafka.errors import TopicAlreadyExistsError
from airflow.contrib.hooks.spark_submit_hook import SparkSubmitHook
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator 


# Define Kafka broker
bootstrap_servers = '125.6.40.186:19092'
spark_master = 'bami-cluster2:7077'
topic_name = 'devices'


# Check with Kafka
def check_connected_kafka():
    try:
        producer = KafkaProducer(bootstrap_servers=bootstrap_servers)
        consumer = KafkaConsumer(topic_name, bootstrap_servers=bootstrap_servers)
        print(f"#####SUCCESS CONNECTING WITH KAFKA#######")
        return True
    except Exception as e:
        print(f"Error connecting to Kafka: {str(e)}")
        return False

with DAG(dag_id='check_connected_kafka', start_date=datetime(2023, 1, 1)) as dag: 
    submit_spark_app = SparkSubmitOperator( 
        task_id='check_connected_kafka', 
        conn_id='spark_default', 
        application='/opt/bitnami/spark.py', 
        name='Streaming_from_Kafka'
    ) 

# Define the DAG
default_args = {
    'owner': 'admin',
    'depends_on_past': False,
    'start_date': datetime(2024, 2, 27),
    'retries': 1,
}
dag = DAG(
    'airflow_kafka_spark_interworking',
    default_args=default_args,
    description='DAG for testing interworking of Airflow and Kafka and Spark',
    schedule_interval=None,
)

# Define tasks
check_connected_kafka_task = PythonOperator(
    task_id='check_connected_kafka',
    python_callable=check_connected_kafka,
    dag=dag,
)

check_connedted_spark_task = PythonOperator(
    task_id='check_connedted_spark',
    python_callable=check_connedted_spark,
    dag=dags,
)