#!/usr/bin/env bash
set -xeuo pipefail
# This script should be idempotent and retriable
# It's used to deploy a new region
# When adding a new dataplane aws account, this script should be
# invoked several times, each with a single region
remote_config=${CONFIG_PATH}
local_config=/tmp/base-config.yaml

bucket_name=${remote_config#s3://}
bucket_name=${bucket_name%%/*}
output_bucket_region=$(aws s3api get-bucket-location --bucket=$bucket_name --query="LocationConstraint" \
  --region=us-west-2 --output=text --no-cli-pager)
aws s3 cp "${remote_config}" "${local_config}" --region=$output_bucket_region

RESOURCE_OWNER_TAG=$(yq '.jobOptions.resourceOwnerTag' "${local_config}")
ACCOUNT_ID=$(yq -r -ojson '.aws.dataplaneAccounts.0' "${local_config}")
PROFILE=${ACCOUNT_ID}
ENV=$(yq '.env' "${local_config}")
REGION=$(yq '.region' "${local_config}")
LOGICAL_REGION_NAME=$(yq '.name' "${local_config}")
ROLE_ARN_FH=$(yq ".jobOptions.sli.firehouseRoleArn // \"arn:aws:iam::${ACCOUNT_ID}:role/tidbcloud-sli-firehose-role\"" "${local_config}")

# ==============================================================================
# Resource Names
# ==============================================================================
BUCKET_NAME="tidbcloud-sli-data-${ENV}-${LOGICAL_REGION_NAME}"
FIREHOSE_STREAM_NAME="tidbcloud-sli-delivery-stream-${ENV}-${LOGICAL_REGION_NAME}"
FIREHOSE_ROLE_NAME="tidbcloud-sli-firehose-delivery-role-${ENV}-${LOGICAL_REGION_NAME}"
FIREHOSE_POLICY_NAME="tidbcloud-sli-firehose-delivery-policy-${ENV}-${LOGICAL_REGION_NAME}"
FIREHOSE_LOG_GROUP_NAME="/aws/kinesisfirehose/${FIREHOSE_STREAM_NAME}"
FIREHOSE_LOG_STREAM_NAME="DestinationDelivery"
GLUE_DATABASE_NAME="${BUCKET_NAME}"
GLUE_TABLE_NAME="${BUCKET_NAME}"

# ==============================================================================
# 1. Create S3 Bucket for SLI Data
# ==============================================================================
echo "Creating S3 bucket: ${BUCKET_NAME}..."
if ! aws --profile ${PROFILE} s3api head-bucket --bucket "${BUCKET_NAME}" &>/dev/null; then
    if [[ "${REGION}" == "us-east-1" ]]; then
        # us-east-1 doesn't need LocationConstraint and will fail if provided
        aws --profile ${PROFILE} s3api create-bucket --bucket "${BUCKET_NAME}" --region "${REGION}" >/dev/null
    else
        # All other regions require LocationConstraint
        aws --profile ${PROFILE} s3api create-bucket \
          --bucket "${BUCKET_NAME}" \
          --region "${REGION}" \
          --create-bucket-configuration LocationConstraint="${REGION}" >/dev/null
    fi
    aws --profile ${PROFILE} s3api put-bucket-tagging --bucket "${BUCKET_NAME}" --tagging "TagSet=[{Key=tidbcloud-auto-deploy-owner,Value=${RESOURCE_OWNER_TAG}}]" >/dev/null
    echo "S3 bucket created."
else
    echo "S3 bucket ${BUCKET_NAME} already exists."
fi

# Set up lifecycle policy for 2-year data retention
echo "Setting up S3 lifecycle policy for 2-year data retention..."
if ! aws --profile ${PROFILE} s3api get-bucket-lifecycle-configuration --bucket "${BUCKET_NAME}" --region "${REGION}" &>/dev/null; then
    LIFECYCLE_POLICY=$(cat <<EOF
{
    "Rules": [
        {
            "ID": "SLI-Data-Retention-2Years",
            "Status": "Enabled",
            "Filter": {},
            "Expiration": {
                "Days": 730
            },
            "AbortIncompleteMultipartUpload": {
                "DaysAfterInitiation": 7
            }
        }
    ]
}
EOF
)
    aws --profile ${PROFILE} s3api put-bucket-lifecycle-configuration --bucket "${BUCKET_NAME}" --lifecycle-configuration "${LIFECYCLE_POLICY}" >/dev/null
    echo "S3 lifecycle policy configured for 2-year retention."
else
    echo "S3 lifecycle policy already exists for bucket ${BUCKET_NAME}."
fi

echo "---"

# ==============================================================================
# 2. Create AWS Glue Database (Required by Firehose for Parquet conversion)
# ==============================================================================
echo "Creating AWS Glue Database: ${GLUE_DATABASE_NAME}..."
aws --profile ${PROFILE} glue create-database --database-input "{\"Name\":\"${GLUE_DATABASE_NAME}\"}" --region "${REGION}" \
 --tags "tidbcloud-auto-deploy-owner=${RESOURCE_OWNER_TAG}" \
 || echo "Glue database '${GLUE_DATABASE_NAME}' may already exist. Continuing."
echo "Glue Database step complete."
echo "---"

# ==============================================================================
# 3. Create AWS Glue Table
# ==============================================================================
echo "Creating AWS Glue Table: ${GLUE_TABLE_NAME}..."
GLUE_TABLE_INPUT=$(cat <<EOF
{
    "Name": "${GLUE_TABLE_NAME}",
    "StorageDescriptor": {
        "Columns": [
            {
                "Name": "timestamp",
                "Type": "bigint"
            },
            {
                "Name": "instance",
                "Type": "string"
            },
            {
                "Name": "metrics",
                "Type": "string"
            },
            {
                "Name": "signature",
                "Type": "string"
            }
        ],
        "Location": "s3://${BUCKET_NAME}/",
        "InputFormat": "org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat",
        "OutputFormat": "org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat",
        "SerdeInfo": {
            "SerializationLibrary": "org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe",
            "Parameters": {
                "serialization.format": "1"
            }
        }
    },
    "PartitionKeys": [],
    "TableType": "EXTERNAL_TABLE"
}
EOF
)
aws --profile ${PROFILE} glue create-table --database-name "${GLUE_DATABASE_NAME}" --table-input "${GLUE_TABLE_INPUT}" --region "${REGION}" || echo "Glue table '${GLUE_TABLE_NAME}' may already exist. Continuing."
echo "Glue Table step complete."
echo "---"

# ==============================================================================
# 4. Create CloudWatch Log Group for Firehose
# ==============================================================================
echo "Creating CloudWatch Log Group for Firehose: ${FIREHOSE_LOG_GROUP_NAME}..."
aws --profile ${PROFILE} logs create-log-group --log-group-name "${FIREHOSE_LOG_GROUP_NAME}" --region "${REGION}" \
  --tags "tidbcloud-auto-deploy-owner=${RESOURCE_OWNER_TAG}" \
  || echo "CloudWatch Log Group '${FIREHOSE_LOG_GROUP_NAME}' may already exist. Continuing."

# Set log group retention to 180 days
echo "Setting CloudWatch Log Group retention to 180 days..."
aws --profile ${PROFILE} logs put-retention-policy --log-group-name "${FIREHOSE_LOG_GROUP_NAME}" --retention-in-days 180 --region "${REGION}" >/dev/null
echo "CloudWatch Log Group retention policy set to 180 days."

echo "Create CloudWatch Log Stream for Firehose: ${FIREHOSE_LOG_STREAM_NAME}..."
aws --profile ${PROFILE} logs create-log-stream --log-group-name "${FIREHOSE_LOG_GROUP_NAME}" --log-stream-name "${FIREHOSE_LOG_STREAM_NAME}" --region "${REGION}" || echo "CloudWatch Log Stream '${FIREHOSE_LOG_STREAM_NAME}' may already exist. Continuing."

echo "CloudWatch Log Group step complete."
echo "---"

# ==============================================================================
# 5. Create Kinesis Data Firehose Delivery Stream
# ==============================================================================
echo "Creating Kinesis Data Firehose stream: ${FIREHOSE_STREAM_NAME}..."
if ! aws --profile ${PROFILE} firehose describe-delivery-stream --delivery-stream-name "${FIREHOSE_STREAM_NAME}" --region "${REGION}" &>/dev/null; then
    # Using a heredoc for readability
    FIREHOSE_CONFIG=$(cat <<EOF
{
    "DeliveryStreamName": "${FIREHOSE_STREAM_NAME}",
    "DeliveryStreamType": "DirectPut",
    "ExtendedS3DestinationConfiguration": {
        "RoleARN": "${ROLE_ARN_FH}",
        "BucketARN": "arn:aws:s3:::${BUCKET_NAME}",
        "BufferingHints": {
            "IntervalInSeconds": 60,
            "SizeInMBs": 64
        },
        "CloudWatchLoggingOptions": {
            "Enabled": true,
            "LogGroupName": "${FIREHOSE_LOG_GROUP_NAME}",
            "LogStreamName": "${FIREHOSE_LOG_STREAM_NAME}"
        },
        "CompressionFormat": "UNCOMPRESSED",
        "DataFormatConversionConfiguration": {
            "Enabled": true,
            "InputFormatConfiguration": {
                "Deserializer": {
                    "OpenXJsonSerDe": {}
                }
            },
            "OutputFormatConfiguration": {
                "Serializer": {
                    "ParquetSerDe": {}
                }
            },
            "SchemaConfiguration": {
                "DatabaseName": "${GLUE_DATABASE_NAME}",
                "Region": "${REGION}",
                "RoleARN": "${ROLE_ARN_FH}",
                "TableName": "${GLUE_TABLE_NAME}",
                "VersionId": "LATEST"
            }
        },
        "DynamicPartitioningConfiguration": {
            "Enabled": true,
            "RetryOptions": {
                "DurationInSeconds": 300
            }
        },
        "ProcessingConfiguration": {
            "Enabled": true,
            "Processors": [
                {
                    "Type": "MetadataExtraction",
                    "Parameters": [
                        {
                            "ParameterName": "JsonParsingEngine",
                            "ParameterValue": "JQ-1.6"
                        },
                        {
                            "ParameterName": "MetadataExtractionQuery",
                            "ParameterValue": "{metrics: .metrics.name, year: (.timestamp | strftime(\"%Y\")), month: (.timestamp | strftime(\"%m\")), day: (.timestamp | strftime(\"%d\")), hour: (.timestamp | strftime(\"%H\")), minute: (.timestamp | strftime(\"%M\"))}"
                        }
                    ]
                }
            ]
        },
        "Prefix": "metrics=!{partitionKeyFromQuery:metrics}/year=!{partitionKeyFromQuery:year}/month=!{partitionKeyFromQuery:month}/day=!{partitionKeyFromQuery:day}/hour=!{partitionKeyFromQuery:hour}/minute=!{partitionKeyFromQuery:minute}/",
        "ErrorOutputPrefix": "errors/year=!{timestamp:yyyy}/month=!{timestamp:MM}/day=!{timestamp:dd}/hour=!{timestamp:HH}/!{firehose:error-output-type}"
    }
}
EOF
)
    aws --profile ${PROFILE} firehose create-delivery-stream --region "${REGION}" --cli-input-json "${FIREHOSE_CONFIG}" \
      --tags "Key=tidbcloud-auto-deploy-owner,Value=${RESOURCE_OWNER_TAG}" >/dev/null
    echo "Kinesis Data Firehose stream created successfully."
else
    echo "Firehose stream ${FIREHOSE_STREAM_NAME} already exists."
fi
echo "---"

remote_output=$(yq '.outputPath' "${local_config}")
local_output=/tmp/output.yaml
cat <<EOF > "${local_output}"
bucket_name: ${BUCKET_NAME}
EOF
echo "Write output to S3..."
aws s3 cp "${local_output}" "${remote_output}" --region=$output_bucket_region