from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.postgres.hooks.postgres import PostgresHook
import pandas as pd
from io import StringIO
from datetime import datetime

# Step 1: Extract data from S3
def extract_data(**context):
    s3 = S3Hook(aws_conn_id="aws_default")
    file_obj = s3.read_key(key="raw/sales.csv", bucket_name="my-data-bucket")
    context['ti'].xcom_push(key='raw_csv', value=file_obj)

# Step 2: Transform with pandas
def transform_data(**context):
    raw_csv = context['ti'].xcom_pull(key='raw_csv', task_ids='extract_task')
    df = pd.read_csv(StringIO(raw_csv))
    df_clean = df.dropna().rename(columns={"amount": "sales_amount"})
    context['ti'].xcom_push(key='clean_data', value=df_clean.to_dict('records'))

# Step 3: Load into Redshift
def load_data(**context):
    records = context['ti'].xcom_pull(key='clean_data', task_ids='transform_task')
    df = pd.DataFrame(records)

    redshift = PostgresHook(postgres_conn_id="redshift_conn")
    engine = redshift.get_sqlalchemy_engine()
    df.to_sql("sales", engine, if_exists="replace", index=False)

# DAG definition
with DAG(
    dag_id="warehouse_load_pipeline",
    description="Daily ETL into Redshift",
    start_date=datetime(2025, 1, 1),
    schedule_interval="@daily",
    catchup=False
) as dag:

    extract_task = PythonOperator(
        task_id="extract_task",
        python_callable=extract_data,
        provide_context=True
    )

    transform_task = PythonOperator(
        task_id="transform_task",
        python_callable=transform_data,
        provide_context=True
    )

    load_task = PythonOperator(
        task_id="load_task",
        python_callable=load_data,
        provide_context=True
    )

    extract_task >> transform_task >> load_task