from datetime import datetime, timedelta
from airflow import DAG
from airflow.providers.amazon.aws.operators.glue import GlueJobOperator
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from airflow.providers.amazon.aws.operators.redshift_data import RedshiftDataOperator
from airflow.operators.python import PythonOperator
from airflow.operators.empty import EmptyOperator
from airflow.utils.task_group import TaskGroup
import logging

logger = logging.getLogger(__name__)

default_args = {
    'owner': 'data-engineering',
    'depends_on_past': False,
    'email_on_failure': True,
    'email': ['data-team@company.com'],
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
    'retry_exponential_backoff': True,
    'max_retry_delay': timedelta(minutes=30),
    'execution_timeout': timedelta(hours=2),
}


def validate_row_count(**context):
    """Validate that the ETL produced expected row counts."""
    ti = context['ti']
    row_count = ti.xcom_pull(task_ids='transform.run_etl_job', key='row_count')
    if row_count is None or row_count < 100:
        raise ValueError(f"Insufficient rows processed: {row_count}")
    logger.info(f"Row count validation passed: {row_count}")


def validate_null_checks(**context):
    """Run null checks on critical columns."""
    import boto3
    redshift = boto3.client('redshift-data')
    result = redshift.execute_statement(
        ClusterIdentifier='analytics-cluster',
        Database='dev',
        Sql="SELECT COUNT(*) FROM analytics.staging WHERE user_id IS NULL"
    )
    # Check result
    logger.info("Null check validation completed")


with DAG(
    'daily_etl_pipeline',
    default_args=default_args,
    description='Production ETL: extract, transform, load, validate',
    schedule_interval='0 2 * * *',
    start_date=datetime(2024, 1, 1),
    catchup=False,
    max_active_runs=1,
    tags=['etl', 'production', 'data-lake'],
    doc_md="""
    ## Daily ETL Pipeline
    - **Extract**: Waits for raw data in S3 landing zone
    - **Transform**: Runs Glue ETL job on partitioned data
    - **Load**: Inserts transformed data into Redshift
    - **Validate**: Runs data quality checks
    """,
) as dag:

    start = EmptyOperator(task_id='start')

    with TaskGroup('extract') as extract_group:
        wait_for_raw = S3KeySensor(
            task_id='wait_for_raw_data',
            bucket_name='data-lake-raw',
            bucket_key='landing/{{ ds }}/',
            wildcard_match=True,
            timeout=3600,
            poke_interval=60,
            mode='reschedule',
        )

    with TaskGroup('transform') as transform_group:
        run_glue = GlueJobOperator(
            task_id='run_etl_job',
            job_name='daily-etl-transform',
            script_location='s3://glue-scripts/etl/transform.py',
            s3_bucket='glue-scripts',
            iam_role_name='GlueETLRole',
            create_job_kwargs={
                'GlueVersion': '4.0',
                'NumberOfWorkers': 10,
                'WorkerType': 'G.1X',
            },
            job_poll_interval=30,
            aws_conn_id='aws_default',
        )

    with TaskGroup('load') as load_group:
        load_redshift = RedshiftDataOperator(
            task_id='load_to_redshift',
            cluster_identifier='analytics-cluster',
            database='dev',
            db_user='admin',
            sql="""
                COPY analytics.fact_daily
                FROM 's3://data-lake-processed/{{ ds }}/'
                IAM_ROLE 'arn:aws:iam::123456789012:role/RedshiftLoadRole'
                FORMAT AS PARQUET;
            """,
            aws_conn_id='aws_default',
        )

    with TaskGroup('validate') as validate_group:
        check_row_count = PythonOperator(
            task_id='validate_row_count',
            python_callable=validate_row_count,
        )
        check_nulls = PythonOperator(
            task_id='validate_null_checks',
            python_callable=validate_null_checks,
        )

    end = EmptyOperator(task_id='end')

    start >> extract_group >> transform_group >> load_group >> validate_group >> end

from airflow.providers.amazon.aws.operators.glue import GlueJobOperator
from airflow.providers.amazon.aws.sensors.glue import GlueJobSensor
from airflow.providers.amazon.aws.operators.lambda_ import LambdaInvokeFunctionOperator
from airflow.providers.amazon.aws.operators.ecs import EcsRunTaskOperator

# Pattern 1: Launch Glue and wait for completion
launch_glue = GlueJobOperator(
    task_id='launch_etl',
    job_name='my-etl-job',
    script_location='s3://scripts/etl.py',
    s3_bucket='my-bucket',
    iam_role_name='GlueRole',
    create_job_kwargs={'NumberOfWorkers': 10},
    aws_conn_id='aws_default',
)

wait_for_glue = GlueJobSensor(
    task_id='wait_for_etl',
    job_name='my-etl-job',
    aws_conn_id='aws_default',
    poke_interval=30,
    timeout=3600,
    mode='reschedule',
)

# Pattern 2: Lambda with payload and response handling
invoke_lambda = LambdaInvokeFunctionOperator(
    task_id='process_record',
    function_name='process-data',
    payload={'bucket': 'my-bucket', 'key': 'data/file.parquet'},
    result_path='$.statusCode',
    aws_conn_id='aws_default',
)

# Pattern 3: ECS Fargate for custom containers
run_container = EcsRunTaskOperator(
    task_id='run_custom_processor',
    cluster='data-cluster',
    task_definition='processor-task',
    launch_type='FARGATE',
    network_configuration={
        'awsvpcConfiguration': {
            'subnets': ['subnet-12345'],
            'securityGroups': ['sg-12345'],
            'assignPublicIp': 'DISABLED',
        }
    },
    aws_conn_id='aws_default',
)