#!/usr/bin/env bash

set -e

# ... (include validate_borg_url, escape_sed, etc. from previous steps)

generate_recovery_file() {
    local repo="$1"
    local pass="$2"
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local recovery_file="$HOME/borg_recovery_$timestamp.txt"

    {
        echo "BORG BACKUP RECOVERY INFORMATION"
        echo "Generated on: $(date)"
        echo "--------------------------------"
        echo "REPOSITORY URL: $repo"
        echo "PASSPHRASE:     $pass"
        echo "--------------------------------"
        echo "Keep this file in a safe, offline location."
    } > "$recovery_file"

    chmod 600 "$recovery_file"
    printf "\nCRITICAL: Recovery information saved to: $recovery_file\n"
    printf "Please move this to a secure location (e.g., a password manager or physical safe).\n"
}

# Validation function for Borg Repository URLs
validate_borg_url() {
    local url="$1"
    # Regex covers: 
    # 1. Local absolute/home paths (/... or ~/...)
    # 2. SSH shortcuts (user@host:path)
    # 3. Explicit protocols (ssh://, file://)
    local regex="^(/|~/|([a-zA-Z0-9._-]+@[a-zA-Z0-9._-]+:)|([a-z]+://)).+"
    
    if [[ $url =~ $regex ]]; then
        return 0
    else
        return 1
    fi
}

escape_sed() {
    printf '%s\n' "$1" | sed -e 's/[]\/$*.^|[]/\\&/g'
}

setup_backup () {
    local repopath="$1"
    local raw_passphrase="$2"

    printf "\nInstalling backup script environment. You will be prompted for your sudo password.\n"
    
    # Check if the source directory exists before copying
    if [ ! -d "borg" ]; then
        printf "Error: 'borg' source directory not found in current location.\n"
        exit 1
    fi

    sudo cp -r borg /opt/
    sudo chown -R "root:root" /opt/borg
    find /opt/borg -type d -exec sudo chmod 755 {} \;
    find /opt/borg -type f -exec sudo chmod 644 {} \;
    sudo chmod -R 755 /opt/borg/bin
    sudo chmod 700 /opt/borg/etc
    sudo chmod 700 /opt/borg/ssh
    sudo ln -s /opt/borg/systemd/backup.service /etc/systemd/system/backup.service
    sudo ln -s /opt/borg/systemd/backup.timer /etc/systemd/system/backup.timer
    sudo systemctl daemon-reload
    
    local escaped_path=$(escape_sed "$repopath")
    local escaped_pass=$(escape_sed "$raw_passphrase")

    sudo sed -i -f - /opt/borg/etc/borg_environment <<EOF
s|REPOPATH|$escaped_path|g
s|REPOPASSPHRASE|$escaped_pass|g
EOF

    sudo chmod 600 /opt/borg/etc/borg_environment
    printf "Installation complete.\n"

    generate_recovery_file "$repopath" "$raw_passphrase"
    sudo systemctl enable --now backup.timer
}

setup_encrypted_backup () {
    local repopath="$1"
    local passphrase1 passphrase2 passphrase

    printf "\n--- Encryption Setup ---\n"
    while true; do
        read -s -p "Enter Passphrase: " passphrase1
        printf "\n"
        read -s -p "Confirm Passphrase: " passphrase2
        printf "\n"
        if [[ -n "$passphrase1" && "$passphrase1" == "$passphrase2" ]]; then
            passphrase="$passphrase1"
            break
        else
            printf "Error: Passphrases do not match or are empty. Try again.\n"
        fi
    done

    setup_backup "$repopath" "$passphrase"
}

main () {
    if [ "$EUID" -eq "0" ]; then
        printf "Please do not run as root or with sudo.\n"
        exit 1
    fi

    local repopath=""
    local valid=0

    # Improved Repo Input Loop
    while [ $valid -eq 0 ]; do
        read -p "Enter BORG_REPO URL: " repopath
        
        if validate_borg_url "$repopath"; then
            read -p "Use '$repopath'? [y/N]: " ans
            if [[ "$ans" =~ ^[yY]$ ]]; then
                valid=1
            fi
        else
            printf "Invalid Format! Examples:\n"
            printf "  - /mnt/backup/repo\n"
            printf "  - user@server:backup_repo\n"
            printf "  - ssh://user@server:2222/path/to/repo\n\n"
        fi
    done

    local encryption=""
    while [[ ! "$encryption" =~ ^[yYnN]$ ]]; do
        read -p "Enable encryption? [y/n]: " encryption
    done

    if [[ "$encryption" =~ ^[yY]$ ]]; then
        setup_encrypted_backup "$repopath"
    else
        setup_backup "$repopath" ""
    fi
}

main "$@"