#!/bin/bash
set -euo pipefail

# nbcrypt - SSH Agent ベースの暗号化/復号ツール
# Ed25519 署名で決定的な鍵導出を行う

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

# 出力用の色
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # リセット

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

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

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

print_source_agent_hint() {
    local env_file="/tmp/.nb_agent_env_${USER:-$(id -un)}"
    echo "" >&2
    info "この agent を現在のシェルで使うには、次を実行:"
    echo "  source $env_file" >&2
}

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

Commands:
    enc <text>                 Encrypt a text string (or use enctext)
    dec <text>                 Decrypt a text string (or use dectext)
    enctext <text>             Encrypt a text string (alias for enc)
    dectext <text>             Decrypt a text string (alias for dec)
    encfile <input> <output>   Encrypt a file (or use encrypt)
    decfile <input> <output>   Decrypt a file (or use decrypt)
    encrypt <input> <output>   Encrypt a file (alias for encfile)
    decrypt [options] <input> <output>   Decrypt a file (alias for decfile)
    keychain [-k|--keep]       Fetch SSH keys from BWS and run ssh-add.
                               Default: remove key files after ssh-add. Use -k/--keep to keep.
    install-bws                Install Bitwarden Secrets Manager CLI (bws)
    check                      Check if SSH Agent has required key
    help                       Show this help message

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

Text commands support stdin:
    echo "secret" | $SCRIPT_NAME enc -
    echo "encrypted" | $SCRIPT_NAME dec -

Requirements:
    - SSH Agent with id_ed25519 key, or ~/.ssh/id_ed25519 file, or BWS_ACCESS_TOKEN
    - ssh-keygen and openssl commands must be available

Examples:
    $SCRIPT_NAME enc "secret text"
    $SCRIPT_NAME dec "ENC:U2FsdGVkX1..."
    $SCRIPT_NAME encfile secrets.txt secrets.enc
    $SCRIPT_NAME decfile secrets.enc secrets.txt
    $SCRIPT_NAME keychain
    $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
}

get_bws_token() {
    # 1. すでに環境変数にある
    if [ -n "${BWS_ACCESS_TOKEN:-}" ]; then
        echo "$BWS_ACCESS_TOKEN"
        return 0
    fi
    # 2. nbconf から取得を試みる
    local nbase_root="${NBASE_HOME:-$(cd "$(dirname "$0")/.." && pwd)}"
    local nbconf_cmd="${nbase_root}/bin/nbconf"
    if [ -f "$nbconf_cmd" ] && [ -x "$nbconf_cmd" ]; then
        local val
        val=$("$nbconf_cmd" get BWS_ACCESS_TOKEN -d 2>/dev/null) || true
        if [ -n "$val" ]; then
            echo "$val"
            return 0
        fi
    fi
    return 1
}

check_ssh_agent() {
    local env_file="/tmp/.nb_agent_env_${USER:-$(id -un)}"
    
    # まず、以前の BWS セットアップで作られた env ファイルがあれば読む
    # SSH_AUTH_SOCK の検査より前にやる（env にソケットパスが入っているため）
    if [ -f "$env_file" ]; then
        source "$env_file"
        info "SSH Agent environment loaded from $env_file"
        
        # 読んだあと ssh-add が動くか確認
        if [ -n "${SSH_AUTH_SOCK:-}" ] && [ -S "${SSH_AUTH_SOCK}" ]; then
            if ssh-add -l >/dev/null 2>&1; then
                # Ed25519 が載っているか確認
                if ssh-add -l 2>/dev/null | grep -q "ED25519"; then
                    info "SSH Agent check passed (Ed25519 key found)"
                    return 0
                fi
            fi
        fi
    fi
    
    # 現在の環境の SSH_AUTH_SOCK がセットかつ有効か（すでにシェルに Agent がある場合）
    if [ -n "${SSH_AUTH_SOCK:-}" ] && [ -S "${SSH_AUTH_SOCK}" ]; then
        # ssh-add で agent に接続できるか
        if ssh-add -l >/dev/null 2>&1; then
            # Ed25519 が載っているか確認
            if ssh-add -l 2>/dev/null | grep -q "ED25519"; then
                info "SSH Agent check passed (Ed25519 key found)"
                return 0
            fi
        fi
    fi
    
    # ここに来た場合: (1) env ファイルがない (2) あっても ssh-add 失敗 (3) Ed25519 がない
    # いずれも BWS から bootstrap を試す
    info "No Ed25519 key found in SSH Agent. Attempting BWS bootstrap..."
    
    # BWS トークン取得（環境変数、nbconf、または入力）
    local bws_token
    bws_token=$(get_bws_token) || true
    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
    
    # BWS からセットアップスクリプトを取得して実行
    load_bws_setup "$bws_token"
    
    # BWS セットアップ後に agent を再確認（env を読み直す）
    if [ -f "$env_file" ]; then
        source "$env_file"
    fi
    
    if [ -n "${SSH_AUTH_SOCK:-}" ] && [ -S "${SSH_AUTH_SOCK}" ]; then
        if ssh-add -l >/dev/null 2>&1; then
            if ssh-add -l 2>/dev/null | grep -q "ED25519"; then
                info "SSH Agent check passed (Ed25519 key found after BWS setup)"
                return 0
            fi
        else
            # BWS セットアップ後に ssh-add が失敗
            if [ -f "$env_file" ]; then
                warn "SSH Agent communication failed after BWS setup."
                warn "Please run: source $env_file"
                error "SSH Agent communication failed. Run 'source $env_file' and try again."
            else
                error "SSH Agent communication failed after BWS setup."
            fi
        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)..."
    
    # bws が無ければインストールを試す
    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
    
    # bws 用にトークンを一時 export
    export BWS_ACCESS_TOKEN="$token"
    
    # 秘密 ID で BWS から secret を取得
    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
        # フォールバック: grep/sed で value を取り出し（単純 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
    
    # ローダースクリプトを実行
    info "Executing BWS setup script..."
    eval "$loader_script"
    
    # 作成されていれば agent 用 env を読む
    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
}

