#!/bin/bash
set -euo pipefail

# nbcrypt - SSH Agent based encryption/decryption tool
# Uses Ed25519 signature for deterministic key derivation

SCRIPT_NAME="$(basename "$0")"
KEY_SEED_MESSAGE="nbase2-secret-key-seed-v1"
KEY_IDENTITY="id_ed25519"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

error() {
    echo -e "${RED}❌ Error: $*${NC}" >&2
    exit 1
}

info() {
    echo -e "${GREEN}ℹ️  $*${NC}" >&2
}

warn() {
    echo -e "${YELLOW}⚠️  $*${NC}" >&2
}

usage() {
    cat <<EOF
Usage: $SCRIPT_NAME <command> [arguments]

Commands:
    encrypt <input> <output>   Encrypt a file using SSH Agent key
    decrypt [options] <input> <output>   Decrypt a file using SSH Agent key
    install-bws                Install Bitwarden Secrets Manager CLI (bws)
    check                      Check if SSH Agent has required key
    help                       Show this help message

Options for decrypt:
    -f, --force                Force download nbloader from BWS and setup SSH Agent
                               (skips SSH Agent check, requires BWS_ACCESS_TOKEN)

Requirements:
    - SSH Agent must be running with id_ed25519 key loaded
    - ssh-keygen and openssl commands must be available

Examples:
    $SCRIPT_NAME encrypt secrets.txt secrets.enc
    $SCRIPT_NAME decrypt secrets.enc secrets.txt
    $SCRIPT_NAME decrypt -f secrets.enc secrets.txt  # Force BWS download
    $SCRIPT_NAME install-bws
    $SCRIPT_NAME check

EOF
    exit 0
}

check_dependencies() {
    local missing=()
    
    for cmd in ssh-keygen openssl ssh-add; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done
    
    if [ ${#missing[@]} -gt 0 ]; then
        error "Missing required commands: ${missing[*]}"
    fi
}

check_ssh_agent() {
    # Check if SSH_AUTH_SOCK is set and valid
    if [ -n "${SSH_AUTH_SOCK:-}" ] && [ -S "${SSH_AUTH_SOCK}" ]; then
        # Check if ssh-add can connect to agent
        if ssh-add -l >/dev/null 2>&1; then
            # Check if Ed25519 key is loaded
            if ssh-add -l 2>/dev/null | grep -q "ED25519"; then
                info "SSH Agent check passed (Ed25519 key found)"
                return 0
            fi
        fi
    fi
    
    # No valid agent or no Ed25519 key - try BWS bootstrap
    info "No Ed25519 key found in SSH Agent. Attempting BWS bootstrap..."
    
    # Get BWS access token
    local bws_token="${BWS_ACCESS_TOKEN:-}"
    if [ -z "$bws_token" ]; then
        echo -n "Enter BWS_ACCESS_TOKEN: " >&2
        read bws_token
        if [ -z "$bws_token" ]; then
            error "BWS_ACCESS_TOKEN is required when SSH Agent has no Ed25519 key"
        fi
    fi
    
    # Try to load setup script from BWS
    load_bws_setup "$bws_token"
    
    # Re-check agent after BWS setup
    if [ -n "${SSH_AUTH_SOCK:-}" ] && [ -S "${SSH_AUTH_SOCK}" ]; then
        if ssh-add -l >/dev/null 2>&1 && ssh-add -l 2>/dev/null | grep -q "ED25519"; then
            info "SSH Agent check passed (Ed25519 key found after BWS setup)"
            return 0
        fi
    fi
    
    error "No Ed25519 key found in SSH Agent. Please add id_ed25519 with: ssh-add ~/.ssh/id_ed25519"
}

load_bws_setup() {
    local token="$1"
    local secret_id="6e70094b-6888-4fde-85f9-b3d6007fd68e"
    
    info "Loading setup script from BWS (secret ID: $secret_id)..."
    
    # Check if bws command exists, if not try to install it
    if ! command -v bws >/dev/null 2>&1; then
        warn "bws command not found. Attempting to install..."
        install_bws || error "Failed to install bws CLI"
    fi
    
    # Export token temporarily for bws command
    export BWS_ACCESS_TOKEN="$token"
    
    # Get the secret from BWS using secret ID
    local loader_script
    if command -v jq >/dev/null 2>&1; then
        loader_script=$(bws secret get "$secret_id" 2>&1 | jq -r '.value // empty')
    elif command -v python3 >/dev/null 2>&1; then
        loader_script=$(bws secret get "$secret_id" 2>&1 | python3 -c "import sys, json; print(json.load(sys.stdin).get('value', ''))" 2>/dev/null)
    else
        # Fallback: try to extract value with grep/sed (fragile but works for simple JSON)
        loader_script=$(bws secret get "$secret_id" 2>&1 | grep -o '"value": "[^"]*"' | sed 's/"value": "//;s/"$//' | head -1)
    fi
    
    if [ -z "$loader_script" ]; then
        error "Failed to retrieve secret from BWS. Check your token and secret ID."
    fi
    
    # Execute the loader script
    info "Executing BWS setup script..."
    eval "$loader_script"
    
    # Load agent environment if it was created
    local env_file="/tmp/.nb_agent_env_${USER:-$(id -un)}"
    if [ -f "$env_file" ]; then
        source "$env_file"
        info "SSH Agent environment loaded from $env_file"
    fi
}

install_bws() {
    local arch
    arch=$(uname -m)
    local bws_version="1.0.0"
    local bws_bin_dir="${HOME}/.local/bin"
    local bws_path="${bws_bin_dir}/bws"
    
    mkdir -p "$bws_bin_dir"
    export PATH="$bws_bin_dir:$PATH"
    
    # Determine architecture
    case "$arch" in
        x86_64)
            arch="x86_64-unknown-linux-gnu"
            ;;
        aarch64|arm64)
            arch="aarch64-unknown-linux-gnu"
            ;;
        *)
            error "Unsupported architecture: $arch"
            ;;
    esac
    
    local zip_name="bws-${arch}-${bws_version}.zip"
    local url="https://github.com/bitwarden/sdk-sm/releases/download/bws-v${bws_version}/${zip_name}"
    
    info "Downloading bws v${bws_version} for ${arch}..."
    
    # Install dependencies if needed
    if ! command -v unzip >/dev/null 2>&1; then
        if command -v apt-get >/dev/null 2>&1; then
            sudo apt-get update -qq && sudo apt-get install -y unzip >/dev/null 2>&1
        elif command -v yum >/dev/null 2>&1; then
            sudo yum install -y unzip >/dev/null 2>&1
        fi
    fi
    
    # Download and extract
    local temp_zip="/tmp/${zip_name}"
    if command -v wget >/dev/null 2>&1; then
        wget -q "$url" -O "$temp_zip" || return 1
    elif command -v curl >/dev/null 2>&1; then
        curl -sL "$url" -o "$temp_zip" || return 1
    else
        error "wget or curl is required to download bws"
    fi
    
    unzip -o "$temp_zip" -d "$bws_bin_dir" >/dev/null 2>&1 || return 1
    chmod +x "$bws_path"
    rm -f "$temp_zip"
    
    info "bws installed successfully at $bws_path"
    return 0
}

