# KMS Client Interface for Multiple Backends
# Provides unified interface for Age, AWS KMS, HashiCorp Vault, and Cosmian KMS

use std log
use ../config/accessor.nu *
use ../utils/error.nu throw-error
use ../utils/interface.nu _print

# KMS Client for encryption/decryption operations
export def kms-encrypt [
    data: string            # Data to encrypt
    key_id?: string         # Key ID (backend-specific)
    --backend: string = ""  # 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
    }

    match $kms_backend {
        "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 = ""  # 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
    }

    match $kms_backend {
        "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 -o errors) != 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 -o errors) != 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
def detect-kms-backend []: nothing -> string {
    let kms_enabled = (get-kms-enabled)

    if not $kms_enabled {
        # Default to Age if KMS not enabled
        return "age"
    }

    let kms_mode = (get-kms-mode)

    match $kms_mode {
        "local" => {
            let local_provider = (get-kms-local-provider)
            $local_provider
        }
        "remote" => {
            # Remote KMS (Cosmian or Vault)
            let kms_server = (get-kms-server)
            if ($kms_server | str contains "vault") {
                "vault"
            } else {
                "cosmian"
            }
        }
        _ => "age"
    }
}

# Test KMS connectivity and functionality
export def kms-test [
    --backend: string = ""  # 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)"

    let test_data = "Hello, KMS encryption test!"
    let initial_result = {
        backend: $kms_backend
        encryption_success: false
        decryption_success: false
        round_trip_success: false
        error: null
    }

    let test_result = (do {
        # Test encryption
        print $"   Testing encryption..."
        let encrypted = (kms-encrypt $test_data --backend=$kms_backend)
        let encryption_result = ($initial_result | upsert encryption_success true)

        # Test decryption
        print $"   Testing decryption..."
        let decrypted = (kms-decrypt $encrypted --backend=$kms_backend)
        let decryption_result = ($encryption_result | upsert decryption_success true)

        # Verify round-trip
        if $decrypted == $test_data {
            print $"   ✅ Round-trip successful"
            $decryption_result | upsert round_trip_success true
        } else {
            print $"   ❌ Round-trip failed: data mismatch"
            $decryption_result | upsert error "Round-trip data mismatch"
        }
    } | complete | if $in.exit_code == 0 {
        $in.stdout
    } else {
        print $"   ❌ Test failed"
        $initial_result | upsert error "Test execution failed"
    })

    print ""
    if $test_result.round_trip_success {
        print $"✅ KMS backend ($kms_backend) is working correctly"
    } 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 {
    let kms_enabled = (get-kms-enabled)
    let backend = (detect-kms-backend)

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

    # Check if backend is configured
    match $backend {
        "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)
        }
        "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"
}