keychain_cmd() {
    local keep_mode="$1"
    info "Fetching SSH keys from BWS and running ssh-add..."

    local bws_token
    bws_token=$(get_bws_token) || true
    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"
        fi
    fi

    load_bws_setup "$bws_token"

    if [ "$keep_mode" != "true" ]; then
        info "Removing key files from disk (use -k/--keep to keep)..."
        rm -f ~/.ssh/id_rsa ~/.ssh/id_ecdsa ~/.ssh/id_ed25519
    fi

    info "SSH agent setup complete."
    ssh-add -l
    print_source_agent_hint
}

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"
    
    # アーキテクチャ判定
    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}..."
    
    # 必要なら unzip 等を入れる
    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
    
    # ダウンロードして展開
    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() {
    # Ed25519 署名から暗号鍵を導出（同じメッセージなら同じ署名で決定的）
    # 優先順位: 1. SSH Agent, 2. 秘密鍵ファイル, 3. BWS (check_ssh_agent で処理)
    
    local temp_message=$(mktemp)
    local temp_sig=$(mktemp)
    local temp_pubkey=$(mktemp)
    local temp_allowed=$(mktemp)
    
    trap "rm -f '$temp_message' '$temp_sig' '$temp_pubkey' '$temp_allowed'" EXIT
    
    # 署名するメッセージを作成
    echo -n "$KEY_SEED_MESSAGE" > "$temp_message"
    
    # 1. SSH Agent から Ed25519 の公開鍵を取得を試みる
    if ssh-add -L 2>/dev/null | grep "ssh-ed25519" | head -1 | awk '{print $1 " " $2}' > "$temp_pubkey" 2>/dev/null && [ -s "$temp_pubkey" ]; then
        # Agent から取得できた場合
        echo "key1 $(cat "$temp_pubkey")" > "$temp_allowed"
        
        # メッセージに署名（Agent経由）
        if ssh-keygen -Y sign -f "$temp_allowed" -n "nbase2" < "$temp_message" > "$temp_sig" 2>/dev/null; then
            # 署名を抜き出してハッシュし 256bit 鍵に
            cat "$temp_sig" | base64 -d 2>/dev/null | sha256sum | head -c 64
            return 0
        fi
    fi
    
    # 2. SSH Agent がない場合、秘密鍵ファイルから直接署名を試みる
    local key_file="${NB_KEY_FILE:-${HOME}/.ssh/id_ed25519}"
    if [ -f "$key_file" ]; then
        # 公開鍵を抽出
        if ssh-keygen -y -f "$key_file" 2>/dev/null | awk '{print $1 " " $2}' > "$temp_pubkey" 2>/dev/null && [ -s "$temp_pubkey" ]; then
            echo "key1 $(cat "$temp_pubkey")" > "$temp_allowed"
            
            # 秘密鍵ファイルを使って署名
            if ssh-keygen -Y sign -f "$key_file" -n "nbase2" < "$temp_message" > "$temp_sig" 2>/dev/null; then
                # 署名を抜き出してハッシュし 256bit 鍵に
                cat "$temp_sig" | base64 -d 2>/dev/null | sha256sum | head -c 64
                return 0
            fi
        fi
    fi
    
    # 3. どちらも失敗した場合、エラー（BWS は check_ssh_agent で処理される）
    error "Could not derive key from SSH Agent or private key file. Please ensure SSH Agent is running with Ed25519 key, or set NB_KEY_FILE to a valid private key path."
}

