import boto3
import json
import logging
import time
from datetime import datetime
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)

class DataBrewManager:
    """Production-grade AWS Glue DataBrew manager."""

    def __init__(self, region: str = 'us-east-1'):
        self.client = boto3.client('databrew', region_name=region)

    def create_dataset(
        self,
        name: str,
        input_location: str,
        format: str = 'PARQUET',
        path_options: Optional[Dict] = None
    ) -> Dict:
        """Create a DataBrew dataset from S3."""
        try:
            kwargs = {
                'Name': name,
                'Input': {
                    'S3InputDefinition': {
                        'Bucket': input_location.split('/')[2],
                        'Key': '/'.join(input_location.split('/')[3:])
                    }
                },
                'FormatOptions': {
                    'Json': {'MultiLine': True} if format == 'JSON' else {}
                },
                'Format': format
            }
            if path_options:
                kwargs['Input']['DataCatalogInputDefinition'] = path_options

            response = self.client.create_dataset(**kwargs)
            logger.info(f"Created dataset: {name}")
            return response
        except Exception as e:
            logger.error(f"Failed to create dataset: {e}")
            raise

    def create_recipe_project(
        self,
        project_name: str,
        dataset_name: str,
        role_arn: str,
        sample_size: str = '10000'
    ) -> Dict:
        """Create a DataBrew project for interactive exploration."""
        try:
            response = self.client.create_project(
                Name=project_name,
                DatasetName=dataset_name,
                RoleArn=role_arn,
                Sample={
                    'Size': int(sample_size)
                }
            )
            logger.info(f"Created project: {project_name}")
            return response
        except Exception as e:
            logger.error(f"Failed to create project: {e}")
            raise

    def create_recipe_job(
        self,
        job_name: str,
        dataset_name: str,
        recipe_version: str = 'LATEST_PUBLISHED',
        output_location: str = '',
        output_format: str = 'PARQUET',
        role_arn: str = '',
        max_capacity: int = 10,
        timeout: int = 3600
    ) -> Dict:
        """Create a DataBrew recipe job."""
        try:
            response = self.client.create_job(
                Name=job_name,
                Type='RECIPE',
                DatasetName=dataset_name,
                RecipeVersion=recipe_version,
                RoleArn=role_arn,
                MaxCapacity=max_capacity,
                Timeout=timeout,
                Output={
                    'Location': {
                        'Bucket': output_location.split('/')[2],
                        'Key': '/'.join(output_location.split('/')[3:])
                    },
                    'Format': output_format,
                    'PartitionColumns': []
                },
                EncryptionOptions={
                    'SseKmsKey': 'arn:aws:kms:us-east-1:*:key/*',
                    'S3SseKmsEnabled': True
                }
            )
            logger.info(f"Created recipe job: {job_name}")
            return response
        except Exception as e:
            logger.error(f"Failed to create recipe job: {e}")
            raise

    def create_profile_job(
        self,
        job_name: str,
        dataset_name: str,
        output_location: str,
        role_arn: str,
        max_samples: int = 10000,
        pattern_rules: Optional[List[Dict]] = None
    ) -> Dict:
        """Create a DataBrew profile job for data profiling."""
        try:
            response = self.client.create_job(
                Name=job_name,
                Type='PROFILE',
                DatasetName=dataset_name,
                RoleArn=role_arn,
                MaxSamples=max_samples,
                Output={
                    'Location': {
                        'Bucket': output_location.split('/')[2],
                        'Key': '/'.join(output_location.split('/')[3:])
                    },
                    'Format': 'JSON'
                }
            )
            logger.info(f"Created profile job: {job_name}")
            return response
        except Exception as e:
            logger.error(f"Failed to create profile job: {e}")
            raise

    def run_job(
        self,
        job_name: str,
        run_id: Optional[str] = None
    ) -> Dict:
        """Start a DataBrew job run."""
        try:
            response = self.client.start_job_run(
                JobName=job_name
            )
            run_id = response['RunId']
            logger.info(f"Started job {job_name}, run: {run_id}")
            return response
        except Exception as e:
            logger.error(f"Failed to start job: {e}")
            raise

    def wait_for_job(
        self,
        job_name: str,
        run_id: str,
        poll_interval: int = 30,
        max_wait: int = 7200
    ) -> Dict:
        """Wait for a job run to complete."""
        start_time = time.time()
        while time.time() - start_time < max_wait:
            response = self.client.get_job_run(
                JobName=job_name,
                RunId=run_id
            )
            state = response['JobRun']['State']
            if state == 'SUCCEEDED':
                logger.info(f"Job {job_name} completed successfully")
                return response['JobRun']
            elif state in ('FAILED', 'STOPPED'):
                logger.error(
                    f"Job {job_name} {state}: "
                    f"{response['JobRun'].get('ErrorMessage', 'Unknown')}"
                )
                return response['JobRun']
            logger.info(f"Job {job_name} state: {state}, waiting...")
            time.sleep(poll_interval)

        raise TimeoutError(f"Job {job_name} timed out after {max_wait}s")

    def publish_recipe(
        self,
        project_name: str,
        recipe_name: str,
        description: str = ''
    ) -> str:
        """Publish a recipe from a project."""
        try:
            response = self.client.publish_project(
                Name=project_name
            )
            version = response['RecipeVersion']
            logger.info(
                f"Published recipe {recipe_name}: version {version}"
            )
            return version
        except Exception as e:
            logger.error(f"Failed to publish recipe: {e}")
            raise

    def batch_delete_recipe_version(
        self,
        recipe_name: str,
        versions: List[str]
    ) -> Dict:
        """Batch delete recipe versions."""
        try:
            response = self.client.batch_delete_recipe_version(
                Name=recipe_name,
                RecipeVersions=versions
            )
            logger.info(
                f"Deleted {len(versions)} versions from {recipe_name}"
            )
            return response
        except Exception as e:
            logger.error(f"Failed to batch delete versions: {e}")
            raise


# Production usage
if __name__ == '__main__':
    databrew = DataBrewManager()

    # Create dataset
    dataset = databrew.create_dataset(
        name='customer-data-raw',
        input_location='s3://data-lake-raw/customers/',
        format='PARQUET'
    )

    # Create project for interactive exploration
    project = databrew.create_recipe_project(
        project_name='customer-cleaning-project',
        dataset_name='customer-data-raw',
        role_arn='arn:aws:iam::*:role/databrew-service-role',
        sample_size='50000'
    )

    print(f"Project created: {project}")

#!/bin/bash
# DataBrew job monitoring and scheduling script

set -euo pipefail

PROJECT_NAME="${1:-customer-cleaning-project}"
DATASET_NAME="${2:-customer-data-raw}"
REGION="${AWS_REGION:-us-east-1}"

echo "=== AWS Glue DataBrew Status ==="
echo "Project: ${PROJECT_NAME}"
echo "Dataset: ${DATASET_NAME}"
echo "Region: ${REGION}"
echo "Timestamp: $(date -u +%Y-%m-%dT%H:%M:%SZ)"

# List datasets
echo ""
echo "=== Datasets ==="
aws databrew list-datasets \
  --region "${REGION}" \
  --query 'datasets[*].[Name,Format,CreateTime]' \
  --output table

# List projects
echo ""
echo "=== Projects ==="
aws databrew list-projects \
  --region "${REGION}" \
  --query 'projects[*].[Name,DatasetName,CreateTime]' \
  --output table

# List jobs
echo ""
echo "=== Jobs ==="
aws databrew list-jobs \
  --region "${REGION}" \
  --query 'jobs[*].[Name,Type,DatasetName,State]' \
  --output table

# Get recent job runs
echo ""
echo "=== Recent Job Runs ==="
aws databrew list-job-runs \
  --job-name "daily-customer-clean" \
  --region "${REGION}" \
  --query 'jobRuns[:5].[*].[RunId,State,StartTime,EndTime,RunBytes]' \
  --output table

# Check data quality metrics
echo ""
echo "=== Data Quality Profile Output ==="
S3_PROFILE_PATH="s3://databrew-profiles/${DATASET_NAME}/latest"
aws s3 ls "${S3_PROFILE_PATH}" --recursive --human-readable 2>/dev/null || echo "No profile output found"

# Monitor CloudWatch metrics
echo ""
echo "=== DataBrew Metrics ==="
aws cloudwatch get-metric-statistics \
  --namespace AWS/DataBrew \
  --metric-name JobRuns \
  --start-time "$(date -u -d '1 hour ago' +%Y-%m-%dT%H:%M:%SZ)" \
  --end-time "$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
  --period 300 \
  --statistics Sum \
  --region "${REGION}" | jq -r '.Datapoints[] | "\(.Timestamp): \(.Sum) runs"'

# Start a profile job
echo ""
echo "=== Starting Profile Job ==="
aws databrew start-job-run \
  --job-name "customer-profile-daily" \
  --region "${REGION}" | jq '.RunId'

echo ""
echo "Monitoring complete."