AWS Infra

from aws_cdk import (
    Stack,
    RemovalPolicy,
    Duration,
    CfnOutput,
    aws_ec2 as ec2,
    aws_iam as iam,
    aws_rds as rds,
    aws_s3 as s3,
    aws_ecr as ecr,
    aws_autoscaling as autoscaling,
    aws_elasticloadbalancingv2 as elbv2,
    aws_secretsmanager as secretsmanager,
    aws_logs as logs,
    # aws_grafana as grafana, # Grafana not used in the final setup, commented out
    aws_cloudfront as cloudfront,
    aws_cloudfront_origins as origins,
    aws_ssm as ssm
)
from constructs import Construct

class CalorieMitraStack(Stack):
    def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)

        # --- Get context variables ---
        app_name = self.node.try_get_context("app_name") or "CalorieMitra"
        app_port = self.node.try_get_context("app_port") or 8000
        container_port = 8000 # Set based on your Dockerfile EXPOSE
        docker_image_tag_context = self.node.try_get_context("docker_image_tag") or "latest"
        app_secret_arn = self.node.try_get_context("app_secret_arn")
        acm_certificate_arn = self.node.try_get_context("acm_certificate_arn")

        if not app_secret_arn:
             raise ValueError("Missing required secret ARN in CDK context (app_secret_arn)")
        if not acm_certificate_arn:
             raise ValueError("Missing required ACM Certificate ARN in CDK context (acm_certificate_arn)") # Ensure ACM ARN is provided for HTTPS

        ssm_image_tag_parameter_name = f"/app/{app_name.lower()}/image-tag"

        # --- Networking ---
        # VPC with Public and Isolated Subnets. Public subnets allow EC2 instances
        # to access the internet (and AWS public endpoints) via the Internet Gateway.
        # Isolated subnets provide a secure location for the RDS database.
        vpc = ec2.Vpc(
            self, f"{app_name}Vpc",
            max_azs=2,
            cidr="10.10.0.0/16",
            nat_gateways=0, # No NAT Gateways needed as EC2 is in Public Subnets
            subnet_configuration=[
                ec2.SubnetConfiguration(name="public", subnet_type=ec2.SubnetType.PUBLIC, cidr_mask=24),
                # Private subnet is ISOLATED for RDS
                ec2.SubnetConfiguration(name="private", subnet_type=ec2.SubnetType.PRIVATE_ISOLATED, cidr_mask=24),
            ]
        )

        # --- VPC Endpoints (Cost Optimization Applied) ---
        # Keep ONLY the S3 Gateway Endpoint. It's free and keeps S3 traffic
        # within the AWS network backbone, enhancing security and potentially performance.
        s3_gateway_endpoint = vpc.add_gateway_endpoint(
            "S3GatewayEndpoint",
            service=ec2.GatewayVpcEndpointAwsService.S3
        )

        # REMOVED: Interface Endpoint Security Group. Not needed as Interface Endpoints are removed.
        # vpc_endpoint_sg = ec2.SecurityGroup(...)

        # REMOVED: All Interface Endpoints (ECR, Secrets Manager, SSM, Logs, EC2/SSM Messages).
        # These were causing high costs. Instances in Public Subnets will now access these
        # services via their public endpoints through the Internet Gateway.
        # vpc.add_interface_endpoint("EcrDockerEndpoint", ...)
        # vpc.add_interface_endpoint("EcrApiEndpoint", ...)
        # vpc.add_interface_endpoint("SecretsManagerEndpoint", ...)
        # vpc.add_interface_endpoint("SsmEndpoint", ...)
        # vpc.add_interface_endpoint("LogsEndpoint", ...)
        # vpc.add_interface_endpoint("Ec2MessagesEndpoint", ...)
        # vpc.add_interface_endpoint("SsmMessagesEndpoint", ...)


        # --- Security Groups ---
        # ALB Security Group: Allows internet traffic on HTTP/HTTPS.
        alb_sg = ec2.SecurityGroup(
            self, f"{app_name}AlbSg",
            vpc=vpc,
            description="ALB SG allowing HTTP/HTTPS ingress",
            allow_all_outbound=True # Standard practice for ALBs
        )
        alb_sg.add_ingress_rule(ec2.Peer.any_ipv4(), ec2.Port.tcp(80), "Allow HTTP for redirect")
        alb_sg.add_ingress_rule(ec2.Peer.any_ipv4(), ec2.Port.tcp(443), "Allow HTTPS from internet")

        # EC2 Instance Security Group: Critical for security.
        ec2_sg = ec2.SecurityGroup(
            self, f"{app_name}Ec2Sg",
            vpc=vpc,
            description="EC2 Instance SG allowing traffic only from ALB and required outbound",
            # IMPORTANT: Allow all outbound is needed for instances in Public Subnets
            # to reach AWS public endpoints (ECR, SSM, etc.) and OS update repos via the IGW.
            # While less strict than fully private, it's necessary for this cost-optimized path.
            allow_all_outbound=True
        )
        # Allow traffic ONLY from the ALB on the application port. Prevents direct access to instances.
        ec2_sg.add_ingress_rule(
            alb_sg, ec2.Port.tcp(app_port),
            f"Allow traffic from ALB on port {app_port}"
        )
        # REMOVED: Ingress rule for SSM via HTTPS. Not needed as SSM agent uses outbound connection.

        # RDS Security Group: Isolates the database.
        rds_sg = ec2.SecurityGroup(
            self, f"{app_name}RdsSg",
            vpc=vpc,
            description="RDS SG allowing traffic only from EC2 SG",
            # Allow outbound is less critical here, but True is acceptable default.
            # Could be restricted further if desired, but usually not necessary.
            allow_all_outbound=True
        )
        # Allow traffic only from application EC2 instances on the standard PostgreSQL port.
        rds_sg.add_ingress_rule(
            ec2_sg, ec2.Port.tcp(5432), # Standard PostgreSQL port
            "Allow DB traffic from EC2"
        )

        # REMOVED: Rule allowing EC2 to connect to VPC Interface Endpoints SG. No longer needed.
        # vpc_endpoint_sg.add_ingress_rule(ec2_sg, ec2.Port.tcp(443), "Allow EC2 instances to connect to Endpoints")


        # --- ECR Repository ---
        # Securely stores Docker images. IAM permissions control access.
        ecr_repository = ecr.Repository(
            self, f"{app_name}EcrRepo", repository_name=f"{app_name.lower()}-app-repo",
            removal_policy=RemovalPolicy.RETAIN, # Keep repo even if stack is deleted
            image_tag_mutability=ecr.TagMutability.MUTABLE # Allows overwriting tags like 'latest'
        )

        # --- S3 Bucket ---
        # Securely stores meal images. Access restricted via CloudFront OAC and IAM.
        meal_images_bucket = s3.Bucket(
            self, f"{app_name}MealImagesBucket",
            block_public_access=s3.BlockPublicAccess.BLOCK_ALL, # Essential for security
            encryption=s3.BucketEncryption.S3_MANAGED, # Server-side encryption
            removal_policy=RemovalPolicy.RETAIN, # Keep bucket even if stack is deleted
            auto_delete_objects=False, # Safer default for production
            lifecycle_rules=[s3.LifecycleRule(id="MoveToIA", enabled=True, transitions=[
                s3.Transition(storage_class=s3.StorageClass.INFREQUENT_ACCESS, transition_after=Duration.days(90))])] # Cost optimization
        )

        # --- Cloudfront Distribution with Origin Access Control (OAC) ---
        # OAC is the recommended secure way for CloudFront to access S3 privately.
        cfn_oac = cloudfront.CfnOriginAccessControl(self, f"{app_name}OAC",
            origin_access_control_config=cloudfront.CfnOriginAccessControl.OriginAccessControlConfigProperty(
                name=f"{app_name}-OAC-{self.region}", origin_access_control_origin_type="s3",
                signing_behavior="always", signing_protocol="sigv4", description="OAC for S3 Bucket"))

        cloudfront_distribution = cloudfront.Distribution(
            self, f"{app_name}CloudfrontDistribution",
            default_behavior=cloudfront.BehaviorOptions(
                origin=origins.S3Origin(meal_images_bucket), # Using the secure S3 origin
                viewer_protocol_policy=cloudfront.ViewerProtocolPolicy.REDIRECT_TO_HTTPS, # Enforce HTTPS
                cache_policy=cloudfront.CachePolicy.CACHING_OPTIMIZED), # Standard caching policy
            comment=f"{app_name} CDN for Meal Images")

        # --- Add OAC configuration to the CloudFront Distribution ---
        # This overrides the default settings to ensure OAC is used.
        cfn_distribution = cloudfront_distribution.node.default_child
        cfn_distribution.add_property_override("DistributionConfig.Origins.0.OriginAccessControlId", cfn_oac.attr_id)
        # Remove the legacy Origin Access Identity if it exists (OAC replaces OAI)
        cfn_distribution.add_property_override("DistributionConfig.Origins.0.S3OriginConfig.OriginAccessIdentity", "")

        # --- Update S3 Bucket Policy to allow CloudFront access via OAC ---
        # This policy grants CloudFront permission to GetObject, restricted by the distribution ARN.
        meal_images_bucket.add_to_resource_policy(
            iam.PolicyStatement(
                actions=["s3:GetObject"],
                resources=[meal_images_bucket.arn_for_objects("*")],
                principals=[iam.ServicePrincipal("cloudfront.amazonaws.com")],
                # Condition ensures only THIS CloudFront distribution can access the bucket
                conditions={"StringEquals": {
                    "AWS:SourceArn": f"arn:aws:cloudfront::{self.account}:distribution/{cloudfront_distribution.distribution_id}"
                }}
            )
        )

        # --- RDS PostgreSQL Instance ---
        # Securely placed in PRIVATE_ISOLATED subnets.
        db_instance_name = "defaultdb"
        db_instance = rds.DatabaseInstance(
            self, f"{app_name}RdsInstance",
            engine=rds.DatabaseInstanceEngine.postgres(version=rds.PostgresEngineVersion.VER_15),
            instance_type=ec2.InstanceType.of(ec2.InstanceClass.BURSTABLE4_GRAVITON, ec2.InstanceSize.MEDIUM),
            vpc=vpc,
            # Place RDS in the PRIVATE_ISOLATED subnets for maximum network security
            vpc_subnets=ec2.SubnetSelection(subnet_type=ec2.SubnetType.PRIVATE_ISOLATED),
            security_groups=[rds_sg], # Apply the restrictive RDS Security Group
            allocated_storage=50,
            storage_type=rds.StorageType.GP3, # Cost-effective GP3 storage
            database_name=db_instance_name,
            multi_az=False, # Set to True for production HA (higher cost)
            backup_retention=Duration.days(7), # Adjust as needed
            delete_automated_backups=False, # Keep backups if instance deleted (safer)
            deletion_protection=True, # Prevent accidental deletion (recommended for prod)
            copy_tags_to_snapshot=True,
            enable_performance_insights=True, # Useful for debugging performance
            removal_policy=RemovalPolicy.RETAIN # Keep DB even if stack is deleted
        )

        # --- SSM Parameter for Docker Image Tag ---
        # Securely stores the deployment tag, accessible via IAM role.
        ssm_image_tag_parameter = ssm.StringParameter(
            self, f"{app_name}ImageTagParam",
            parameter_name=ssm_image_tag_parameter_name,
            string_value=docker_image_tag_context,
            description=f"Docker image tag for {app_name}",
            tier=ssm.ParameterTier.STANDARD
        )

        # --- IAM Role for EC2 Instances ---
        # Defines permissions for EC2 instances, following least privilege principle.
        ec2_role = iam.Role(
            self, f"{app_name}Ec2Role",
            assumed_by=iam.ServicePrincipal("ec2.amazonaws.com"),
            managed_policies=[
                # Required for SSM Agent to function (used for management, patching, etc.)
                iam.ManagedPolicy.from_aws_managed_policy_name("AmazonSSMManagedInstanceCore"),
                # Required for CloudWatch Agent to send logs/metrics
                iam.ManagedPolicy.from_aws_managed_policy_name("CloudWatchAgentServerPolicy"),
                # Required to pull images from ECR
                iam.ManagedPolicy.from_aws_managed_policy_name("AmazonEC2ContainerRegistryReadOnly")
            ]
        )

        # Grant specific permissions needed by the application and startup script.
        # Log permissions (scoped to specific log groups)
        ec2_role.add_to_policy(iam.PolicyStatement(
            effect=iam.Effect.ALLOW,
            actions=[
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents",
                "logs:DescribeLogStreams",
                "logs:DescribeLogGroups"
            ],
            resources=[
                f"arn:aws:logs:{self.region}:{self.account}:log-group:/app/{app_name.lower()}/*",
                f"arn:aws:logs:{self.region}:{self.account}:log-group:/app/{app_name.lower()}/*:*",
                f"arn:aws:logs:{self.region}:{self.account}:log-group:/var/log/user-data*",
                f"arn:aws:logs:{self.region}:{self.account}:log-group:/var/log/user-data*:*"
            ]
        ))
        # SSM Parameter Read permission (scoped to the specific parameter)
        ec2_role.add_to_policy(iam.PolicyStatement(
            effect=iam.Effect.ALLOW,
            actions=["ssm:GetParameter", "ssm:GetParameters", "ssm:GetParametersByPath"],
            resources=[ssm_image_tag_parameter.parameter_arn] # Use parameter ARN for least privilege
        ))
        # Autoscaling permissions (needed for UserData error handling to set instance unhealthy)
        ec2_role.add_to_policy(iam.PolicyStatement(
            effect=iam.Effect.ALLOW,
            actions=["autoscaling:SetInstanceHealth", "autoscaling:DescribeAutoScalingInstances"],
            # Scoping this further requires knowing the ASG ARN or using conditions,
            # "*" is often acceptable here for simplicity.
            resources=["*"]
        ))

        # Grant resource-based permissions (simpler way to grant access)
        ecr_repository.grant_pull(ec2_role) # Allow role to pull from this ECR repo
        meal_images_bucket.grant_read_write(ec2_role) # Allow role to R/W to this S3 bucket
        app_secret = secretsmanager.Secret.from_secret_complete_arn(self, "AppSecret", app_secret_arn)
        app_secret.grant_read(ec2_role) # Allow role to read the specified secret
        # Note: ssm_image_tag_parameter.grant_read(ec2_role) could also be used instead of the policy above

        # --- CloudWatch Log Group for Application Logs ---
        # Centralized logging for the Docker container via awslogs driver.
        app_log_group = logs.LogGroup(
            self, f"{app_name}AppLogGroup",
            log_group_name=f"/app/{app_name.lower()}/docker-logs",
            retention=logs.RetentionDays.TWO_WEEKS, # Adjust retention as needed
            removal_policy=RemovalPolicy.DESTROY # OK to destroy log group if stack deleted
        )
        # Granting write is implicitly handled by CloudWatchAgentServerPolicy and docker config,
        # but explicit grant doesn't hurt.
        app_log_group.grant_write(ec2_role)

        # --- User Data Script ---
        # This script runs on instance launch to set up the environment and start the application.
        # No changes needed here due to endpoint removal, as AWS CLI/SDK calls will
        # automatically use the IGW path from the public subnet.
        app_log_group_name = app_log_group.log_group_name
        custom_ami_id = "ami-04c5ae52f42a72334"

        user_data_script = f"""#!/bin/bash -xeu
# -x: print commands, -e: exit on error, -u: treat unset variables as error
set -o pipefail

# Redirect stdout/stderr to a log file and to the system logger
exec > >(tee /var/log/user-data.log | logger -t user-data -s 2>/dev/console) 2>&1

# --- Basic Logging Function (Output to stderr) ---
log() {{
    # Use >&2 to redirect echo output to standard error
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" >&2
}}

log "--- UserData Script Start (Using Custom AMI: {custom_ami_id}) ---"

# --- Essential Variables (Injected by CDK f-string) ---
APP_PORT="{app_port}"
CONTAINER_PORT="{container_port}"
AWS_REGION="{self.region}"
AWS_ACCOUNT_ID="{self.account}"
APP_SECRET_ARN="{app_secret_arn}"
SSM_PARAM_NAME="{ssm_image_tag_parameter_name}"
ECR_REPOSITORY_NAME="{ecr_repository.repository_name}"
DB_ENDPOINT_ADDRESS="{db_instance.db_instance_endpoint_address}"
DB_ENDPOINT_PORT="{db_instance.db_instance_endpoint_port}"
DB_NAME="{db_instance_name}"
S3_BUCKET="{meal_images_bucket.bucket_name}"
APP_LOG_GROUP_NAME="{app_log_group_name}" # Use the variable passed from CDK
# -----------------------------------------------------

# --- Error Handling Setup ---
handle_error() {{
    local error_msg="$1"
    local exit_code="${{2:-1}}"
    log "ERROR: $error_msg (Exit Code: $exit_code)"
    TOKEN=$(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 60" || true)
    INSTANCE_ID=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" http://169.254.169.254/latest/meta-data/instance-id || curl -s http://169.254.169.254/latest/meta-data/instance-id || echo "unknown")
    if [[ "$INSTANCE_ID" != "unknown" ]] && command -v aws &> /dev/null; then
        log "Attempting to signal ASG Unhealthy for instance $INSTANCE_ID in region $AWS_REGION..."
        aws autoscaling set-instance-health --instance-id "$INSTANCE_ID" --health-status Unhealthy --region "$AWS_REGION" || \\
            log "Warning: Failed to set instance health via AWS CLI. Check IAM permissions (autoscaling:SetInstanceHealth)."
    else
        log "Warning: Could not get Instance ID or AWS CLI not available. Cannot signal ASG."
    fi
    exit $exit_code
}}
trap 'handle_error "Script failed on line $LINENO" $?' ERR

# --- Verify Prerequisites (Should be installed in AMI) ---
log "Verifying prerequisite tools (should be installed in AMI)..."
command -v docker >/dev/null || handle_error "Docker command not found. AMI build might have failed."
command -v aws >/dev/null || handle_error "AWS CLI command not found. AMI build might have failed."
command -v jq >/dev/null || handle_error "jq command not found. AMI build might have failed."
# command -v git >/dev/null || handle_error "git command not found. AMI build might have failed." # Check if needed
command -v nc >/dev/null || handle_error "netcat (nc) command not found. AMI build might have failed."
log "Prerequisite tools verified."

# --- Configure and Start Docker ---
log "Configuring Docker log driver (awslogs)..."
sudo mkdir -p /etc/docker # Ensure directory exists
sudo tee /etc/docker/daemon.json > /dev/null <<EOF
{{
  "log-driver": "awslogs",
  "log-opts": {{
    "awslogs-region": "$AWS_REGION",
    "awslogs-group": "$APP_LOG_GROUP_NAME",
    "awslogs-create-group": "true",
    "tag": "{{{{.Name}}}}/{{{{.ID}}}}"
  }}
}}
EOF
log "Docker daemon.json configured for awslogs."

log "Ensuring Docker service is enabled and restarting it..."
# Docker service should be enabled by the AMI build process
# Restart is needed to apply the daemon.json configuration
sudo systemctl restart docker
sleep 5 # Give docker a moment to restart
sudo systemctl is-active --quiet docker || handle_error "Docker service failed to restart after configuration. Check /etc/docker/daemon.json syntax and system logs (journalctl -u docker.service)."
# Add ubuntu user to docker group if not already added (useful for potential debugging)
sudo usermod -a -G docker ubuntu || log "Warning: Failed to add 'ubuntu' user to docker group (might already exist)."
log "Docker service configured and running."


# --- Fetch Secrets Function (with retries) ---
function get_secret {{
    local secret_arn="$1"
    local retries=3
    local count=0
    local delay=5
    local secret_json=""

    log "Attempting to fetch secret: $secret_arn (Retry logic active)"
    while [ $count -lt $retries ]; do
        # AWS CLI call will use IAM role and network path (IGW in this case)
        secret_json=$(aws secretsmanager get-secret-value --secret-id "$secret_arn" --region "$AWS_REGION" --query SecretString --output text 2>/var/log/awscli_secret_error.log)
        local aws_exit_code=$?
        if [ $aws_exit_code -eq 0 ] && [ -n "$secret_json" ] && [ "$secret_json" != "null" ]; then
            # Validate JSON structure
            if echo "$secret_json" | jq -e . > /dev/null; then
                log "Successfully fetched and validated secret JSON (attempt $((count+1))/$retries)."
                rm -f /var/log/awscli_secret_error.log
                echo "$secret_json" # Output ONLY the JSON
                return 0
            else
                log "Warning: Fetched secret for $secret_arn is not valid JSON (attempt $((count+1))/$retries). Retrying in $delay seconds..."
            fi
        else
            log "Warning: Failed to fetch secret $secret_arn (AWS CLI exit code: $aws_exit_code, attempt $((count+1))/$retries). Check /var/log/awscli_secret_error.log. Retrying in $delay seconds..."
        fi
        count=$((count+1))
        sleep $delay
    done
    log "ERROR: Failed to fetch valid secret JSON for $secret_arn after $retries attempts."
    cat /var/log/awscli_secret_error.log # Log the actual error
    return 1
}}

log "Fetching App secrets from ARN: $APP_SECRET_ARN"
APP_SECRET_JSON=$(get_secret "$APP_SECRET_ARN") || handle_error "Failed to fetch App secrets from $APP_SECRET_ARN after retries" $?
# Avoid logging secrets to UserData logs if possible
log "Successfully fetched secrets."

# Parse secrets using jq (ensure required fields are present)
log "Parsing essential secrets..."
GEMINI_API_KEY=$(echo "$APP_SECRET_JSON" | jq -r '.GEMINI_API_KEY // ""')
ADMIN_EMAILS=$(echo "$APP_SECRET_JSON" | jq -r '.ADMIN_EMAILS // ""')
RATE_LIMIT_MAX=$(echo "$APP_SECRET_JSON" | jq -r '.RATE_LIMIT_MAX // ""')
ENABLE_SWAGGER=$(echo "$APP_SECRET_JSON" | jq -r '.ENABLE_SWAGGER // ""')
CLOUDFRONT_DOMAIN_NAME=$(echo "$APP_SECRET_JSON" | jq -r '.CLOUDFRONT_DOMAIN_NAME // ""')
ENDPOINT=$(echo "$APP_SECRET_JSON" | jq -r '.ENDPOINT // ""')
FIREBASE_PROJECT_ID=$(echo "$APP_SECRET_JSON" | jq -r '.FIREBASE_PROJECT_ID // ""')
FIREBASE_CLIENT_EMAIL=$(echo "$APP_SECRET_JSON" | jq -r '.FIREBASE_CLIENT_EMAIL // ""')
FIREBASE_PRIVATE_KEY=$(echo "$APP_SECRET_JSON" | jq -r '.FIREBASE_PRIVATE_KEY // ""')
DB_USERNAME=$(echo "$APP_SECRET_JSON" | jq -r '.DB_USERNAME // ""')
DB_PASSWORD=$(echo "$APP_SECRET_JSON" | jq -r '.DB_PASSWORD // ""')

# Validate essential secrets that are critical for startup
if [ -z "$DB_USERNAME" ]; then handle_error "DB_USERNAME missing from secrets"; fi
if [ -z "$DB_PASSWORD" ]; then handle_error "DB_PASSWORD missing from secrets"; fi
# Add checks for other critical secrets if needed
log "Secrets parsed."

# --- Fetch Docker Image Tag from SSM ---
log "Fetching image tag from SSM parameter: $SSM_PARAM_NAME"
# AWS CLI call uses IAM role and network path (IGW)
TARGET_IMAGE_TAG=$(aws ssm get-parameter --name "$SSM_PARAM_NAME" --query "Parameter.Value" --output text --region "$AWS_REGION" 2>/var/log/awscli_ssm_error.log)
ssm_exit_code=$?
if [ $ssm_exit_code -ne 0 ] || [ -z "$TARGET_IMAGE_TAG" ] || [ "$TARGET_IMAGE_TAG" == "None" ] || [ "$TARGET_IMAGE_TAG" == "null" ]; then
    log "Warning: Failed to get tag from SSM $SSM_PARAM_NAME (Exit code: $ssm_exit_code). Check /var/log/awscli_ssm_error.log. Using 'latest' as fallback."
    cat /var/log/awscli_ssm_error.log # Log the actual error
    TARGET_IMAGE_TAG="latest" # Define a safe fallback
fi
rm -f /var/log/awscli_ssm_error.log
log "Using image tag: $TARGET_IMAGE_TAG"

# --- ECR Login ---
ECR_LOGIN_SERVER="$AWS_ACCOUNT_ID.dkr.ecr.$AWS_REGION.amazonaws.com"
log "Logging into ECR: $ECR_LOGIN_SERVER"
# AWS CLI call uses IAM role and network path (IGW)
aws ecr get-login-password --region "$AWS_REGION" | sudo docker login --username AWS --password-stdin "$ECR_LOGIN_SERVER" || \\
    handle_error "Failed to log into ECR server $ECR_LOGIN_SERVER. Check IAM permissions (ecr:GetAuthorizationToken) and network connectivity."
log "ECR login successful."

# --- Prepare and Run Container ---
ECR_REPO_URI="$ECR_LOGIN_SERVER/$ECR_REPOSITORY_NAME"
IMAGE_URI="$ECR_REPO_URI:$TARGET_IMAGE_TAG"
CONTAINER_NAME="myapp" # Use a variable for the container name

log "Stopping and removing existing container named '$CONTAINER_NAME' (if any)..."
sudo docker stop "$CONTAINER_NAME" 2>/dev/null || true
sudo docker rm -f "$CONTAINER_NAME" 2>/dev/null || true

log "Pulling image: $IMAGE_URI"
PULL_SUCCESS="false"
for i in {{1..3}}; do
    if sudo docker pull "$IMAGE_URI"; then
        log "Image pull successful (attempt $i)."
        PULL_SUCCESS="true"
        break
    else
        log "Warning: Pull attempt $i failed for $IMAGE_URI. Retrying in 5 seconds..."
        sleep 5;
    fi
done
if [ "$PULL_SUCCESS" != "true" ]; then
    handle_error "FATAL: Failed to pull image $IMAGE_URI after 3 attempts. Check ECR repository, tag, network, and IAM permissions (ecr:BatchGetImage, ecr:GetDownloadUrlForLayer)."
fi
log "Image pull successful."

log "Running container '$CONTAINER_NAME' from image $IMAGE_URI..."
# Build DATABASE_URL securely (avoiding logging password if possible)
DB_URL_CREDENTIALS=""
if [ -n "$DB_USERNAME" ]; then
    # Note: Password is still visible in process list (`ps aux`) during run.
    # More secure methods involve reading from mounted files or container secrets management.
    DB_URL_CREDENTIALS="$DB_USERNAME:$DB_PASSWORD@"
fi
DATABASE_URL="postgresql://$DB_URL_CREDENTIALS$DB_ENDPOINT_ADDRESS:$DB_ENDPOINT_PORT/$DB_NAME"

# Build docker run options array (passing secrets as environment variables)
# Note: Environment variables can potentially be inspected. Consider alternatives for highly sensitive data.
declare -a docker_env_opts
docker_env_opts+=( "-e" "DATABASE_HOST=$DB_ENDPOINT_ADDRESS" )
docker_env_opts+=( "-e" "DATABASE_PORT=$DB_ENDPOINT_PORT" )
docker_env_opts+=( "-e" "DATABASE_NAME=$DB_NAME" )
docker_env_opts+=( "-e" "DATABASE_URL=$DATABASE_URL" ) # App might use this directly
docker_env_opts+=( "-e" "S3_BUCKET_NAME=$S3_BUCKET" )
docker_env_opts+=( "-e" "AWS_REGION=$AWS_REGION" ) # Useful for SDK inside container if needed
docker_env_opts+=( "-e" "NODE_ENV=production" ) # Example for Node.js apps
docker_env_opts+=( "-e" "PORT=$CONTAINER_PORT" ) # Port inside the container

# Conditionally add other secrets if they exist
[ -n "$DB_USERNAME" ] && docker_env_opts+=( "-e" "DATABASE_USERNAME=$DB_USERNAME" )
[ -n "$DB_PASSWORD" ] && docker_env_opts+=( "-e" "DATABASE_PASSWORD=$DB_PASSWORD" ) # Sensitive
[ -n "$GEMINI_API_KEY" ] && docker_env_opts+=( "-e" "GEMINI_API_KEY=$GEMINI_API_KEY" ) # Sensitive
[ -n "$ADMIN_EMAILS" ] && docker_env_opts+=( "-e" "ADMIN_EMAILS=$ADMIN_EMAILS" )
[ -n "$RATE_LIMIT_MAX" ] && docker_env_opts+=( "-e" "RATE_LIMIT_MAX=$RATE_LIMIT_MAX" )
[ -n "$ENABLE_SWAGGER" ] && docker_env_opts+=( "-e" "ENABLE_SWAGGER=$ENABLE_SWAGGER" )
[ -n "$CLOUDFRONT_DOMAIN_NAME" ] && docker_env_opts+=( "-e" "CLOUDFRONT_DOMAIN_NAME=$CLOUDFRONT_DOMAIN_NAME" )
[ -n "$ENDPOINT" ] && docker_env_opts+=( "-e" "ENDPOINT=$ENDPOINT" )
[ -n "$FIREBASE_PROJECT_ID" ] && docker_env_opts+=( "-e" "FIREBASE_PROJECT_ID=$FIREBASE_PROJECT_ID" )
[ -n "$FIREBASE_CLIENT_EMAIL" ] && docker_env_opts+=( "-e" "FIREBASE_CLIENT_EMAIL=$FIREBASE_CLIENT_EMAIL" )
[ -n "$FIREBASE_PRIVATE_KEY" ] && docker_env_opts+=( "-e" "FIREBASE_PRIVATE_KEY=$FIREBASE_PRIVATE_KEY" ) # Sensitive

# Execute docker run
# Log driver is set globally in daemon.json
sudo docker run -d --name "$CONTAINER_NAME" \\
    --restart always \\
    -p "$APP_PORT":"$CONTAINER_PORT" \\
    "${{docker_env_opts[@]}}" \\
    "$IMAGE_URI" || handle_error "Failed to execute 'docker run' command for image $IMAGE_URI."

# Brief pause and check if container started and is running
log "Waiting 10 seconds for container '$CONTAINER_NAME' to stabilize..."
sleep 10
if ! sudo docker ps -q --filter name="^/${{CONTAINER_NAME}}$"; then
    log "--- Docker Status After Failed Start ---"
    sudo docker ps -a || log "Warning: Failed to run 'docker ps -a'."
    log "--- Last 100 lines of '$CONTAINER_NAME' logs (attempting fetch) ---"
    # Attempt to get logs; might fail if container never started properly or logging driver issue
    sudo docker logs --tail 100 "$CONTAINER_NAME" || log "Could not get logs for '$CONTAINER_NAME' via docker logs command."
    handle_error "Container '$CONTAINER_NAME' is not running shortly after start. Check CloudWatch Logs ($APP_LOG_GROUP_NAME), system logs (/var/log/syslog or journalctl -u docker.service), and container logs (if fetchable)."
fi

log "Container '$CONTAINER_NAME' started successfully. Application should be initializing."
log "Logs should appear in CloudWatch Group: $APP_LOG_GROUP_NAME"
log "--- UserData Script Finished Successfully ---"
exit 0
"""

        # --- User Data Object ---
        user_data = ec2.UserData.custom(user_data_script)

        # --- Use the Custom AMI ---
        # Ensure the specified AMI exists in the target region and has necessary tools (Docker, AWS CLI, jq, nc)
        custom_calorie_mitra_ami = ec2.MachineImage.generic_linux({
            self.region: custom_ami_id
        })

        # --- EC2 Launch Template ---
        # Defines the configuration for instances launched by the Auto Scaling Group.
        launch_template = ec2.LaunchTemplate(
            self, f"{app_name}LaunchTemplate",
            instance_type=ec2.InstanceType.of(ec2.InstanceClass.BURSTABLE4_GRAVITON, ec2.InstanceSize.MEDIUM), # Graviton for cost/performance
            machine_image=custom_calorie_mitra_ami, # Use the prepared custom AMI
            role=ec2_role, # Attach the defined IAM role
            security_group=ec2_sg, # Attach the EC2 Security Group
            user_data=user_data, # Provide the startup script
            # Configure the root EBS volume
            block_devices=[ec2.BlockDevice(
                device_name="/dev/sda1", # Verify this is the correct root device name for your AMI
                volume=ec2.BlockDeviceVolume.ebs(
                    volume_size=20, # Adjust size as needed
                    volume_type=ec2.EbsDeviceVolumeType.GP3, # Cost-effective and performant
                    encrypted=True, # Encrypt EBS volume for security at rest
                    delete_on_termination=True # Ensure volume is deleted with instance
                )
            )],
            # Enable detailed CloudWatch monitoring for instances (optional, incurs cost)
            # detailed_monitoring=True,
            # Ensure IMDSv2 is required for better security (metadata service access)
            http_tokens=ec2.LaunchTemplateHttpTokens.REQUIRED,
            http_put_response_hop_limit=2 # Standard setting when requiring IMDSv2
        )

        # --- Define Auto Scaling Group ---
        # Manages the EC2 instances, ensuring desired capacity and scaling.
        asg = autoscaling.AutoScalingGroup(
            self, f"{app_name}Asg",
            vpc=vpc,
            launch_template=launch_template, # Use the defined Launch Template
            # Launch instances into PUBLIC subnets to utilize the Internet Gateway
            vpc_subnets=ec2.SubnetSelection(subnet_type=ec2.SubnetType.PUBLIC),
            min_capacity=1,
            max_capacity=4, # Adjust max capacity based on expected load
            desired_capacity=1,
            # Use ELB health checks PLUS EC2 health checks for robustness
            health_check=autoscaling.HealthCheck.elb(grace=Duration.minutes(5)), # Shorter grace period maybe? Depends on app start time
        )
        # Simple CPU-based scaling policy
        asg.scale_on_cpu_utilization(
            f"{app_name}CpuScaling",
            target_utilization_percent=75, # Target 75% CPU before scaling out (adjust as needed)
            cooldown=Duration.minutes(5) # Cooldown period after scaling activity
        )

        # --- Application Load Balancer ---
        # Securely distributes incoming traffic across instances.
        alb = elbv2.ApplicationLoadBalancer(
            self, f"{app_name}Alb",
            vpc=vpc,
            internet_facing=True, # Receives traffic from the internet
            security_group=alb_sg, # Apply the ALB Security Group
            # Enable deletion protection for production ALBs
            deletion_protection=True,
            # Enable access logs for security auditing and analysis (requires an S3 bucket)
            # access_logs_bucket=s3.Bucket(self, f"{app_name}AlbLogsBucket", ...) # Define an S3 bucket for logs
        )
        # Enable ALB Access Logs (Optional but recommended for security)
        # You would need to create an S3 bucket first
        # alb_logs_bucket = s3.Bucket(self, f"{app_name}AlbLogsBucket", ...)
        # alb.log_access_logs(alb_logs_bucket, prefix=f"{app_name}-alb-logs")

        # --- Target Group ---
        # Defines how the ALB routes traffic to the registered instances.
        target_group = elbv2.ApplicationTargetGroup(
            self, f"{app_name}TargetGroup",
            vpc=vpc,
            port=app_port, # Port the application listens on inside the instance
            protocol=elbv2.ApplicationProtocol.HTTP, # ALB communicates with instances over HTTP
            target_type=elbv2.TargetType.INSTANCE,
            targets=[asg], # Automatically register instances from the ASG
            health_check=elbv2.HealthCheck(
                enabled=True,
                path="/healthz", # Specific health check endpoint in your application
                port=str(app_port), # Health check port must match target port
                protocol=elbv2.Protocol.HTTP,
                interval=Duration.seconds(60), # Check more frequently
                timeout=Duration.seconds(30), # Shorter timeout
                healthy_threshold_count=2,
                unhealthy_threshold_count=10, # Slightly more tolerant to transient issues
                healthy_http_codes="200-299" # Be specific about healthy response code
            ),
            # Adjust deregistration delay (time to drain connections before termination)
            deregistration_delay=Duration.seconds(60)
        )

        # --- ALB Listener ---
        # Handles incoming HTTPS traffic on port 443.
        certificate = elbv2.ListenerCertificate.from_arn(acm_certificate_arn)
        https_listener = alb.add_listener(f"{app_name}HttpsListener",
            port=443,
            protocol=elbv2.ApplicationProtocol.HTTPS,
            certificates=[certificate], # Attach the ACM SSL/TLS certificate
            # Set a secure default action (e.g., fixed response or redirect) if needed,
            # but usually default_target_groups is correct.
            default_target_groups=[target_group],
            # Recommended: Use a secure SSL policy from AWS
            ssl_policy=elbv2.SslPolicy.RECOMMENDED_TLS # Uses modern TLS versions/ciphers
        )

        # --- ALB HTTP to HTTPS Redirect ---
        # Automatically redirects any HTTP traffic to HTTPS for security.
        alb.add_redirect(
            source_protocol=elbv2.ApplicationProtocol.HTTP,
            source_port=80,
            target_protocol=elbv2.ApplicationProtocol.HTTPS,
            target_port=443,
            # redirect_status_code="HTTP_301" # Permanent redirect (default)
        )


        # --- Outputs ---
        # Provide useful information about the created resources.
        CfnOutput(self, "Region", value=self.region)
        CfnOutput(self, "VpcId", value=vpc.vpc_id)
        CfnOutput(self, "EcrRepositoryUri", value=ecr_repository.repository_uri)
        CfnOutput(self, "MealImagesBucketName", value=meal_images_bucket.bucket_name)
        CfnOutput(self, "CloudfrontDomainName", value=cloudfront_distribution.distribution_domain_name)
        CfnOutput(self, "LoadBalancerDns", value=alb.load_balancer_dns_name)
        CfnOutput(self, "LoadBalancerUrlHttps", value=f"https://{alb.load_balancer_dns_name}")
        CfnOutput(self, "RdsEndpointAddress", value=db_instance.db_instance_endpoint_address)
        CfnOutput(self, "RdsEndpointPort", value=db_instance.db_instance_endpoint_port)
        CfnOutput(self, "EC2RoleArn", value=ec2_role.role_arn)
        CfnOutput(self, "AppLogGroupName", value=app_log_group.log_group_name)
        CfnOutput(self, "SsmImageTagParameterName", value=ssm_image_tag_parameter.parameter_name)
        CfnOutput(self, "SelectedAmiId", value=custom_ami_id)
        CfnOutput(self, "PublicSubnetIds", value=", ".join(vpc.select_subnets(subnet_type=ec2.SubnetType.PUBLIC).subnet_ids))
        CfnOutput(self, "PrivateIsolatedSubnetIds", value=", ".join(vpc.select_subnets(subnet_type=ec2.SubnetType.PRIVATE_ISOLATED).subnet_ids))