from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
from datetime import datetime, timedelta

# Personalized for our elite data engineers 😉
default_args = {
    'owner': 'pritha_queen_of_the_world',
    'depends_on_past': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

with DAG(
    'daily_weather_ingestion',
    default_args=default_args,
    description='Extracts weather data from API and loads to Snowflake',
    schedule_interval='@daily',
    start_date=datetime(2026, 1, 1),
    catchup=False
) as dag:

    extract_data = PythonOperator(
        task_id='fetch_weatherstack_api',
        python_callable=fetch_weather_data_func
    )

    load_to_snowflake = SnowflakeOperator(
        task_id='load_to_warehouse',
        snowflake_conn_id='snowflake_default',
        sql="""
            COPY INTO raw.weather_data 
            FROM @my_s3_stage/weather/ 
            FILE_FORMAT = (TYPE = JSON);
        """
    )

    # Define Task Dependencies (DAG structure)
    extract_data >> load_to_snowflake

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, to_timestamp, avg

# Initialize Spark Session for Distributed Computing
spark = SparkSession.builder \
    .appName("BigData_Transformation") \
    .getOrCreate()

# Read massive dataset from Data Lake (S3)
df = spark.read.parquet("s3a://data-lake/raw/events/")

# Transform data distributed across cluster nodes
clean_df = df.filter(col("status") == "success") \
    .withColumn("event_time", to_timestamp(col("timestamp"))) \
    .groupBy("user_id") \
    .agg(avg("session_duration").alias("avg_duration"))

# Write results back to Lakehouse partitioned by date
clean_df.write \
    .partitionBy("date") \
    .format("delta") \
    .mode("append") \
    .save("s3a://data-lake/processed/user_metrics/")

spark.stop()