#!/bin/bash
set -e

# Function to extract domain from Git SSH URL
extract_domain() {
    local url=$1

    # Handle different SSH URL formats
    if [[ $url == ssh://* ]]; then
        # For URLs starting with ssh://, extract the domain
        echo "$url" | sed -E 's|ssh://[^@]+@([^:/]+).*|\1|'
    else
        # For URLs without ssh://, extract the domain after '@' and before ':'
        echo "$url" | sed -E 's|[^@]+@([^:/]+).*|\1|'
    fi
}

# Check if the required environment variables are set
if [ -z "$REPO_URL" ]; then
    echo "Error: REPO_URL environment variable must be set."
    exit 1
fi

if [ ! -f /root/.ssh/id_rsa ]; then
    echo "Error: SSH private key not found at /root/.ssh/id_rsa"
    exit 1
fi

if [ -z "$GIT_USER" ]; then
    echo "Error: GIT_USER environment variable must be set."
    exit 1
fi

if [ -z "$GIT_EMAIL" ]; then
    echo "Error: GIT_EMAIL environment variable must be set."
    exit 1
fi

# Set default package manager if not provided
PACKAGE_MANAGER=${PACKAGE_MANAGER:-npm}

add_teleport_labels() {
    # Default config path if not provided
    CONFIG_FILE=${CONFIG_FILE:-/etc/teleport.yaml}

    if [ -z "$TELEPORT_LABELS" ]; then
        echo "No TELEPORT_LABELS environment variable set, skipping label configuration"
        return 0
    fi

    if [ ! -f "$CONFIG_FILE" ]; then
        echo "Error: Teleport config not found at $CONFIG_FILE"
        return 1
    fi

    echo "Adding Teleport labels from TELEPORT_LABELS environment variable..."

    # Check if ssh_service section exists
    if ! grep -q "ssh_service:" "$CONFIG_FILE"; then
        echo "Error: ssh_service section not found in teleport.yaml"
        return 1
    fi

    # Extract existing labels from the config file
    declare -A EXISTING_LABELS
    IN_LABELS_SECTION=0
    IN_SSH_SERVICE=0
    while IFS= read -r line; do
        # Detect when we enter the ssh_service section
        if [[ "$line" =~ ^[[:space:]]*ssh_service:[[:space:]]*$ ]]; then
            IN_SSH_SERVICE=1
            continue
        fi
        # Detect when we enter the labels section
        if [ "$IN_SSH_SERVICE" -eq 1 ] && [[ "$line" =~ ^[[:space:]]*labels:[[:space:]]*$ ]]; then
            IN_LABELS_SECTION=1
            continue
        fi
        # Detect when we leave the ssh_service section (next top-level key)
        if [ "$IN_SSH_SERVICE" -eq 1 ] && [[ "$line" =~ ^[[:space:]]*[a-zA-Z_][a-zA-Z0-9_]*:[[:space:]]* ]] && [ "$IN_LABELS_SECTION" -eq 0 ]; then
            IN_SSH_SERVICE=0
        fi
        # Detect when we leave the labels section (next key at same or higher indentation)
        if [ "$IN_LABELS_SECTION" -eq 1 ] && [[ "$line" =~ ^[[:space:]]*[a-zA-Z_][a-zA-Z0-9_]*:[[:space:]]* ]]; then
            # Check if this is a key at the same or less indentation as labels:
            if [[ ! "$line" =~ ^[[:space:]]{4,} ]]; then
                IN_LABELS_SECTION=0
                IN_SSH_SERVICE=0
            fi
        fi
        # Extract label key:value pairs
        if [ "$IN_LABELS_SECTION" -eq 1 ]; then
            if [[ "$line" =~ ^[[:space:]]+([^:]+):[[:space:]]*(.+)$ ]]; then
                existing_key=$(echo "${BASH_REMATCH[1]}" | xargs)
                existing_value=$(echo "${BASH_REMATCH[2]}" | xargs)
                EXISTING_LABELS["$existing_key"]="$existing_value"
            fi
        fi
    done < "$CONFIG_FILE"

    # Check if all labels from TELEPORT_LABELS are already present
    IFS=',' read -ra LABELS <<< "$TELEPORT_LABELS"
    ALL_LABELS_EXIST=true
    for label in "${LABELS[@]}"; do
        # Trim whitespace
        label=$(echo "$label" | xargs)

        # Check if label contains colon
        if [[ "$label" == *":"* ]]; then
            key=$(echo "$label" | cut -d':' -f1 | xargs)

            if [ -n "$key" ]; then
                # Check if this label key doesn't exist
                if [ -z "${EXISTING_LABELS[$key]}" ]; then
                    ALL_LABELS_EXIST=false
                    break
                fi
            fi
        fi
    done

    # If all labels already exist, skip file modification
    if [ "$ALL_LABELS_EXIST" = true ]; then
        echo "All Teleport labels from TELEPORT_LABELS are already present, skipping update"
        return 0
    fi

    # Build the new ssh_service section with existing and new labels
    NEW_SSH_SECTION="ssh_service:"
    NEW_SSH_SECTION="$NEW_SSH_SECTION"$'\n'"  labels:"

    # First, add existing labels that are not being replaced
    for existing_key in "${!EXISTING_LABELS[@]}"; do
        NEW_SSH_SECTION="$NEW_SSH_SECTION"$'\n'"    $existing_key: ${EXISTING_LABELS[$existing_key]}"
    done

    IFS=',' read -ra LABELS <<< "$TELEPORT_LABELS"
    for label in "${LABELS[@]}"; do
        # Trim whitespace
        label=$(echo "$label" | xargs)

        # Check if label contains colon
        if [[ "$label" == *":"* ]]; then
            key=$(echo "$label" | cut -d':' -f1 | xargs)
            value=$(echo "$label" | cut -d':' -f2- | xargs)

            if [ -n "$key" ] && [ -n "$value" ]; then
                # Only add if the label key doesn't already exist
                if [ -z "${EXISTING_LABELS[$key]}" ]; then
                    NEW_SSH_SECTION="$NEW_SSH_SECTION"$'\n'"    $key: $value"
                    echo "Added label: $key = $value"
                else
                    echo "Label '$key' already exists with value '${EXISTING_LABELS[$key]}', skipping"
                fi
            else
                echo "Warning: Invalid label format '$label'. Expected format: key:value"
            fi
        else
            echo "Warning: Label '$label' does not contain colon separator. Expected format: key:value"
        fi
    done

    # Create a backup
    cp "$CONFIG_FILE" "$CONFIG_FILE.backup"

    # If ssh_service exists, replace it. Otherwise, append to the end.
    if grep -qE '^[[:space:]]*ssh_service:' "$CONFIG_FILE"; then
        echo "Found existing ssh_service section, updating labels..."
        awk -v new_section="$NEW_SSH_SECTION" '
        BEGIN {
            in_ssh_service = 0
            section_replaced = 0
        }
        /^[[:space:]]*ssh_service:/ {
            in_ssh_service = 1
            if (!section_replaced) {
                print new_section
                section_replaced = 1
            }
            next
        }
        in_ssh_service && /^[[:space:]]*[a-zA-Z_][a-zA-Z0-9_]*:/ {
            in_ssh_service = 0
        }
        in_ssh_service {
            next
        }
        { print $0 }
        ' "$CONFIG_FILE" > "$CONFIG_FILE.tmp" && mv "$CONFIG_FILE.tmp" "$CONFIG_FILE"
        echo "Teleport labels updated successfully"
    else
        echo "ssh_service section not found; appending to bottom of $CONFIG_FILE..."
        # Ensure file ends with a newline before appending
        tail -c1 "$CONFIG_FILE" | read -r _ || echo >> "$CONFIG_FILE"
        printf "%s\n" "$NEW_SSH_SECTION" >> "$CONFIG_FILE"
        echo "Teleport ssh_service section appended successfully"
    fi
}


# Optional Teleport setup
if [ -n "$TELEPORT_VERSION" ] && [ -n "$TELEPORT_URL" ] && [ -n "$TELEPORT_TOKEN" ]; then
    echo "Teleport configuration detected. Setting up Teleport..."

    # Set default edition if not provided
    TELEPORT_EDITION=${TELEPORT_EDITION:-oss}

    # Only install Teleport if it's not already installed
    if ! command -v teleport &> /dev/null; then
        echo "Installing Teleport version $TELEPORT_VERSION, edition $TELEPORT_EDITION..."
        curl https://goteleport.com/static/install.sh | bash -s "${TELEPORT_VERSION}" "${TELEPORT_EDITION}"
    fi

    # Create teleport configuration
    if [ ! -f /etc/teleport.yaml ]; then
        teleport configure --roles=node --token="$TELEPORT_TOKEN" --proxy="$TELEPORT_URL" --no-acme -o file
        echo "Created teleport configuration file"
    else
        echo "Teleport configuration file already exists, skipping configuration creation"
    fi

    # Get the hostname of the machine
    HOSTNAME=$(hostname)

    add_teleport_labels
    # Update nodename in teleport config
    sed -i "s/^  nodename:.*/  nodename: $HOSTNAME/" /etc/teleport.yaml

    # Start Teleport in the background
    echo "Starting Teleport service..."
    cat /etc/teleport.yaml
    nohup teleport start &
else
    echo "Skipping Teleport setup - credentials not provided"
fi

# Ensure the SSH key has the correct permissions
chmod 600 /root/.ssh/id_rsa

# Extract the domain from the REPO_URL
DOMAIN=$(extract_domain "$REPO_URL")

# Add the SSH configuration for the domain
echo "Host $DOMAIN
  HostName $DOMAIN
  IdentityFile /root/.ssh/id_rsa
  StrictHostKeyChecking no
" >> /root/.ssh/config

# Clone the repository if it doesn't exist
if [ ! -d "/workspace/.git" ]; then
    git clone "$REPO_URL" /workspace
fi

# Add the public SSH key to authorized_keys
if [ -f /root/.ssh/hostkey.pub ]; then
    cat /root/.ssh/hostkey.pub >> /root/.ssh/authorized_keys
fi

# Change to the repository directory
cd /workspace

# Install dependencies using the specified package manager
case $PACKAGE_MANAGER in
    npm)
        npm install
        ;;
    yarn)
        yarn install
        ;;
    pnpm)
        pnpm install
        ;;
    bun)
        bun install
        ;;
    *)
        echo "Error: Unsupported package manager: $PACKAGE_MANAGER"
        exit 1
        ;;
esac

# Copy the environment file to the repository's environments directory
if [ -f /environment/.env ]; then
    mkdir -p /workspace/environments
    cp /environment/.env /workspace/.env
else
    echo "Warning: Environment file not found at /environment/.env"
fi

# Set the git user and email
git config --global user.name "$GIT_USER"
git config --global user.email "$GIT_EMAIL"

# Start SSH server
service ssh start

# Determine the run command
if [ -z "$RUN_COMMAND" ]; then
    # If RUN_COMMAND is not set, try to find a suitable command from package.json
    if [ -f "package.json" ]; then
        if grep -q '"dev"' package.json; then
            RUN_COMMAND="dev"
        elif grep -q '"start"' package.json; then
            RUN_COMMAND="start"
        else
            echo "Error: Could not determine a suitable run command from package.json"
            exit 1
        fi
    else
        echo "Error: package.json not found and RUN_COMMAND not set"
        exit 1
    fi
fi

# Start the application with pm2 using the specified package manager and run command
exec pm2-runtime "$PACKAGE_MANAGER run $RUN_COMMAND"