#!/bin/bash
# ----------------------------------------------------
# Cirries - EC2 Docker Test (ECR + IAM Role, Ubuntu 22.04)
# ----------------------------------------------------
# Usage: ./ec2_final_v2.sh <collector_ip>
# ----------------------------------------------------

set -e
if [ $# -ne 1 ]; then
  echo "❌ Usage: $0 <collector_ip>"
  exit 1
fi

COLLECTOR_IP=$1
REGION="us-east-1"
#AMI_ID="ami-04a81a99f5ec58529"
AMI_ID="ami-053b0d53c279acc90"
INSTANCE_TYPE="t3.large"
KEY_PAIR="PMCPerformance"
KEY_FILE="./${KEY_PAIR}.pem"
SECURITY_GROUP_NAME="cirries-docker-role-sg"
ROLE_NAME="CirriesEC2DockerRole"
INSTANCE_PROFILE_NAME="CirriesEC2DockerProfile"
ECR_REGISTRY="709825985650.dkr.ecr.us-east-1.amazonaws.com"
ECR_IMAGE="$ECR_REGISTRY/cirries/dart_ai_vpb_agent_10gbps:0.0.1"

echo "🚀 Launching EC2 + Docker Container (Collector: $COLLECTOR_IP)"

# ---------- IAM Role ----------
set +e
if ! aws iam get-role --role-name "$ROLE_NAME" --region $REGION >/dev/null 2>&1; then
  aws iam create-role \
    --role-name "$ROLE_NAME" \
    --assume-role-policy-document '{
      "Version": "2012-10-17",
      "Statement": [{
        "Effect": "Allow",
        "Principal": {"Service": "ec2.amazonaws.com"},
        "Action": "sts:AssumeRole"
      }]
    }'
  aws iam attach-role-policy --role-name "$ROLE_NAME" --policy-arn arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryFullAccess
  aws iam attach-role-policy --role-name "$ROLE_NAME" --policy-arn arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
  aws iam attach-role-policy --role-name "$ROLE_NAME" --policy-arn arn:aws:iam::aws:policy/CloudWatchLogsFullAccess
fi

if ! aws iam get-instance-profile --instance-profile-name "$INSTANCE_PROFILE_NAME" >/dev/null 2>&1; then
  aws iam create-instance-profile --instance-profile-name "$INSTANCE_PROFILE_NAME"
  sleep 5
  aws iam add-role-to-instance-profile --instance-profile-name "$INSTANCE_PROFILE_NAME" --role-name "$ROLE_NAME"
fi
set -e

# ---------- Security Group ----------
VPC_ID=$(aws ec2 describe-vpcs --query "Vpcs[0].VpcId" --output text --region $REGION)
SG_ID=$(aws ec2 create-security-group \
  --group-name $SECURITY_GROUP_NAME \
  --description "Cirries Test SG" \
  --vpc-id $VPC_ID \
  --region $REGION \
  --query 'GroupId' --output text)
aws ec2 authorize-security-group-ingress --group-id $SG_ID --protocol tcp --port 22 --cidr 0.0.0.0/0 --region $REGION
aws ec2 authorize-security-group-ingress --group-id $SG_ID --protocol udp --port 4789 --cidr 0.0.0.0/0 --region $REGION

# ---------- EC2 Launch ----------
USER_DATA=$(cat <<EOF
#!/bin/bash
set -e
exec > /var/log/cirries-init.log 2>&1

echo "=== Starting setup at \$(date) ==="

# Docker setup
apt-get update -y
apt-get install -y ca-certificates curl gnupg lsb-release awscli jq
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
chmod a+r /etc/apt/keyrings/docker.asc
echo "deb [arch=\$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] \
https://download.docker.com/linux/ubuntu \$(. /etc/os-release && echo "\${UBUNTU_CODENAME:-\$VERSION_CODENAME}") stable" \
| tee /etc/apt/sources.list.d/docker.list > /dev/null
apt-get update -y
apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

systemctl enable docker
systemctl start docker

# Wait for Docker
for i in {1..10}; do
  if systemctl is-active --quiet docker; then
    echo "✅ Docker active"
    break
  fi
  echo "⏳ Waiting for Docker... (\$i)"
  sleep 30
done

# Kernel modules
modprobe vxlan sch_ingress act_mirred cls_u32 sch_clsact || true

# Login to ECR
for i in {1..5}; do
  echo "🔐 Logging in to ECR (attempt \$i)..."
  if aws ecr get-login-password --region $REGION | docker login --username AWS --password-stdin $ECR_REGISTRY; then
    break
  fi
  sleep 30
done

# Pull + Run container
for i in {1..5}; do
  if docker pull $ECR_IMAGE; then
    echo "✅ Image pulled"
    break
  fi
  sleep 30
done

docker run -d  --privileged --name cirries-diverter --net=host \\
  -e REMOTE_IP=$COLLECTOR_IP \\
  -e MIRROR_TYPE=2 -e VXLAN_ID=100 -e VXLAN_PORT=4789 \\
  $ECR_IMAGE

sleep 30
docker ps
echo "=== Setup complete at \$(date) ==="
docker logs -f cirries-diverter
EOF
)

SUBNET_ID=$(aws ec2 describe-subnets --query "Subnets[0].SubnetId" --output text --region $REGION)
INSTANCE_ID=$(aws ec2 run-instances \
  --image-id $AMI_ID \
  --count 1 \
  --instance-type $INSTANCE_TYPE \
  --key-name $KEY_PAIR \
  --security-group-ids $SG_ID \
  --subnet-id $SUBNET_ID \
  --associate-public-ip-address \
  --iam-instance-profile Name=$INSTANCE_PROFILE_NAME \
  --user-data "$USER_DATA" \
  --region $REGION \
  --query 'Instances[0].InstanceId' --output text)

aws ec2 wait instance-running --instance-ids $INSTANCE_ID --region $REGION
INSTANCE_IP=$(aws ec2 describe-instances --instance-ids $INSTANCE_ID --region $REGION --query 'Reservations[0].Instances[0].PublicIpAddress' --output text)
echo "🌍 Public IP: $INSTANCE_IP"

echo "⏳ Waiting for startup (1 min)..."
sleep 180

echo "🔍 Checking container..."
ssh -o StrictHostKeyChecking=no -i "$KEY_FILE" ubuntu@"$INSTANCE_IP" "sudo docker ps -a; echo '--- Logs ---'; sudo cat /var/log/cirries-init.log | tail -n 30"

echo ""
read -p '🧹 Clean up instance, SG, and IAM role? (y/n): ' CLEANUP
if [ "\$CLEANUP" == "y" ]; then
  aws ec2 terminate-instances --instance-ids "\$INSTANCE_ID" --region $REGION >/dev/null
  aws ec2 wait instance-terminated --instance-ids "\$INSTANCE_ID" --region $REGION
  aws ec2 delete-security-group --group-id "\$SG_ID" --region $REGION >/dev/null
  aws iam remove-role-from-instance-profile --instance-profile-name "$INSTANCE_PROFILE_NAME" --role-name "$ROLE_NAME" >/dev/null 2>&1 || true
  aws iam delete-instance-profile --instance-profile-name "$INSTANCE_PROFILE_NAME" >/dev/null 2>&1 || true
  aws iam delete-role --role-name "$ROLE_NAME" >/dev/null 2>&1 || true
  echo "✅ Cleanup complete."
else
  echo "⚠️ Instance retained. SSH with:"
  echo "ssh -i $KEY_FILE ubuntu@$INSTANCE_IP"
  echo "Check logs: sudo cat /var/log/cirries-init.log"
fi