encrypt_file() {
    local input="$1"
    local output="$2"
    
    if [ ! -f "$input" ]; then
        error "Input file not found: $input"
    fi
    
    # 秘密鍵ファイルがある場合は check_ssh_agent をスキップ
    local key_file="${NB_KEY_FILE:-${HOME}/.ssh/id_ed25519}"
    if [ ! -f "$key_file" ]; then
        check_ssh_agent
    fi
    
    info "Deriving encryption key..."
    local key=$(derive_key)
    
    if [ -z "$key" ]; then
        error "Failed to derive encryption key"
    fi
    
    info "Encrypting $input -> $output..."
    
    # OpenSSL で AES-256-CBC 暗号化（鍵は hex で渡す）
    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: SSH Agent チェックを飛ばし BWS から直接読む
    if [ "$force_mode" = "true" ]; then
        info "Force mode: Loading setup script from BWS..."
        local bws_token
        bws_token=$(get_bws_token) || true
        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"
        # BWS セットアップ後に agent を再確認
        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 をスキップ
        local key_file="${NB_KEY_FILE:-${HOME}/.ssh/id_ed25519}"
        if [ ! -f "$key_file" ]; then
            check_ssh_agent
        fi
    fi
    
    info "Deriving decryption key..."
    local key=$(derive_key)
    
    if [ -z "$key" ]; then
        error "Failed to derive decryption key"
    fi
    
    info "Decrypting $input -> $output..."
    
    # OpenSSL で 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
}

encrypt_text() {
    local input_text="$1"
    
    # 秘密鍵ファイルがある場合は check_ssh_agent をスキップ
    local key_file="${NB_KEY_FILE:-${HOME}/.ssh/id_ed25519}"
    if [ ! -f "$key_file" ]; then
        check_ssh_agent
    fi
    
    info "Deriving encryption key..."
    local key=$(derive_key)
    
    if [ -z "$key" ]; then
        error "Failed to derive encryption key"
    fi
    
    # 標準入力から読み取る場合
    if [ "$input_text" = "-" ]; then
        local temp_input=$(mktemp)
        cat > "$temp_input"
        input_text=$(cat "$temp_input")
        rm -f "$temp_input"
    fi
    
    # OpenSSL で暗号化して Base64 エンコード
    echo -n "$input_text" | openssl enc -aes-256-cbc -salt -pbkdf2 -pass "pass:$key" | base64 -w 0
    echo ""  # 改行を追加
}

decrypt_text() {
    local input_text="$1"
    
    # 秘密鍵ファイルがある場合は check_ssh_agent をスキップ
    local key_file="${NB_KEY_FILE:-${HOME}/.ssh/id_ed25519}"
    if [ ! -f "$key_file" ]; then
        check_ssh_agent
    fi
    
    info "Deriving decryption key..."
    local key=$(derive_key)
    
    if [ -z "$key" ]; then
        error "Failed to derive decryption key"
    fi
    
    # 標準入力から読み取る場合
    if [ "$input_text" = "-" ]; then
        input_text=$(cat)
    fi
    
    # Base64 デコードして OpenSSL で復号
    echo -n "$input_text" | base64 -d 2>/dev/null | openssl enc -aes-256-cbc -d -salt -pbkdf2 -pass "pass:$key" 2>/dev/null || error "Decryption failed (wrong key or corrupted data)"
}

# メインのコマンド分岐
main() {
    if [ $# -eq 0 ]; then
        usage
    fi
    
    local command="$1"
    shift
    
    # install-bws, keychain の場合は check_dependencies をスキップ
    if [ "$command" != "install-bws" ] && [ "$command" != "keychain" ] && [ "$command" != "help" ] && [ "$command" != "--help" ] && [ "$command" != "-h" ]; then
        check_dependencies
    fi
    
    case "$command" in
        enc|enctext)
            if [ $# -lt 1 ]; then
                error "enc requires 1 argument: <text> or '-' for stdin"
            fi
            encrypt_text "${1:-}"
            ;;
        dec|dectext)
            if [ $# -lt 1 ]; then
                error "dec requires 1 argument: <text> or '-' for stdin"
            fi
            decrypt_text "${1:-}"
            ;;
        encfile|encrypt)
            if [ $# -ne 2 ]; then
                error "encfile requires 2 arguments: <input> <output>"
            fi
            encrypt_file "$1" "$2"
            ;;
        decfile|decrypt)
            local force_mode=false
            local input=""
            local output=""
            
            # 引数解析
            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 "decfile requires 2 arguments: <input> <output>"
            fi
            
            decrypt_file "$force_mode" "$input" "$output"
            ;;
        check)
            check_ssh_agent
            echo "✅ All checks passed"
            ;;
        keychain)
            local keep_mode=false
            while [ $# -gt 0 ]; do
                case "$1" in
                    -k|--keep)
                        keep_mode=true
                        shift
                        ;;
                    -*)
                        error "Unknown option: $1"
                        ;;
                    *)
                        error "keychain accepts only -k/--keep"
                        ;;
                esac
            done
            keychain_cmd "$keep_mode"
            ;;
        install-bws)
            install_bws
            ;;
        help|--help|-h)
            usage
            ;;
        *)
            error "Unknown command: $command (use 'help' for usage)"
            ;;
    esac
}

main "$@"