derive_key() {
    # Derive encryption key from Ed25519 signature
    # Ed25519 signatures are deterministic, so same message = same signature
    
    local temp_message=$(mktemp)
    local temp_sig=$(mktemp)
    local temp_pubkey=$(mktemp)
    
    trap "rm -f '$temp_message' '$temp_sig' '$temp_pubkey'" EXIT
    
    # Create message to sign
    echo -n "$KEY_SEED_MESSAGE" > "$temp_message"
    
    # Get public key from agent for the Ed25519 key
    ssh-add -L 2>/dev/null | grep "ssh-ed25519" | head -1 > "$temp_pubkey"
    
    if [ ! -s "$temp_pubkey" ]; then
        error "Could not extract Ed25519 public key from SSH Agent"
    fi
    
    # Sign the message using ssh-keygen
    # Note: ssh-keygen -Y sign requires the public key in "allowed signers" format
    local temp_allowed=$(mktemp)
    echo "key1 $(cat "$temp_pubkey")" > "$temp_allowed"
    
    # Sign the message
    if ! ssh-keygen -Y sign -f "$temp_allowed" -n "nbase2" < "$temp_message" > "$temp_sig" 2>/dev/null; then
        # If -Y sign doesn't work (older ssh-keygen), try alternative method
        # We'll use ssh-agent's signing capability directly through a workaround
        
        # Extract just the signature data using ssh-agent protocol
        # Since we can't easily do that in pure bash, we'll use a deterministic approach
        # based on the public key itself as a fallback
        
        warn "ssh-keygen -Y sign not available, using fallback key derivation"
        
        # Fallback: Use public key hash as seed
        cat "$temp_pubkey" | sha256sum | head -c 64
        rm -f "$temp_allowed"
        return 0
    fi
    
    rm -f "$temp_allowed"
    
    # Extract signature and hash it to create 256-bit key
    # The signature file contains base64-encoded signature data
    cat "$temp_sig" | base64 -d 2>/dev/null | sha256sum | head -c 64
}

