"""
AWS Glue ETL Job - Bronze Layer Ingestion
Ingests raw data from multiple sources into the Bronze layer.
"""
import sys
import logging
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job
from pyspark.sql.functions import current_timestamp, input_file_name, lit, col

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

args = getResolvedOptions(sys.argv, ['JOB_NAME', 'SOURCE_PATH', 'DATABASE_NAME'])
sc = SparkContext()
glueContext = GlueContext(sc)
spark = glueContext.spark_session
job = Job(glueContext)

try:
    logger.info(f"Starting Bronze ingestion from {args['SOURCE_PATH']}")
    
    raw_df = spark.read.json(args['SOURCE_PATH'])
    record_count = raw_df.count()
    logger.info(f"Read {record_count} records from source")
    
    if record_count == 0:
        logger.warning("No records found. Skipping ingestion.")
        sys.exit(0)
    
    bronze_df = raw_df \
        .withColumn("ingestion_timestamp", current_timestamp()) \
        .withColumn("source_file", input_file_name()) \
        .withColumn("batch_id", lit(f"batch_{current_timestamp().strftime('%Y%m%d_%H%M%S')}"))
    
    bronze_df.write.mode("append") \
        .partitionBy("event_date") \
        .parquet("s3://data-lake/bronze/events/")
    
    logger.info(f"Successfully ingested {record_count} records to Bronze layer")
    
except Exception as e:
    logger.error(f"Bronze ingestion failed: {str(e)}")
    raise

job.commit()

"""
Bronze to Silver transformation with deduplication and quality checks.
"""
from pyspark.sql.functions import col, row_number, when, count, sum as spark_sum
from pyspark.sql.window import Window

def process_silver(spark, bronze_path, silver_path):
    """Transform Bronze data to Silver with quality validation."""
    try:
        bronze_df = spark.read.parquet(bronze_path)
        logger.info(f"Read {bronze_df.count()} records from Bronze")
        
        # Deduplicate by primary key
        window_spec = Window.partitionBy("event_id").orderBy(col("ingestion_timestamp").desc())
        deduped_df = bronze_df \
            .withColumn("row_num", row_number().over(window_spec)) \
            .filter(col("row_num") == 1) \
            .drop("row_num")
        
        # Type casting and validation
        validated_df = deduped_df \
            .withColumn("event_id", col("event_id").cast("string")) \
            .withColumn("user_id", col("user_id").cast("long")) \
            .withColumn("amount", col("amount").cast("decimal(10,2)"))
        
        # Data quality filters
        valid_df = validated_df.filter(
            col("event_id").isNotNull() &
            col("user_id").isNotNull() &
            (col("amount") >= 0) &
            (col("amount") <= 1000000)
        )
        
        invalid_count = deduped_df.count() - valid_df.count()
        if invalid_count > 0:
            logger.warning(f"Filtered {invalid_count} invalid records")
        
        valid_df.write.mode("overwrite") \
            .partitionBy("event_date", "event_type") \
            .parquet(silver_path)
        
        logger.info(f"Successfully wrote {valid_df.count()} records to Silver layer")
        
    except Exception as e:
        logger.error(f"Silver processing failed: {str(e)}")
        raise

"""
Silver to Gold transformation with business aggregations.
"""
from pyspark.sql.functions import sum as spark_sum, count, avg, date_trunc

def process_gold(spark, silver_path, gold_path):
    """Aggregate Silver data to Gold layer for BI consumption."""
    try:
        silver_df = spark.read.parquet(silver_path)
        logger.info(f"Read {silver_df.count()} records from Silver")
        
        # Daily revenue aggregation
        daily_revenue = silver_df.groupBy("event_date", "product_category") \
            .agg(
                spark_sum("amount").alias("total_revenue"),
                count("event_id").alias("transaction_count"),
                avg("amount").alias("avg_order_value")
            )
        
        daily_revenue.write.mode("overwrite") \
            .partitionBy("event_date") \
            .parquet(f"{gold_path}/daily_revenue/")
        
        # Dimension tables
        product_dim = silver_df.select(
            "product_id", "product_name", "product_category"
        ).distinct()
        
        product_dim.write.mode("overwrite") \
            .parquet(f"{gold_path}/dimensions/product_dim/")
        
        logger.info("Successfully wrote Gold layer aggregations")
        
    except Exception as e:
        logger.error(f"Gold processing failed: {str(e)}")
        raise

#!/bin/bash
# Configure S3 data lake bucket with best practices

BUCKET_NAME="company-data-lake-prod"
REGION="us-east-1"

# Create bucket
aws s3api create-bucket \
    --bucket "$BUCKET_NAME" \
    --region "$REGION" \
    --create-bucket-configuration LocationConstraint="$REGION"

# Enable versioning
aws s3api put-bucket-versioning \
    --bucket "$BUCKET_NAME" \
    --versioning-configuration Status=Enabled

# Enable default encryption
aws s3api put-bucket-encryption \
    --bucket "$BUCKET_NAME" \
    --server-side-encryption-configuration '{
        "Rules": [
            {
                "ApplyServerSideEncryptionByDefault": {
                    "SSEAlgorithm": "aws:kms",
                    "KMSMasterKeyID": "arn:aws:kms:us-east-1:ACCOUNT:key/KEY_ID"
                }
            }
        ]
    }'

# Lifecycle policies for cost optimization
aws s3api put-bucket-lifecycle-configuration \
    --bucket "$BUCKET_NAME" \
    --lifecycle-configuration '{
        "Rules": [
            {
                "ID": "TieredStorage",
                "Status": "Enabled",
                "Transitions": [
                    {"Days": 90, "StorageClass": "STANDARD_IA"},
                    {"Days": 180, "StorageClass": "GLACIER"},
                    {"Days": 365, "StorageClass": "DEEP_ARCHIVE"}
                ]
            }
        ]
    }'

# Block public access
aws s3api put-public-access-block \
    --bucket "$BUCKET_NAME" \
    --public-access-block-configuration \
        BlockPublicAcls=true,IgnorePublicAcls=true,BlockPublicPolicy=true,RestrictPublicBuckets=true