#!/bin/bash
set -euxo pipefail

...
# Install Docker Compose
mkdir -p /usr/local/lib/docker/cli-plugins
curl -L "<https://github.com/docker/compose/releases/download/v2.29.2/docker-compose-linux-x86_64>" -o /usr/local/lib/docker/cli-plugins/docker-compose
chmod +x /usr/local/lib/docker/cli-plugins/docker-compose

# Setup directories
mkdir -p /data/postgres /opt/prefect /var/lib/tailscale
chown -R ec2-user:ec2-user /opt/prefect

# Mount Postgres volume
if ! blkid /dev/xvdb >/dev/null 2>&1; then
  mkfs -t xfs /dev/xvdb
fi
if ! grep -q "/dev/xvdb" /etc/fstab; then
  echo "/dev/xvdb /data/postgres xfs defaults,nofail 0 2" >> /etc/fstab
fi
mount -a
chown -R 999:999 /data/postgres

# Fetch or generate secrets from SSM
cd /opt/prefect
rm -f .env

# Create docker-compose.yml
cat > docker-compose.yml << 'COMPOSE_EOF'
services:
  postgres:
    image: postgres:15
    container_name: prefect-postgres
    environment:
      POSTGRES_USER: prefect
      POSTGRES_PASSWORD: $${DB_PASSWORD}
      POSTGRES_DB: prefect
    volumes:
      - /data/postgres:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U prefect"]
      interval: 5s
      timeout: 5s
      retries: 5
    restart: always
    networks:
      - prefect

  redis:
    image: redis:7-alpine
    container_name: prefect-redis
    healthcheck:
      test: ["CMD-SHELL", "redis-cli ping"]
      interval: 5s
      timeout: 5s
      retries: 5
    restart: always
    networks:
      - prefect

  prefect-server:
    image: prefecthq/prefect:3-latest
    container_name: prefect-server
    depends_on:
      postgres:
        condition: service_healthy
      redis:
        condition: service_healthy
    environment:
      # Database
      
      # Server config
      
      # Redis messaging
      # UI config - point to ALB for external access
      
    command: prefect server start
    ports:
      - "4200:4200"
    restart: always
    networks:
      - prefect

  tailscale:
    image: tailscale/tailscale:latest
    container_name: tailscale
    hostname: $${TAILSCALE_HOSTNAME}
    environment:
      - TS_AUTHKEY=$${TAILSCALE_AUTH_KEY}
      - TS_STATE_DIR=/var/lib/tailscale
      - TS_USERSPACE=false
      - TS_ACCEPT_DNS=true
      - TS_EXTRA_ARGS=--advertise-tags=tag:server
    volumes:
      - /var/lib/tailscale:/var/lib/tailscale
      - /dev/net/tun:/dev/net/tun
    cap_add:
      - NET_ADMIN
      - SYS_MODULE
    network_mode: host
    restart: always
networks:
  prefect:
    driver: bridge
COMPOSE_EOF

# Create .env file with secrets
cat > .env << ENV_EOF
DB_PASSWORD=$DB_PASSWORD
PREFECT_API_KEY=$PREFECT_API_KEY
TAILSCALE_AUTH_KEY=$TAILSCALE_AUTH_KEY
TAILSCALE_HOSTNAME=$TAILSCALE_HOSTNAME
APP_DOMAIN=$APP_DOMAIN
AWS_REGION=$AWS_REGION
PROJECT=$PROJECT
ENVIRONMENT=$ENVIRONMENT
ENV_EOF

# Ensure proper permissions
chown ec2-user:ec2-user docker-compose.yml .env
chmod 600 .env

# Configure Docker daemon to use CloudWatch logs driver
mkdir -p /etc/docker
cat > /etc/docker/daemon.json << DAEMON_EOF
{
  "log-driver": "awslogs",
  "log-opts": {
    "awslogs-region": "$AWS_REGION"
  }
}
DAEMON_EOF
systemctl restart docker

# Start services
docker compose up -d

# Wait for Tailscale to connect
echo "Waiting for Tailscale to establish connection..."
sleep 10

# If routes are specified, advertise them
if [ -n "$ADVERTISE_ROUTES" ] && [ "$ADVERTISE_ROUTES" != "" ]; then
  echo "Advertising routes: $ADVERTISE_ROUTES"
  docker exec tailscale tailscale set --advertise-routes="$ADVERTISE_ROUTES" || true
fi

# Display Tailscale status
docker exec tailscale tailscale status || true

# Create systemd service for auto-start
cat > /etc/systemd/system/prefect.service << 'SYSTEMD_EOF'
[Unit]
Description=Prefect Server Docker Compose
Requires=docker.service
After=docker.service

[Service]
Type=oneshot
WorkingDirectory=/opt/prefect
ExecStart=/usr/bin/docker compose up -d
ExecStop=/usr/bin/docker compose down
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
SYSTEMD_EOF

systemctl daemon-reload
systemctl enable prefect.service