encrypt_file() {
    local input="$1"
    local output="$2"
    
    if [ ! -f "$input" ]; then
        error "Input file not found: $input"
    fi
    
    check_ssh_agent
    
    info "Deriving encryption key from SSH Agent..."
    local key=$(derive_key)
    
    if [ -z "$key" ]; then
        error "Failed to derive encryption key"
    fi
    
    info "Encrypting $input -> $output..."
    
    # Use OpenSSL to encrypt with AES-256-CBC
    # Pass the hex key directly to openssl
    openssl enc -aes-256-cbc -salt -pbkdf2 -in "$input" -out "$output" -pass "pass:$key"
    
    if [ $? -eq 0 ]; then
        info "✅ Encryption successful: $output"
    else
        error "Encryption failed"
    fi
}

decrypt_file() {
    local force_mode="$1"
    local input="$2"
    local output="$3"
    
    if [ ! -f "$input" ]; then
        error "Input file not found: $input"
    fi
    
    # Force mode: skip SSH Agent check and directly load from BWS
    if [ "$force_mode" = "true" ]; then
        info "Force mode: Loading setup script from BWS..."
        local bws_token="${BWS_ACCESS_TOKEN:-}"
        if [ -z "$bws_token" ]; then
            echo -n "Enter BWS_ACCESS_TOKEN: " >&2
            read bws_token
            if [ -z "$bws_token" ]; then
                error "BWS_ACCESS_TOKEN is required in force mode"
            fi
        fi
        load_bws_setup "$bws_token"
        # After BWS setup, check SSH Agent again
        if [ -z "${SSH_AUTH_SOCK:-}" ] || [ ! -S "${SSH_AUTH_SOCK}" ]; then
            error "SSH Agent not available after BWS setup"
        fi
        if ! ssh-add -l >/dev/null 2>&1 || ! ssh-add -l 2>/dev/null | grep -q "ED25519"; then
            error "Ed25519 key not found in SSH Agent after BWS setup"
        fi
    else
        check_ssh_agent
    fi
    
    info "Deriving decryption key from SSH Agent..."
    local key=$(derive_key)
    
    if [ -z "$key" ]; then
        error "Failed to derive decryption key"
    fi
    
    info "Decrypting $input -> $output..."
    
    # Use OpenSSL to decrypt with AES-256-CBC
    if openssl enc -aes-256-cbc -d -salt -pbkdf2 -in "$input" -out "$output" -pass "pass:$key" 2>/dev/null; then
        info "✅ Decryption successful: $output"
    else
        error "Decryption failed (wrong key or corrupted file)"
    fi
}

# Main command dispatcher
main() {
    if [ $# -eq 0 ]; then
        usage
    fi
    
    local command="$1"
    shift
    
    # install-bws コマンドの場合は check_dependencies をスキップ
    if [ "$command" != "install-bws" ] && [ "$command" != "help" ] && [ "$command" != "--help" ] && [ "$command" != "-h" ]; then
        check_dependencies
    fi
    
    case "$command" in
        encrypt)
            if [ $# -ne 2 ]; then
                error "encrypt requires 2 arguments: <input> <output>"
            fi
            encrypt_file "$1" "$2"
            ;;
        decrypt)
            local force_mode=false
            local input=""
            local output=""
            
            # Parse arguments
            while [ $# -gt 0 ]; do
                case "$1" in
                    -f|--force)
                        force_mode=true
                        shift
                        ;;
                    -*)
                        error "Unknown option: $1"
                        ;;
                    *)
                        if [ -z "$input" ]; then
                            input="$1"
                        elif [ -z "$output" ]; then
                            output="$1"
                        else
                            error "Too many arguments for decrypt"
                        fi
                        shift
                        ;;
                esac
            done
            
            if [ -z "$input" ] || [ -z "$output" ]; then
                error "decrypt requires 2 arguments: <input> <output>"
            fi
            
            decrypt_file "$force_mode" "$input" "$output"
            ;;
        check)
            check_ssh_agent
            echo "✅ All checks passed"
            ;;
        install-bws)
            install_bws
            ;;
        help|--help|-h)
            usage
            ;;
        *)
            error "Unknown command: $command (use 'help' for usage)"
            ;;
    esac
}

main "$@"