# KMS Client Interface for Multiple Backends
# Provides unified interface for Age, AWS KMS, HashiCorp Vault, and Cosmian KMS
# Prioritizes plugin-based implementations for 10x performance improvement

use std log
use ../config/accessor.nu *
use ../utils/error.nu throw-error
use ../utils/interface.nu _print
use ../plugins/kms.nu [plugin-kms-encrypt plugin-kms-decrypt plugin-kms-status plugin-kms-info]

# KMS Client for encryption/decryption operations
export def kms-encrypt [
    data: string            # Data to encrypt
    key_id?: string         # Key ID (backend-specific)
    --backend: string = ""  # rustyvault, age, aws-kms, vault, cosmian (auto-detect if empty)
    --output-format: string = "base64"  # base64, hex, binary
]: nothing -> string {
    let kms_backend = if ($backend | is-empty) {
        detect-kms-backend
    } else {
        $backend
    }

    # Try plugin-based encryption first (10x faster for rustyvault and age)
    let plugin_info = (plugin-kms-info)

    if $plugin_info.plugin_available and $plugin_info.plugin_enabled {
        if $kms_backend in ["rustyvault", "age"] {
            let result = (do -i {
                plugin-kms-encrypt $data --backend $kms_backend --key-id ($key_id | default "")
            })

            if $result != null {
                if ($result | describe) == "record" and "ciphertext" in $result {
                    return $result.ciphertext
                } else if ($result | describe) == "string" {
                    return $result
                }
            } else {
                # Plugin failed, fall through to CLI implementation
                let debug_result = (do -i { get-debug-mode })
                if $debug_result != null and $debug_result {
                    print $"⚠️  Plugin encryption failed for ($kms_backend), using CLI fallback"
                }
            }
        }
    }

    # Fallback to CLI implementations
    match $kms_backend {
        "rustyvault" => {
            # Rustyvault fallback to vault CLI
            kms-encrypt-vault $data $key_id --output-format=$output_format
        }
        "age" => {
            kms-encrypt-age $data $key_id --output-format=$output_format
        }
        "aws-kms" => {
            kms-encrypt-aws $data $key_id --output-format=$output_format
        }
        "vault" => {
            kms-encrypt-vault $data $key_id --output-format=$output_format
        }
        "cosmian" => {
            kms-encrypt-cosmian $data $key_id --output-format=$output_format
        }
        _ => {
            error make {
                msg: $"Unsupported KMS backend: ($kms_backend)"
            }
        }
    }
}

# Decrypt data using KMS
export def kms-decrypt [
    encrypted_data: string  # Encrypted data
    key_id?: string         # Key ID (backend-specific)
    --backend: string = ""  # rustyvault, age, aws-kms, vault, cosmian (auto-detect if empty)
    --input-format: string = "base64"  # base64, hex, binary
]: nothing -> string {
    let kms_backend = if ($backend | is-empty) {
        detect-kms-backend
    } else {
        $backend
    }

    # Try plugin-based decryption first (10x faster for rustyvault and age)
    let plugin_info = (plugin-kms-info)

    if $plugin_info.plugin_available and $plugin_info.plugin_enabled {
        if $kms_backend in ["rustyvault", "age"] {
            let result = (do -i {
                plugin-kms-decrypt $encrypted_data --backend $kms_backend --key-id ($key_id | default "")
            })

            if $result != null {
                return $result
            }

            # Plugin failed, fall through to CLI implementation
            let debug_result = (do -i { get-debug-mode })
            if $debug_result != null and $debug_result {
                print $"⚠️  Plugin decryption failed for ($kms_backend), using CLI fallback"
            }
        }
    }

    # Fallback to CLI implementations
    match $kms_backend {
        "rustyvault" => {
            # Rustyvault fallback to vault CLI
            kms-decrypt-vault $encrypted_data $key_id --input-format=$input_format
        }
        "age" => {
            kms-decrypt-age $encrypted_data $key_id --input-format=$input_format
        }
        "aws-kms" => {
            kms-decrypt-aws $encrypted_data $key_id --input-format=$input_format
        }
        "vault" => {
            kms-decrypt-vault $encrypted_data $key_id --input-format=$input_format
        }
        "cosmian" => {
            kms-decrypt-cosmian $encrypted_data $key_id --input-format=$input_format
        }
        _ => {
            error make {
                msg: $"Unsupported KMS backend: ($kms_backend)"
            }
        }
    }
}

