#!/usr/bin/env nu

# KMS Service CLI Integration
# Provides commands to interact with the KMS service API

# Get KMS service base URL from environment or use default
def kms-url [] -> string {
    $env.KMS_SERVICE_URL? | default "http://localhost:8081"
}

# Build KMS API endpoint URL
def kms-endpoint [path: string] -> string {
    $"(kms-url)/api/v1/kms/(path)"
}

# Encrypt data using KMS
#
# Examples:
#   "secret-data" | kms encrypt
#   "api-key" | kms encrypt --context "env=prod,service=api"
export def "kms encrypt" [
    --context: string      # Encryption context (key1=val1,key2=val2)
    --backend: string = "vault"  # KMS backend (vault or aws-kms)
] -> string {
    let input = $in

    # Encode plaintext to base64
    let plaintext_b64 = $input | encode base64

    # Prepare request body
    let body = {
        plaintext: $plaintext_b64
        context: $context
    }

    # Make API request
    let response = http post (kms-endpoint "encrypt") $body | complete

    if $response.exit_code != 0 {
        error make {
            msg: "KMS encrypt request failed"
            label: {
                text: $response.stderr
                span: (metadata $input).span
            }
        }
    }

    let result = $response.stdout | from json

    if ($result | get -i error | is-not-empty) {
        error make {
            msg: "KMS encryption failed"
            label: {
                text: $result.error
                span: (metadata $input).span
            }
        }
    }

    $result.ciphertext
}

# Decrypt data using KMS
#
# Examples:
#   "vault:v1:..." | kms decrypt
#   $ciphertext | kms decrypt --context "env=prod,service=api"
export def "kms decrypt" [
    --context: string      # Encryption context (must match encryption context)
] -> string {
    let ciphertext = $in

    # Prepare request body
    let body = {
        ciphertext: $ciphertext
        context: $context
    }

    # Make API request
    let response = http post (kms-endpoint "decrypt") $body | complete

    if $response.exit_code != 0 {
        error make {
            msg: "KMS decrypt request failed"
            label: {
                text: $response.stderr
                span: (metadata $ciphertext).span
            }
        }
    }

    let result = $response.stdout | from json

    if ($result | get -i error | is-not-empty) {
        error make {
            msg: "KMS decryption failed"
            label: {
                text: $result.error
                span: (metadata $ciphertext).span
            }
        }
    }

    # Decode base64 plaintext
    $result.plaintext | decode base64
}

# Generate a data key for envelope encryption
#
# Examples:
#   kms generate-key
#   kms generate-key --key-spec AES_128
export def "kms generate-key" [
    --key-spec: string = "AES_256"  # Key specification (AES_128, AES_256, RSA_2048, RSA_4096)
] -> record {
    let body = {
        key_spec: $key_spec
    }

    let response = http post (kms-endpoint "generate-key") $body | complete

    if $response.exit_code != 0 {
        error make {
            msg: "KMS generate-key request failed"
            label: {
                text: $response.stderr
            }
        }
    }

    let result = $response.stdout | from json

    if ($result | get -i error | is-not-empty) {
        error make {
            msg: "KMS key generation failed"
            label: {
                text: $result.error
            }
        }
    }

    {
        plaintext: $result.plaintext
        ciphertext: $result.ciphertext
        key_id: $result.key_id
    }
}

# Rotate KMS key
#
# Examples:
#   kms rotate-key "arn:aws:kms:..."
export def "kms rotate-key" [
    key_id: string  # Key ID to rotate
] -> record {
    let body = {
        key_id: $key_id
    }

    let response = http post (kms-endpoint "rotate-key") $body | complete

    if $response.exit_code != 0 {
        error make {
            msg: "KMS rotate-key request failed"
            label: {
                text: $response.stderr
                span: (metadata $key_id).span
            }
        }
    }

    let result = $response.stdout | from json

    if ($result | get -i error | is-not-empty) {
        error make {
            msg: "KMS key rotation failed"
            label: {
                text: $result.error
                span: (metadata $key_id).span
            }
        }
    }

    {
        new_key_id: $result.new_key_id
        success: $result.success
    }
}

