#!/bin/bash

# Standalone TPM Certificate Signing Request automation script
# Automates TPM CSR process for cisco-hf-agents daemonset with no external dependencies

set -euo pipefail

# Constants
readonly DRAKECTL_PATH="/opt/cisco/bin/tortuga/drakectl"
readonly DAEMONSET="${DAEMONSET:-cisco-hf-agents}"
readonly NAMESPACE="${NAMESPACE:-cisco-hf}"

# Global variables
FORCE_FLAG=""
USE_PROXY=""
INSECURE=""
HYPERFABRIC_API_URL=""
K8S_CLI=""

# ============================================================================
# LOGGING FUNCTIONS
# ============================================================================

# Simple logging with timestamp
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') [$1] $2" >&2
}

# ============================================================================
# UTILITY FUNCTIONS
# ============================================================================

# Check if required commands are available
check_required_commands() {
    local missing_commands=()
    
    for cmd in "$@"; do
        if ! command -v "$cmd" &> /dev/null; then
            missing_commands+=("$cmd")
        fi
    done
    
    if [[ ${#missing_commands[@]} -gt 0 ]]; then
        for cmd in "${missing_commands[@]}"; do
            log "ERROR" "$cmd is not installed"
        done
        return 1
    fi
    
    return 0
}

# Check if bearer token is set
check_bearer_token() {
    if [[ -z "${CISCO_HF_BEARER_TOKEN:-}" ]]; then
        log "ERROR" "CISCO_HF_BEARER_TOKEN environment variable is required"
        return 1
    fi
    return 0
}

# ============================================================================
# JSON AND CERTIFICATE FUNCTIONS
# ============================================================================

# Create JSON payload from CSR
create_csr_json_payload() {
    local csr="$1"
    local escaped_csr
    escaped_csr=$(echo "$csr" | jq -Rs .)
    jq -n --argjson csr "$escaped_csr" '{csr: $csr}'
}

# Extract certificate from API response
extract_certificate() {
    local response="$1"
    local cert_chain
    
    if ! cert_chain=$(echo "$response" | jq -r '.certChain' 2>/dev/null); then
        log "ERROR" "Failed to parse API response JSON"
        return 1
    fi

    if [[ ! "$cert_chain" =~ -----BEGIN\ CERTIFICATE----- ]]; then
        log "ERROR" "Invalid or missing certificate in API response"
        return 1
    fi
    
    echo "$cert_chain"
    return 0
}

# ============================================================================
# API FUNCTIONS
# ============================================================================

# Determine API URL based on tortuga package
get_api_url_from_package() {
    local package="$1"
    
    if [[ "$package" == "tortuga-agents-prod" ]]; then
        echo "https://hyperfabric.cisco.com/api/v1/deviceCert"
    elif [[ "$package" == "tortuga-agents" ]]; then
        echo "https://dev.hyperfabric.cisco.com/api/v1/deviceCert"
    else
        log "ERROR" "Unknown tortuga package: $package"
        return 1
    fi
}

# Check if API endpoint is reachable
check_api_health() {
    local api_url="$1"
    local https_proxy="$2"
    local insecure_flag="$3"
    
    # Extract root URL (eg. https://hyperfabric.cisco.com)
    local root_url
    if [[ "$api_url" =~ ^(https?://[^/]+) ]]; then
        root_url="${BASH_REMATCH[1]}"
    else
        log "ERROR" "Invalid API URL format: $api_url"
        return 1
    fi

    local curl_cmd=(curl -s -o /dev/null -w "%{http_code}"
        --connect-timeout 10
        --max-time 30)

    if [[ -n "$https_proxy" ]]; then
        curl_cmd+=(--proxy "$https_proxy")
    fi

    if [[ -n "$insecure_flag" ]]; then
        curl_cmd+=(-k)
    fi

    curl_cmd+=("$root_url")

    local http_code
    if http_code=$("${curl_cmd[@]}" 2>/dev/null); then
        if [[ "$http_code" -ge 200 && "$http_code" -lt 400 ]]; then
            return 0
        else
            log "ERROR" "API health check failed - received HTTP $http_code from $root_url"
            return 1
        fi
    else
        log "ERROR" "Failed to connect to API at: $root_url (no response received)"
        return 1
    fi
}

# Submit CSR to API with retry logic
submit_csr_to_api() {
    local csr="$1"
    local api_url="$2"
    local bearer_token="$3"
    local https_proxy="$4"
    local insecure_flag="$5"
    local context="$6"
    
    local json_payload
    if ! json_payload=$(create_csr_json_payload "$csr"); then
        log "ERROR" "Failed to create JSON payload for $context"
        return 1
    fi
    
    local curl_cmd=(curl -s -w "%{http_code}"
        -X POST
        -H "Authorization: Bearer $bearer_token"
        -H "Content-Type: application/json"
        -d "$json_payload"
        --connect-timeout 60
        --max-time 180)
    
    if [[ -n "$https_proxy" ]]; then
        curl_cmd+=(--proxy "$https_proxy")
    fi
    
    if [[ -n "$insecure_flag" ]]; then
        curl_cmd+=(-k)
    fi
    
    curl_cmd+=("$api_url")
    
    # Retry logic - attempt up to 3 times
    local retry_count=0
    local max_retries=3
    local response=""
    local http_code=""
    
    while [[ $retry_count -lt $max_retries ]]; do
        if response=$("${curl_cmd[@]}" 2>/dev/null); then
            http_code="${response: -3}"
            response="${response%???}"
            
            if [[ "$http_code" -ge 200 && "$http_code" -lt 300 ]]; then
                log "INFO" "API request successful (HTTP $http_code) for $context"
                echo "$response"
                return 0
            else
                log "WARN" "API request failed with HTTP $http_code for $context (attempt $((retry_count + 1))/$max_retries)"
                if [[ -n "$response" ]]; then
                    log "WARN" "Response: $response"
                fi
            fi
        else
            log "WARN" "API request failed for $context (attempt $((retry_count + 1))/$max_retries)"
        fi
        
        retry_count=$((retry_count + 1))
        
        if [[ $retry_count -lt $max_retries ]]; then
            sleep 2
        fi
    done
    
    log "ERROR" "API request failed after $max_retries attempts for $context"
    return 1
}

# ============================================================================
# KUBERNETES CLI DETECTION
# ============================================================================

# Detect or validate the Kubernetes CLI (oc or kubectl).
# If K8S_CLI is already set (via --cli flag), validate it exists.
# Otherwise auto-detect: prefer oc, fall back to kubectl.
detect_k8s_cli() {
    if [[ -n "$K8S_CLI" ]]; then
        if ! command -v "$K8S_CLI" &>/dev/null; then
            log "ERROR" "Specified CLI '$K8S_CLI' not found in PATH"
            return 1
        fi
        log "INFO" "Using specified CLI: $K8S_CLI"
        return 0
    fi

    if command -v oc &>/dev/null; then
        K8S_CLI="oc"
        log "INFO" "Auto-detected CLI: oc"
    elif command -v kubectl &>/dev/null; then
        K8S_CLI="kubectl"
        log "INFO" "Auto-detected CLI: kubectl (oc not found)"
    else
        log "ERROR" "No Kubernetes CLI found in PATH. Install 'oc' or 'kubectl', or use --cli to specify one."
        return 1
    fi
    return 0
}

# ============================================================================
# ENVIRONMENT DETECTION
# ============================================================================

# Detect environment and set HYPERFABRIC_API_URL
detect_environment() {
    local exec_cmd="$1"
    local tortuga_package

    # Query tortuga package from target
    if ! tortuga_package=$($exec_cmd dpkg -l 2>/dev/null | grep tortuga | awk '{print $2}' | head -1); then
        log "ERROR" "Failed to query tortuga package from target"
        return 1
    fi

    if [[ -z "$tortuga_package" ]]; then
        log "ERROR" "No tortuga package found in target"
        return 1
    fi

    # Set API URL based on package name
    if ! HYPERFABRIC_API_URL=$(get_api_url_from_package "$tortuga_package"); then
        return 1
    fi

    # Verify API is reachable
    if ! check_api_health "$HYPERFABRIC_API_URL" "$USE_PROXY" "$INSECURE"; then
        return 1
    fi

    return 0
}

# ============================================================================
# TPM CERTIFICATE PROCESSING
# ============================================================================

# Process TPM certificate for a single target (pod)
process_tpm_certificate() {
    local exec_cmd="$1"
    local exec_stdin_cmd="$2"
    local target_name="$3"
    
    local csr response cert_chain temp_cert_file
    
    # Generate CSR
    log "INFO" "Generating CSR from target: $target_name"
    if ! csr=$($exec_cmd "$DRAKECTL_PATH" -q gencsr "$FORCE_FLAG" 2>/dev/null); then
        log "ERROR" "Failed to generate CSR from target: $target_name"
        return 1
    fi

    # Validate CSR format
    if ! echo "$csr" | openssl req -text -noout > /dev/null 2>&1; then
        log "ERROR" "Invalid CSR format from target: $target_name"
        log "ERROR" "CSR output: $csr"
        return 1
    fi
    
    log "INFO" "CSR generated successfully from target: $target_name"
    
    # Submit CSR to API
    log "INFO" "Submitting CSR to Cisco Hyperfabric API for target: $target_name"
    if ! response=$(submit_csr_to_api "$csr" "$HYPERFABRIC_API_URL" "$CISCO_HF_BEARER_TOKEN" "$USE_PROXY" "$INSECURE" "$target_name"); then
        return 1
    fi
    
    # Extract certificate
    log "INFO" "Extracting certificate from API response for target: $target_name"
    if ! cert_chain=$(extract_certificate "$response"); then
        log "ERROR" "Failed to extract certificate from API response for target: $target_name"
        return 1
    fi
    
    # Import certificate
    log "INFO" "Importing certificate to target: $target_name"
    
    # Create temp file path
    temp_cert_file="/tmp/cert_${target_name}_$(date +%s).pem"
    
    # Write certificate to temp file inside the target using stdin
    if ! echo -e "$cert_chain" | $exec_stdin_cmd sh -c "cat > '$temp_cert_file'" 2>/dev/null; then
        log "ERROR" "Failed to write certificate file to target: $target_name"
        return 1
    fi
    
    # Import certificate using the temp file
    if ! $exec_cmd "$DRAKECTL_PATH" -q import --infile "$temp_cert_file" "$FORCE_FLAG" >/dev/null 2>&1; then
        log "ERROR" "Failed to import certificate to target: $target_name"
        $exec_cmd rm -f "$temp_cert_file" 2>/dev/null || true
        return 1
    fi
    
    # Clean up temp file
    $exec_cmd rm -f "$temp_cert_file" 2>/dev/null || true
    
    log "INFO" "Certificate imported successfully to target: $target_name"
    
    # Verify certificate
    log "INFO" "Verifying certificate installation on target: $target_name"
    if ! $exec_cmd "$DRAKECTL_PATH" -q certcheck >/dev/null 2>&1; then
        log "ERROR" "Certificate verification failed for target: $target_name"
        return 1
    fi
    
    log "INFO" "Certificate verification successful for target: $target_name"
    return 0
}

# ============================================================================
# COMMAND LINE PARSING
# ============================================================================

# Usage information
usage() {
    cat << EOF
Usage: $0 [OPTIONS]

Automate TPM CSR process for cisco-hf-agents daemonset

OPTIONS:
    -f, --force          Force regeneration of existing keys
    -p, --proxy URL      HTTPS proxy for API calls
    -h, --help           Show this help
    -k, --insecure       Skip TLS certificate verification
    -c, --cli [oc|kubectl]
                         Kubernetes CLI to use (default: auto-detect, prefers oc)

ENVIRONMENT:
    CISCO_HF_BEARER_TOKEN    Bearer token for Cisco Hyperfabric API (required)
    DAEMONSET                Daemonset name to process (default: cisco-hf-agents)
    NAMESPACE                Kubernetes namespace (default: cisco-hf)

EXAMPLE:
    export CISCO_HF_BEARER_TOKEN="your-token"
    $0 --proxy http://proxy.example.com:8080
    $0 --cli kubectl

EOF
}

# Parse command-line arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -f|--force)
                FORCE_FLAG="--force"
                shift
                ;;
            -p|--proxy)
                USE_PROXY="$2"
                shift 2
                ;;
            -k|--insecure)
                INSECURE="-k"
                shift
                ;;
            -c|--cli)
                if [[ -z "${2:-}" ]]; then
                    log "ERROR" "--cli requires a value: 'oc' or 'kubectl'"
                    usage
                    exit 1
                fi
                if [[ "$2" != "oc" && "$2" != "kubectl" ]]; then
                    log "WARN" "Unknown CLI '$2' for --cli, falling back to auto-detection"
                else
                    K8S_CLI="$2"
                fi
                shift 2
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            *)
                log "ERROR" "Unknown option: $1"
                usage
                exit 1
                ;;
        esac
    done
}

# ============================================================================
# MAIN SCRIPT
# ============================================================================

main() {
    # Parse arguments
    parse_args "$@"

    # Check prerequisites
    if ! check_bearer_token; then
        exit 1
    fi

    if ! detect_k8s_cli; then
        exit 1
    fi

    if ! check_required_commands curl jq openssl; then
        exit 1
    fi

    log "INFO" "Starting TPM CSR automation for $DAEMONSET in $NAMESPACE"

    # Get pods
    pods=$($K8S_CLI get pods -n "$NAMESPACE" -l app=cisco_hf_agents -o jsonpath='{.items[*].metadata.name}')
    if [[ -z "$pods" ]]; then
        log "ERROR" "No pods found for daemonset $DAEMONSET in namespace $NAMESPACE"
        exit 1
    fi

    read -ra pod_array <<< "$pods"
    total_pods=${#pod_array[@]}
    log "INFO" "Found $total_pods pods to process"

    # Detect environment and set appropriate API endpoint
    first_pod=$($K8S_CLI get pods -n "$NAMESPACE" -l app=cisco_hf_agents --field-selector=status.phase=Running -o jsonpath='{.items[0].metadata.name}')
    if [[ -z "$first_pod" ]]; then
        log "ERROR" "No running pods found to detect environment"
        exit 1
    fi
    
    if ! detect_environment "$K8S_CLI exec -t $first_pod -n $NAMESPACE --"; then
        exit 1
    fi

    log "INFO" "Using API endpoint: $HYPERFABRIC_API_URL"

    # Process each pod
    success_count=0
    failure_count=0

    for pod in "${pod_array[@]}"; do
        current_num=$((success_count + failure_count + 1))
        log "INFO" "Processing pod $current_num/$total_pods: $pod"
     
        # Check pod status
        if ! status=$($K8S_CLI get pod "$pod" -n "$NAMESPACE" -o jsonpath='{.status.phase}' 2>/dev/null); then
            log "ERROR" "Failed to get status for pod: $pod"
            failure_count=$((failure_count + 1))
            continue
        fi

        # Skip CSR generation if cert already present and --force not specified
        if [[ -z "$FORCE_FLAG" ]]; then
            if $K8S_CLI exec -t "$pod" -n "$NAMESPACE" -- "$DRAKECTL_PATH" -q certcheck >/dev/null 2>&1; then
                log "INFO" "Certificate already present in pod: $pod, skipping (use --force to regenerate)"
                success_count=$((success_count + 1))
                continue
            fi
        fi

        if [[ "$status" != "Running" ]]; then
            log "WARN" "Pod $pod is not running (status: $status), skipping"
            failure_count=$((failure_count + 1))
            continue
        fi

        # Process TPM certificate
        # Use subshell to isolate errors and prevent loop termination
        if ! (
            process_tpm_certificate \
                "$K8S_CLI exec -t $pod -n $NAMESPACE --" \
                "$K8S_CLI exec -i $pod -n $NAMESPACE --" \
                "$pod"
        ); then
            failure_count=$((failure_count + 1))
            continue
        fi

        success_count=$((success_count + 1))
        log "INFO" "Pod $current_num/$total_pods completed: $pod"
    done

    # Summary
    log "INFO" "Processing complete - Success: $success_count, Failed: $failure_count"

    if [[ $failure_count -gt 0 ]]; then
        exit 1
    fi

    log "INFO" "All $success_count pods processed successfully!"
}

# Run main function
main "$@"