# Age backend encryption
def kms-encrypt-age [
    data: string
    key_id?: string
    --output-format: string = "base64"
]: nothing -> string {
    # Get Age recipients
    let recipients = if ($key_id | is-not-empty) {
        $key_id
    } else {
        $env.SOPS_AGE_RECIPIENTS? | default (get-sops-age-recipients)
    }

    if ($recipients | is-empty) {
        error make {
            msg: "No Age recipients configured"
        }
    }

    # Encrypt with age
    let encrypted = ($data | ^age -r $recipients -a | complete)

    if $encrypted.exit_code != 0 {
        error make {
            msg: $"Age encryption failed: ($encrypted.stderr)"
        }
    }

    $encrypted.stdout
}

# Age backend decryption
def kms-decrypt-age [
    encrypted_data: string
    key_id?: string
    --input-format: string = "base64"
]: nothing -> string {
    # Get Age key file
    let key_file = if ($key_id | is-not-empty) {
        $key_id
    } else {
        get-sops-age-key-file
    }

    if ($key_file | is-empty) {
        error make {
            msg: "No Age key file configured"
        }
    }

    if not ($key_file | path exists) {
        error make {
            msg: $"Age key file not found: ($key_file)"
        }
    }

    # Decrypt with age
    let decrypted = ($encrypted_data | ^age -d -i $key_file | complete)

    if $decrypted.exit_code != 0 {
        error make {
            msg: $"Age decryption failed: ($decrypted.stderr)"
        }
    }

    $decrypted.stdout
}

