from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from datetime import datetime, timedelta
import pandas as pd

default_args = {
    "owner": "data-team",
    "depends_on_past": False,
    "email_on_failure": True,
    "email": ["alerts@company.com"],
    "retries": 3,
    "retry_delay": timedelta(minutes=5),
}

def ingest_raw_events(**context):
    execution_date = context["ds"]
    hook = PostgresHook(postgres_conn_id="source_db")
    df = hook.get_pandas_df(
        f"SELECT * FROM events WHERE date = '{execution_date}'"
    )
    df.to_parquet(f"/tmp/raw_events_{execution_date}.parquet")
    return f"Ingested {len(df)} rows"

def clean_and_enrich(**context):
    execution_date = context["ds"]
    df = pd.read_parquet(f"/tmp/raw_events_{execution_date}.parquet")
    
    # Clean
    df = df.dropna(subset=["user_id", "event_type"])
    df["event_type"] = df["event_type"].str.lower().str.strip()
    
    # Enrich
    geo_hook = PostgresHook(postgres_conn_id="geo_db")
    geo_df = geo_hook.get_pandas_df("SELECT ip, country FROM ip_lookup")
    df = df.merge(geo_df, on="ip", how="left")
    
    df.to_parquet(f"/tmp/clean_events_{execution_date}.parquet")
    return f"Cleaned {len(df)} rows"

def aggregate_metrics(**context):
    execution_date = context["ds"]
    df = pd.read_parquet(f"/tmp/clean_events_{execution_date}.parquet")
    
    metrics = df.groupby(["event_type", "country"]).agg(
        event_count=("event_id", "count"),
        unique_users=("user_id", "nunique"),
    ).reset_index()
    metrics["date"] = execution_date
    
    metrics.to_parquet(f"/tmp/metrics_{execution_date}.parquet")

def export_to_warehouse(**context):
    execution_date = context["ds"]
    metrics = pd.read_parquet(f"/tmp/metrics_{execution_date}.parquet")
    
    wh_hook = PostgresHook(postgres_conn_id="warehouse")
    wh_hook.insert_rows(
        table="daily_metrics",
        rows=metrics.values.tolist(),
        target_fields=metrics.columns.tolist(),
        replace=True,
    )

with DAG(
    dag_id="event_pipeline",
    default_args=default_args,
    description="Daily event processing pipeline",
    schedule_interval="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=True,
    tags=["events", "daily"],
) as dag:

    ingest = PythonOperator(
        task_id="ingest_raw_events",
        python_callable=ingest_raw_events,
    )

    clean = PythonOperator(
        task_id="clean_and_enrich",
        python_callable=clean_and_enrich,
    )

    aggregate = PythonOperator(
        task_id="aggregate_metrics",
        python_callable=aggregate_metrics,
    )

    export = PythonOperator(
        task_id="export_to_warehouse",
        python_callable=export_to_warehouse,
    )

    ingest >> clean >> aggregate >> export

from dagster import (
    asset, AssetIn, DailyPartitionsDefinition, Output,
    define_asset_job, AssetSelection, ScheduleDefinition,
    Definitions, MetadataValue, FreshnessPolicy,
)
import pandas as pd
from dagster_postgres import PostgreSQLIOManager
from datetime import timedelta

daily_partitions = DailyPartitionsDefinition(start_date="2024-01-01")

@asset(
    partitions_def=daily_partitions,
    group_name="event_pipeline",
    description="Raw events ingested from source database",
    metadata={"source": "postgres://source_db/events"},
)
def raw_events(context) -> pd.DataFrame:
    partition_date = context.partition_key
    # In production: use resource injection for DB connection
    df = pd.DataFrame({
        "event_id": range(1000),
        "user_id": [f"user_{i % 100}" for i in range(1000)],
        "event_type": ["click", "view", "purchase"][i % 3] for i in range(1000),
        "ip": [f"1.2.3.{i % 255}" for i in range(1000)],
        "date": partition_date,
    })
    context.add_output_metadata({
        "row_count": MetadataValue.int(len(df)),
        "partition": MetadataValue.text(partition_date),
    })
    return df

@asset(
    ins={"raw_events": AssetIn()},
    partitions_def=daily_partitions,
    group_name="event_pipeline",
    description="Cleaned and geo-enriched events",
    freshness_policy=FreshnessPolicy(maximum_lag_minutes=60),
)
def clean_events(context, raw_events: pd.DataFrame) -> pd.DataFrame:
    df = raw_events.dropna(subset=["user_id", "event_type"]).copy()
    df["event_type"] = df["event_type"].str.lower().str.strip()
    
    # Geo enrichment (simplified)
    df["country"] = df["ip"].apply(lambda ip: "DE" if ip.startswith("1.") else "US")
    
    context.add_output_metadata({
        "row_count": MetadataValue.int(len(df)),
        "null_drop_count": MetadataValue.int(len(raw_events) - len(df)),
    })
    return df