# Check KMS service health
#
# Examples:
#   kms health
export def "kms health" [] -> record {
    let response = http get (kms-endpoint "health") | complete

    if $response.exit_code != 0 {
        return {
            status: "unhealthy"
            error: $response.stderr
        }
    }

    $response.stdout | from json
}

# Get KMS service status
#
# Examples:
#   kms status
export def "kms status" [] -> record {
    let response = http get (kms-endpoint "status") | complete

    if $response.exit_code != 0 {
        error make {
            msg: "KMS status request failed"
            label: {
                text: $response.stderr
            }
        }
    }

    $response.stdout | from json
}

# Encrypt file using KMS
#
# Examples:
#   kms encrypt-file config.yaml
#   kms encrypt-file secrets.json --output secrets.enc --context "env=prod"
export def "kms encrypt-file" [
    input_file: path       # File to encrypt
    --output: path         # Output file (default: <input>.enc)
    --context: string      # Encryption context
] {
    if not ($input_file | path exists) {
        error make {
            msg: "Input file not found"
            label: {
                text: $input_file
                span: (metadata $input_file).span
            }
        }
    }

    let output_file = if ($output | is-empty) {
        $"($input_file).enc"
    } else {
        $output
    }

    print $"Encrypting ($input_file)..."

    # Read file and encrypt
    let plaintext = open --raw $input_file
    let ciphertext = $plaintext | kms encrypt --context $context

    # Save encrypted data
    $ciphertext | save --raw $output_file

    print $"✓ Encrypted file saved to ($output_file)"
}

# Decrypt file using KMS
#
# Examples:
#   kms decrypt-file config.yaml.enc
#   kms decrypt-file secrets.enc --output secrets.json --context "env=prod"
export def "kms decrypt-file" [
    input_file: path       # File to decrypt
    --output: path         # Output file (default: remove .enc extension)
    --context: string      # Encryption context (must match)
] {
    if not ($input_file | path exists) {
        error make {
            msg: "Input file not found"
            label: {
                text: $input_file
                span: (metadata $input_file).span
            }
        }
    }

    let output_file = if ($output | is-empty) {
        if ($input_file | str ends-with ".enc") {
            $input_file | str replace ".enc" ""
        } else {
            $"($input_file).dec"
        }
    } else {
        $output
    }

    print $"Decrypting ($input_file)..."

    # Read encrypted data and decrypt
    let ciphertext = open --raw $input_file
    let plaintext = $ciphertext | kms decrypt --context $context

    # Save decrypted data
    $plaintext | save --raw $output_file

    print $"✓ Decrypted file saved to ($output_file)"
}

# Envelope encrypt data (AWS KMS only)
#
# Examples:
#   "large-data" | kms envelope encrypt
export def "kms envelope encrypt" [] -> string {
    let input = $in
    let plaintext_b64 = $input | encode base64

    let body = { plaintext: $plaintext_b64 }
    let response = http post (kms-endpoint "envelope/encrypt") $body | complete

    if $response.exit_code != 0 {
        error make {
            msg: "KMS envelope encrypt request failed"
            label: { text: $response.stderr }
        }
    }

    let result = $response.stdout | from json
    if ($result | get -i error | is-not-empty) {
        error make {
            msg: "KMS envelope encryption failed"
            label: { text: $result.error }
        }
    }

    $result.ciphertext
}

# Envelope decrypt data (AWS KMS only)
#
# Examples:
#   $envelope_ciphertext | kms envelope decrypt
export def "kms envelope decrypt" [] -> string {
    let ciphertext = $in

    let body = { ciphertext: $ciphertext }
    let response = http post (kms-endpoint "envelope/decrypt") $body | complete

    if $response.exit_code != 0 {
        error make {
            msg: "KMS envelope decrypt request failed"
            label: { text: $response.stderr }
        }
    }

    let result = $response.stdout | from json
    if ($result | get -i error | is-not-empty) {
        error make {
            msg: "KMS envelope decryption failed"
            label: { text: $result.error }
        }
    }

    $result.plaintext | decode base64
}