# AWS KMS backend encryption
def kms-encrypt-aws [
    data: string
    key_id?: string
    --output-format: string = "base64"
]: nothing -> string {
    # Get KMS key ID from config or parameter
    let kms_key = if ($key_id | is-not-empty) {
        $key_id
    } else {
        get-config-value "kms.aws.key_id" ""
    }

    if ($kms_key | is-empty) {
        error make {
            msg: "No AWS KMS key ID configured"
        }
    }

    # Check if AWS CLI is available
    let aws_check = (^which aws | complete)
    if $aws_check.exit_code != 0 {
        error make {
            msg: "AWS CLI not installed"
        }
    }

    # Encrypt with AWS KMS
    let encrypted = ($data | ^aws kms encrypt --key-id $kms_key --plaintext fileb:///dev/stdin --output text --query CiphertextBlob | complete)

    if $encrypted.exit_code != 0 {
        error make {
            msg: $"AWS KMS encryption failed: ($encrypted.stderr)"
        }
    }

    $encrypted.stdout | str trim
}

# AWS KMS backend decryption
def kms-decrypt-aws [
    encrypted_data: string
    key_id?: string
    --input-format: string = "base64"
]: nothing -> binary {
    # Check if AWS CLI is available
    let aws_check = (^which aws | complete)
    if $aws_check.exit_code != 0 {
        error make {
            msg: "AWS CLI not installed"
        }
    }

    # Decrypt with AWS KMS
    let decrypted = ($encrypted_data | ^aws kms decrypt --ciphertext-blob fileb:///dev/stdin --output text --query Plaintext | complete)

    if $decrypted.exit_code != 0 {
        error make {
            msg: $"AWS KMS decryption failed: ($decrypted.stderr)"
        }
    }

    $decrypted.stdout | str trim | decode base64
}

# HashiCorp Vault backend encryption
def kms-encrypt-vault [
    data: string
    key_id?: string
    --output-format: string = "base64"
]: nothing -> string {
    # Get Vault configuration
    let vault_addr = $env.VAULT_ADDR? | default (get-config-value "kms.vault.address" "")
    let vault_token = $env.VAULT_TOKEN? | default (get-config-value "kms.vault.token" "")
    let transit_key = if ($key_id | is-not-empty) { $key_id } else { get-config-value "kms.vault.transit_key" "provisioning" }

    if ($vault_addr | is-empty) {
        error make {
            msg: "Vault address not configured"
        }
    }

    # Encode data to base64 for Vault transit
    let plaintext_b64 = ($data | encode base64)

    # Encrypt with Vault transit
    let vault_payload = {plaintext: $plaintext_b64} | to json

    let encrypted = (^curl -s -X POST -H $"X-Vault-Token: ($vault_token)" -H "Content-Type: application/json" -d $vault_payload $"($vault_addr)/v1/transit/encrypt/($transit_key)" | complete)

    if $encrypted.exit_code != 0 {
        error make {
            msg: $"Vault encryption failed: ($encrypted.stderr)"
        }
    }

    let response = ($encrypted.stdout | from json)

    if ($response | get errors? | default null) != null {
        error make {
            msg: $"Vault encryption error: ($response.errors | to json)"
        }
    }

    $response.data.ciphertext
}

# HashiCorp Vault backend decryption
def kms-decrypt-vault [
    encrypted_data: string
    key_id?: string
    --input-format: string = "base64"
]: nothing -> binary {
    # Get Vault configuration
    let vault_addr = $env.VAULT_ADDR? | default (get-config-value "kms.vault.address" "")
    let vault_token = $env.VAULT_TOKEN? | default (get-config-value "kms.vault.token" "")
    let transit_key = if ($key_id | is-not-empty) { $key_id } else { get-config-value "kms.vault.transit_key" "provisioning" }

    if ($vault_addr | is-empty) {
        error make {
            msg: "Vault address not configured"
        }
    }

    # Decrypt with Vault transit
    let vault_payload = {ciphertext: $encrypted_data} | to json

    let decrypted = (^curl -s -X POST -H $"X-Vault-Token: ($vault_token)" -H "Content-Type: application/json" -d $vault_payload $"($vault_addr)/v1/transit/decrypt/($transit_key)" | complete)

    if $decrypted.exit_code != 0 {
        error make {
            msg: $"Vault decryption failed: ($decrypted.stderr)"
        }
    }

    let response = ($decrypted.stdout | from json)

    if ($response | get errors? | default null) != null {
        error make {
            msg: $"Vault decryption error: ($response.errors | to json)"
        }
    }

    $response.data.plaintext | decode base64
}

# Cosmian KMS backend encryption
def kms-encrypt-cosmian [
    data: string
    key_id?: string
    --output-format: string = "base64"
]: nothing -> string {
    # Get Cosmian KMS configuration
    let kms_server = get-kms-server

    if ($kms_server | is-empty) {
        error make {
            msg: "Cosmian KMS server not configured"
        }
    }

    # Use curl to call Cosmian KMS REST API
    let encrypted = ($data | ^curl -s -X POST -H "Content-Type: application/octet-stream" --data-binary @- $"($kms_server)/encrypt" | complete)

    if $encrypted.exit_code != 0 {
        error make {
            msg: $"Cosmian KMS encryption failed: ($encrypted.stderr)"
        }
    }

    $encrypted.stdout | encode base64
}

# Cosmian KMS backend decryption
def kms-decrypt-cosmian [
    encrypted_data: string
    key_id?: string
    --input-format: string = "base64"
]: nothing -> string {
    # Get Cosmian KMS configuration
    let kms_server = get-kms-server

    if ($kms_server | is-empty) {
        error make {
            msg: "Cosmian KMS server not configured"
        }
    }

    # Decode from base64 first
    let binary_data = ($encrypted_data | decode base64)

    # Use curl to call Cosmian KMS REST API
    let decrypted = ($binary_data | ^curl -s -X POST -H "Content-Type: application/octet-stream" --data-binary @- $"($kms_server)/decrypt" | complete)

    if $decrypted.exit_code != 0 {
        error make {
            msg: $"Cosmian KMS decryption failed: ($decrypted.stderr)"
        }
    }

    $decrypted.stdout
}

# Detect KMS backend from configuration
# Priority: rustyvault (fastest) > age (fastest local) > vault > aws-kms > cosmian
def detect-kms-backend []: nothing -> string {
    let kms_enabled = (get-kms-enabled)

    # Check if plugin is available to prefer native backends
    let plugin_info = (do -i { plugin-kms-info })
    let plugin_info = if $plugin_info != null {
        $plugin_info
    } else {
        { plugin_available: false, default_backend: "age" }
    }

    if not $kms_enabled {
        # Default to Age if KMS not enabled (plugin accelerated if available)
        return "age"
    }

    let kms_mode = (get-kms-mode)

    match $kms_mode {
        "local" => {
            let local_provider = (get-kms-local-provider)
            # Prefer rustyvault if plugin available, otherwise use configured provider
            if $plugin_info.plugin_available and $local_provider in ["vault", "rustyvault"] {
                "rustyvault"
            } else if $local_provider == "age" {
                "age"
            } else {
                $local_provider
            }
        }
        "remote" => {
            # Remote KMS (prefer rustyvault for Vault, otherwise Cosmian or AWS)
            let kms_server = (get-kms-server)
            if ($kms_server | str contains "vault") {
                # Use rustyvault native client if plugin available
                if $plugin_info.plugin_available {
                    "rustyvault"
                } else {
                    "vault"
                }
            } else if ($kms_server | str contains "cosmian") {
                "cosmian"
            } else if ($kms_server | str contains "aws") or ($kms_server | str contains "kms") {
                "aws-kms"
            } else {
                "cosmian"
            }
        }
        _ => "age"
    }
}

# Test KMS connectivity and functionality
export def kms-test [
    --backend: string = ""  # rustyvault, age, aws-kms, vault, cosmian (auto-detect if empty)
]: nothing -> record {
    print $"🧪 Testing KMS backend..."

    let kms_backend = if ($backend | is-empty) {
        detect-kms-backend
    } else {
        $backend
    }

    print $"   Backend: ($kms_backend)"

    # Check if using plugin
    let plugin_info = (plugin-kms-info)
    if $plugin_info.plugin_available and $kms_backend in ["rustyvault", "age"] {
        print $"   Mode: ⚡ Plugin-accelerated (10x faster)"
    } else {
        print $"   Mode: CLI/HTTP fallback"
    }

    let test_data = "Hello, KMS encryption test!"
    let initial_result = {
        backend: $kms_backend
        plugin_used: ($plugin_info.plugin_available and $kms_backend in ["rustyvault", "age"])
        encryption_success: false
        decryption_success: false
        round_trip_success: false
        encryption_time: 0ms
        decryption_time: 0ms
        error: null
    }

    mut test_result = $initial_result

    # Test encryption
    print $"   Testing encryption..."
    let enc_start = (date now)
    let encrypted = (do -i {
        kms-encrypt $test_data --backend=$kms_backend
    })

    if $encrypted != null {
        let enc_time = ((date now) - $enc_start)
        $test_result = ($test_result | upsert encryption_success true | upsert encryption_time $enc_time)
        print $"   Encrypted in ($enc_time)"

        # Test decryption
        print $"   Testing decryption..."
        let dec_start = (date now)
        let decrypted = (do -i {
            kms-decrypt $encrypted --backend=$kms_backend
        })

        if $decrypted != null {
            let dec_time = ((date now) - $dec_start)
            $test_result = ($test_result | upsert decryption_success true | upsert decryption_time $dec_time)
            print $"   Decrypted in ($dec_time)"

            # Verify round-trip
            if $decrypted == $test_data {
                print $"   ✅ Round-trip successful"
                $test_result = ($test_result | upsert round_trip_success true)
            } else {
                print $"   ❌ Round-trip failed: data mismatch"
                $test_result = ($test_result | upsert error "Round-trip data mismatch")
            }
        } else {
            print $"   ❌ Decryption failed"
            $test_result = ($test_result | upsert error "Decryption failed")
        }
    } else {
        print $"   ❌ Encryption failed"
        $test_result = ($test_result | upsert error "Encryption failed")
    }

    print ""
    if $test_result.round_trip_success {
        print $"✅ KMS backend ($kms_backend) is working correctly"
        print $"   Total time: ($test_result.encryption_time + $test_result.decryption_time)"
    } else {
        print $"❌ KMS backend ($kms_backend) test failed"
    }

    $test_result
}

# List available KMS backends
export def kms-list-backends [] {
    print "Available KMS Backends:"
    print "======================="
    print ""
    print "Local Backends:"
    print "  • age          - Age encryption (file-based keys)"
    print ""
    print "Cloud Backends:"
    print "  • aws-kms      - AWS Key Management Service"
    print ""
    print "Enterprise Backends:"
    print "  • vault        - HashiCorp Vault Transit Engine"
    print "  • cosmian      - Cosmian KMS (confidential computing)"
    print ""
    print "Current Configuration:"

    let kms_enabled = (get-kms-enabled)
    print $"  KMS Enabled: ($kms_enabled)"

    if $kms_enabled {
        let kms_mode = (get-kms-mode)
        print $"  Mode: ($kms_mode)"

        let detected_backend = (detect-kms-backend)
        print $"  Detected Backend: ($detected_backend)"
    } else {
        print "  Detected Backend: age (default)"
    }
}

# Get KMS backend status
export def kms-status []: nothing -> record {
    # Try plugin status first
    let plugin_info = (do -i { plugin-kms-info })
    let plugin_info = if $plugin_info != null {
        $plugin_info
    } else {
        { plugin_available: false, default_backend: "age" }
    }

    if $plugin_info.plugin_available {
        let plugin_status = (do -i { plugin-kms-status })
        if $plugin_status != null {
            return ($plugin_status | upsert plugin_accelerated true)
        }
        # Plugin failed, continue with manual status check
    }

    # Manual status check fallback
    let kms_enabled = (get-kms-enabled)
    let backend = (detect-kms-backend)

    mut status = {
        enabled: $kms_enabled
        backend: $backend
        plugin_accelerated: false
        configured: false
        available: false
    }

    # Check if backend is configured
    match $backend {
        "rustyvault" | "age" => {
            let key_file = (get-sops-age-key-file)
            let recipients = ($env.SOPS_AGE_RECIPIENTS? | default "")

            $status = ($status | upsert configured (($key_file | is-not-empty) and ($recipients | is-not-empty)))
            let available = if ($key_file | is-not-empty) { $key_file | path exists } else { false }
            $status = ($status | upsert available $available)
            $status = ($status | upsert plugin_accelerated $plugin_info.plugin_available)
        }
        "aws-kms" => {
            let aws_check = (^which aws | complete)
            $status = ($status | upsert available ($aws_check.exit_code == 0))

            let key_id = (get-config-value "kms.aws.key_id" "")
            $status = ($status | upsert configured ($key_id | is-not-empty))
        }
        "vault" => {
            let vault_addr = ($env.VAULT_ADDR? | default "")
            let vault_token = ($env.VAULT_TOKEN? | default "")

            $status = ($status | upsert configured (($vault_addr | is-not-empty) and ($vault_token | is-not-empty)))

            # Test connectivity
            if $status.configured {
                let health_check = (^curl -s $"($vault_addr)/v1/sys/health" | complete)
                $status = ($status | upsert available ($health_check.exit_code == 0))
            }
        }
        "cosmian" => {
            let kms_server = (get-kms-server)
            $status = ($status | upsert configured ($kms_server | is-not-empty))

            # Test connectivity
            if $status.configured {
                let health_check = (^curl -s $"($kms_server)/health" | complete)
                $status = ($status | upsert available ($health_check.exit_code == 0))
            }
        }
    }

    $status
}

# Get configuration value with fallback
def get-config-value [
    path: string
    default_value: any
]: nothing -> any {
    # This would integrate with the config accessor
    # For now, return default
    $default_value
}

# Main help
export def main [] {
    print "KMS Client Interface"
    print "===================="
    print ""
    print "Commands:"
    print "  kms-encrypt <data>        Encrypt data using KMS"
    print "  kms-decrypt <data>        Decrypt data using KMS"
    print "  kms-test                  Test KMS backend functionality"
    print "  kms-list-backends         List available KMS backends"
    print "  kms-status                Show KMS backend status"
    print ""
    print "Supported Backends:"
    print "  age, aws-kms, vault, cosmian"
}