@asset(
    ins={"clean_events": AssetIn()},
    partitions_def=daily_partitions,
    group_name="event_pipeline",
    description="Daily aggregated metrics by event type and country",
)
def daily_metrics(context, clean_events: pd.DataFrame) -> pd.DataFrame:
    metrics = clean_events.groupby(["event_type", "country"]).agg(
        event_count=("event_id", "count"),
        unique_users=("user_id", "nunique"),
    ).reset_index()
    metrics["date"] = context.partition_key
    
    context.add_output_metadata({
        "metric_count": MetadataValue.int(len(metrics)),
        "preview": MetadataValue.md(metrics.head().to_markdown()),
    })
    return metrics

# Define job and schedule
event_pipeline_job = define_asset_job(
    name="event_pipeline_job",
    selection=AssetSelection.groups("event_pipeline"),
    partitions_def=daily_partitions,
)

daily_schedule = ScheduleDefinition(
    job=event_pipeline_job,
    cron_schedule="0 6 * * *",
    execution_timezone="UTC",
)

defs = Definitions(
    assets=[raw_events, clean_events, daily_metrics],
    jobs=[event_pipeline_job],
    schedules=[daily_schedule],
    resources={
        "io_manager": PostgreSQLIOManager(
            host="warehouse-host",
            db_name="analytics",
        ),
    },
)

from prefect import flow, task, get_run_logger
from prefect.tasks import task_input_hash
from prefect.deployments import Deployment
from prefect.server.schemas.schedules import CronSchedule
from datetime import timedelta, date
import pandas as pd

@task(
    retries=3,
    retry_delay_seconds=60,
    cache_key_fn=task_input_hash,
    cache_expiration=timedelta(hours=12),
    description="Ingest raw events from source DB",
)
def ingest_raw_events(execution_date: str) -> pd.DataFrame:
    logger = get_run_logger()
    logger.info(f"Ingesting events for {execution_date}")
    
    # Simulated ingestion
    df = pd.DataFrame({
        "event_id": range(1000),
        "user_id": [f"user_{i % 100}" for i in range(1000)],
        "event_type": ["click", "view", "purchase"][i % 3] for i in range(1000),
        "ip": [f"1.2.3.{i % 255}" for i in range(1000)],
    })
    logger.info(f"Ingested {len(df)} rows")
    return df

@task(
    retries=2,
    description="Clean and geo-enrich events",
)
def clean_and_enrich(df: pd.DataFrame) -> pd.DataFrame:
    logger = get_run_logger()
    initial_count = len(df)
    
    df = df.dropna(subset=["user_id", "event_type"]).copy()
    df["event_type"] = df["event_type"].str.lower().str.strip()
    df["country"] = df["ip"].apply(lambda ip: "DE" if ip.startswith("1.") else "US")
    
    logger.info(f"Cleaned: {initial_count} → {len(df)} rows ({initial_count - len(df)} dropped)")
    return df

@task(description="Aggregate daily metrics")
def aggregate_metrics(df: pd.DataFrame, execution_date: str) -> pd.DataFrame:
    metrics = df.groupby(["event_type", "country"]).agg(
        event_count=("event_id", "count"),
        unique_users=("user_id", "nunique"),
    ).reset_index()
    metrics["date"] = execution_date
    return metrics

@task(description="Export metrics to data warehouse")
def export_to_warehouse(metrics: pd.DataFrame) -> None:
    logger = get_run_logger()
    # In production: use Prefect block for DB credentials
    logger.info(f"Exporting {len(metrics)} metric rows to warehouse")
    # metrics.to_sql("daily_metrics", con=engine, if_exists="append", index=False)

@flow(
    name="event-pipeline",
    description="Daily event processing: ingest → clean → aggregate → export",
    retries=1,
    retry_delay_seconds=300,
)
def event_pipeline(execution_date: str = None):
    if execution_date is None:
        execution_date = date.today().isoformat()
    
    logger = get_run_logger()
    logger.info(f"Starting event pipeline for {execution_date}")
    
    # Tasks execute in dependency order
    raw = ingest_raw_events(execution_date)
    clean = clean_and_enrich(raw)
    metrics = aggregate_metrics(clean, execution_date)
    export_to_warehouse(metrics)
    
    logger.info("Pipeline complete ✓")

# Deploy with schedule
if __name__ == "__main__":
    deployment = Deployment.build_from_flow(
        flow=event_pipeline,
        name="daily-event-pipeline",
        schedule=CronSchedule(cron="0 6 * * *", timezone="UTC"),
        work_queue_name="default",
        tags=["events", "daily"],
    )
    deployment.apply()

flowchart TD
    A[New orchestration project] --> B{Do you have existing Airflow?}
    B -->|Yes + happy| C[Stay on Airflow<br/>consider Astronomer]
    B -->|Yes + unhappy| D{Main pain point?}
    B -->|No| E{How do you think about pipelines?}
    
    D -->|Dev experience, boilerplate| F[Migrate to Prefect]
    D -->|Data quality, asset tracking| G[Migrate to Dagster]
    D -->|Scale, cost| H[Consider MWAA or Astro]
    
    E -->|Assets and data products| I[Dagster]
    E -->|Tasks and workflows| J{Team background?}
    
    J -->|Data engineers, enterprise| K[Airflow]
    J -->|Python devs, ML team, startups| L[Prefect]
    
    style I fill:#6bcb77,color:#333
    style L fill:#4d96ff,color:#fff
    style K fill:#ff6b6b,color:#fff