from __future__ import annotations

import typing import Any, Callable

import pendulum
@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example"],
)

def example_pyspark():
    @task.pyspark(conn_id="spark-default")
    def spark_task(spark: SparkSession, sc: SparkContext) -> pd.DataFrame:
        df = spark.createDataFrame(
            [
                (1, "John Doe", 21),
                (2, "Jane Doe", 22),
                (3, "Joe Bloggs", 23),
            ],
            ["id", "name", "age"],
        )
        df.show()
        return df.toPandas()
    @task
    def print_df(df: pd.DataFrame):
        print(df)
    df = spark_task()
    print_df(df)
# work around pre-commit
dag = example_pyspark()  # type: ignore
test_run = get_test_run(dag)


from __future__ import annotations

import logging
import os
from datetime import timedelta
from typing import TYPE_CHECKING, Callable

from tabulate import tabulate

from airflow.utils.state import State

if TYPE_CHECKING:
    from airflow.utils.context import Context

def get_test_run(dag):
    def callback(context: Context):
        try:
            ti = context["dag_run"].get_task_instances()
            if not ti:
                logging.warning("Could not retrieve tasks that ran in the DAG, cannot display a summary")
                return

            ti.sort(key=lambda x: x.end_date)

            headers = ["Task ID", "Status", "Duration (s)"]
            results = []
            prev_time = ti[0].end_date - timedelta(seconds=ti[0].duration)
            for t in ti:
                results.append([t.task_id, t.state, f"{(t.end_date - prev_time).total_seconds():.1f}"])
                prev_time = t.end_date

            logging.info("EXECUTION SUMMARY:\n" + tabulate(results, headers=headers, tablefmt="fancy_grid"))
        except Exception as e:
            logging.error(f"Error occurred during callback execution: {e}")

    def add_callback(current: List[Callable], new: Callable) -> List[Callable]:
        return current + [new]

    def test_run():
        dag.on_failure_callback = add_callback(dag.on_failure_callback or [], callback)
        dag.on_success_callback = add_callback(dag.on_success_callback or [], callback)
        dag.clear(dag_run_state=State.QUEUED)
        dag.run()

    return test_run

def get_test_env_id(env_var_name: str = "SYSTEM_TESTS_ENV_ID"):
    return os.environ.get(env_var_name)

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

def initialize_spark_session():
    spark = SparkSession \
        .builder \
        .appName("pipeline") \
        .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() 

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2024, 3, 5),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

# Instantiate the DAG
dag = DAG(
    'spark_data_pipeline',
    default_args=default_args,
    description='A DAG to process data from Kafka using Spark',
    schedule_interval=None,
)

# Define tasks
initialize_spark_task = PythonOperator(
    task_id='initialize_spark_session',
    python_callable=initialize_spark_session,
    dag=dag,
)


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

writing_df.awaitTermination()

airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator