import boto3
import os
import json

def lambda_handler(event, context):
    RDS_CLUSTER_NAME = os.getenv("RDS_CLUSTER_NAME")
    S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
    EXPORT_IAM_ROLE = os.getenv("EXPORT_IAM_ROLE")
    KMS_KEY_ID = os.getenv("KMS_KEY_ID")
    SNAPSHOT_ID = os.getenv("SNAPSHOT_ID")  # 생성할 스냅샷 아이디
    S3_PREFIX = f"rds_snapshots/{SNAPSHOT_ID}.sql"

    # RDS 및 S3 클라이언트 생성
    rds_client = boto3.client("rds", region_name="ap-northeast-2")

    # RDS Aurora 클러스터 스냅샷 생성
    rds_client.create_db_cluster_snapshot(DBClusterSnapshotIdentifier=SNAPSHOT_ID, DBClusterIdentifier=RDS_CLUSTER_NAME)

    # RDS Aurora 클러스터 스냅샷이 사용 가능할 때까지 대기 (상태가 "available"이 될 때까지)
    waiter = rds_client.get_waiter("db_cluster_snapshot_available")
    waiter.wait(DBClusterSnapshotIdentifier=SNAPSHOT_ID)

    # RDS Aurora 클러스터 스냅샷을 S3 버킷으로 내보내기
    rds_cluster_snapshot_arn = rds_client.describe_db_cluster_snapshots(DBClusterSnapshotIdentifier=SNAPSHOT_ID)[
        "DBClusterSnapshots"
    ][0]["DBClusterSnapshotArn"]
    response = rds_client.start_export_task(
        ExportTaskIdentifier=SNAPSHOT_ID,
        SourceArn=rds_cluster_snapshot_arn,
        S3BucketName=S3_BUCKET_NAME,
        IamRoleArn=EXPORT_IAM_ROLE,
        KmsKeyId=KMS_KEY_ID,
        S3Prefix=S3_PREFIX,
    )
    return {"statusCode": 200, "body": json.dumps({"message": "OK"})}