#!/bin/bash

# Standalone TPM Certificate Signing Request automation script for a single server
# running cisco_hf_agent container in Docker, Podman, or containerd (ctr)
# This script has no external dependencies - all functions are included inline

set -euo pipefail

# Configuration
CONTAINER_NAME="${CONTAINER_NAME:-cisco_hf_agent}"
CONTAINER_RUNTIME="${CONTAINER_RUNTIME:-}"

# Runtime type (set by init_container_runtime)
RUNTIME_TYPE=""

# Constants
DRAKECTL_PATH="/opt/cisco/bin/tortuga/drakectl"

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

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

# 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
}

# 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
}

# Determine API URL based on tortuga package or controller URL
get_api_url_from_package() {
    local package="$1"
    
    # If controller URL is specified, use it and add the API path
    if [[ -n "$CONTROLLER_URL" ]]; then
        # Remove trailing slash if present and add API path
        local controller_url="${CONTROLLER_URL%/}"
        echo "${controller_url}/api/v1/deviceCert"
        return 0
    fi
    
    # Default behavior based on tortuga package
    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"
    
    # For controller URLs, test the base URL (without /api/v1/deviceCert)
    # For default URLs, test the root URL as before
    local test_url="$api_url"
    
    # If using default tortuga URLs, extract root URL for health check
    if [[ "$api_url" == *"hyperfabric.cisco.com"* ]] || [[ "$api_url" == *"dev.hyperfabric.cisco.com"* ]]; then
        if [[ "$api_url" =~ ^(https?://[^/]+) ]]; then
            test_url="${BASH_REMATCH[1]}"
        else
            log "ERROR" "Invalid API URL format: $api_url"
            return 1
        fi
    else
        # For custom controller URLs, extract base URL (remove /api/v1/deviceCert)
        if [[ "$api_url" =~ ^(https?://[^/]+) ]]; then
            test_url="${BASH_REMATCH[1]}"
        else
            log "ERROR" "Invalid API URL format: $api_url"
            return 1
        fi
    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+=("$test_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 $test_url"
            return 1
        fi
    else
        log "ERROR" "Failed to connect to API at: $test_url (no response received)"
        return 1
    fi
}

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

    # If controller URL is specified, use it directly
    if [[ -n "$CONTROLLER_URL" ]]; then
        log "INFO" "Using specified controller URL: $CONTROLLER_URL"
        export HYPERFABRIC_API_URL=$(get_api_url_from_package "")
        
        # Auto-set insecure flag for custom controllers
        if [[ -z "$INSECURE" ]]; then
            export INSECURE="-k"
            log "INFO" "Auto-setting --insecure flag for custom controller URL"
        fi
        
        # Verify API is reachable
        if ! check_api_health "$HYPERFABRIC_API_URL" "$USE_PROXY" "$INSECURE"; then
            return 1
        fi
        return 0
    fi

    # 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

    export HYPERFABRIC_API_URL

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

    return 0
}

# 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
}

# Parse command-line arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -f|--force)
                export FORCE_FLAG="--force"
                shift
                ;;
            -p|--proxy)
                export USE_PROXY="$2"
                shift 2
                ;;
            -k|--insecure)
                export INSECURE="-k"
                shift
                ;;
            -u|--controller-url)
                export CONTROLLER_URL="$2"
                shift 2
                ;;
            -h|--help)
                return 2  # Special return code to indicate help requested
                ;;
            *)
                log "ERROR" "Unknown option: $1"
                return 1
                ;;
        esac
    done
    
    return 0
}

# Process TPM certificate for a single target
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; 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"
    
    # Restart drake to apply the new certificate
    log "INFO" "Restarting drake service on target: $target_name"
    if ! $exec_cmd supervisorctl restart drake >/dev/null 2>&1; then
        log "WARN" "Failed to restart drake service (this may be expected in some environments)"
    else
        log "INFO" "Drake service restarted successfully"
    fi
    
    return 0
}

# Validate and initialize container runtime
init_container_runtime() {
    # If CONTAINER_RUNTIME is explicitly set, use it
    if [[ -n "$CONTAINER_RUNTIME" ]]; then
        if ! command -v "$CONTAINER_RUNTIME" &> /dev/null; then
            log "ERROR" "Specified container runtime '$CONTAINER_RUNTIME' not found in PATH"
            exit 1
        fi
        log "INFO" "Using specified container runtime: $CONTAINER_RUNTIME"
    else
        # Auto-detect available container runtimes in order of preference
        local runtimes=("docker" "podman" "ctr")
        local found=false
        
        for runtime in "${runtimes[@]}"; do
            if command -v "$runtime" &> /dev/null; then
                CONTAINER_RUNTIME="$runtime"
                log "INFO" "Auto-detected container runtime: $runtime"
                found=true
                break
            fi
        done
        
        if [[ "$found" == false ]]; then
            log "ERROR" "No supported container runtime found in PATH"
            log "ERROR" "Please install one of: docker, podman, ctr"
            exit 1
        fi
    fi

    local runtime_bin
    runtime_bin=$(basename "$CONTAINER_RUNTIME")

    case "$runtime_bin" in
        docker)  RUNTIME_TYPE="docker" ;;
        podman)  RUNTIME_TYPE="podman" ;;
        ctr)     RUNTIME_TYPE="containerd" ;;
        *)
            log "ERROR" "Unsupported container runtime: $CONTAINER_RUNTIME"
            exit 1
            ;;
    esac
}

# Check if container is running (runtime-aware)
check_container_running() {
    case "$RUNTIME_TYPE" in
        "docker"|"podman")
            if ! $CONTAINER_RUNTIME inspect "$CONTAINER_NAME" &> /dev/null; then
                log "ERROR" "Container '$CONTAINER_NAME' not found"
                return 1
            fi
            local status
            status=$($CONTAINER_RUNTIME inspect "$CONTAINER_NAME" --format '{{.State.Status}}')
            if [[ "$status" != "running" ]]; then
                log "ERROR" "Container '$CONTAINER_NAME' is not running (status: $status)"
                return 1
            fi
            ;;
        "containerd")
            if ! $CONTAINER_RUNTIME tasks ls 2>/dev/null | grep -q "$CONTAINER_NAME.*RUNNING"; then
                log "ERROR" "Container '$CONTAINER_NAME' is not running (check: ctr tasks ls)"
                return 1
            fi
            ;;
    esac
    return 0
}

# Exec wrapper for regular commands
container_exec() {
    case "$RUNTIME_TYPE" in
        "docker"|"podman")
            $CONTAINER_RUNTIME exec "$CONTAINER_NAME" "$@"
            ;;
        "containerd")
            $CONTAINER_RUNTIME tasks exec --exec-id "exec-$$-$(date +%s%N)" "$CONTAINER_NAME" "$@"
            ;;
    esac
}

# Exec wrapper for stdin operations
container_exec_stdin() {
    case "$RUNTIME_TYPE" in
        "docker"|"podman")
            $CONTAINER_RUNTIME exec -i "$CONTAINER_NAME" "$@"
            ;;
        "containerd")
            $CONTAINER_RUNTIME tasks exec --exec-id "cp-$$-$(date +%s%N)" "$CONTAINER_NAME" "$@"
            ;;
    esac
}

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

Automate TPM CSR process for cisco_hf_agent container on a single server
(Standalone version with no external dependencies)

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
    -u, --controller-url Controller API URL

ENVIRONMENT VARIABLES:
    CISCO_HF_BEARER_TOKEN    Bearer token for Cisco Hyperfabric API (required)
    CONTAINER_NAME           Container name (default: cisco_hf_agent)
    CONTAINER_RUNTIME        Container runtime (auto-detected if not specified)
                             Supported: docker, podman, ctr
                             Detection order: docker -> podman -> ctr

EXAMPLES:
    # Auto-detect container runtime:
    export CISCO_HF_BEARER_TOKEN="your-token"
    $0

    # Force specific runtime (overrides auto-detection):
    export CISCO_HF_BEARER_TOKEN="your-token"
    export CONTAINER_RUNTIME="podman"
    $0

    # Using custom controller URL:
    export CISCO_HF_BEARER_TOKEN="your-token"
    $0 --controller-url https://custom-controller.example.com

    # Using containerd (ctr):
    export CISCO_HF_BEARER_TOKEN="your-token"
    export CONTAINER_RUNTIME="ctr"
    $0

EOF
}

# Parse arguments
if ! parse_args "$@"; then
    case $? in
        2) usage; exit 0 ;;  # Help requested
        *) usage; exit 1 ;;  # Parse error
    esac
fi

if ! check_bearer_token; then
    exit 1
fi

if ! check_required_commands curl jq; then
    exit 1
fi

# Initialize and validate container runtime
init_container_runtime

log "INFO" "Starting TPM CSR automation for container: $CONTAINER_NAME"
log "INFO" "Using container runtime: $CONTAINER_RUNTIME ($RUNTIME_TYPE)"

# Check if container exists and is running
if ! check_container_running; then
    exit 1
fi

log "INFO" "Container '$CONTAINER_NAME' is running"

# Detect environment and set HYPERFABRIC_API_URL
if ! detect_environment "container_exec"; then
    exit 1
fi

# Skip CSR generation if cert already present and --force not specified
if [[ -z "$FORCE_FLAG" ]]; then
    if container_exec "$DRAKECTL_PATH" -q certcheck >/dev/null 2>&1; then
        log "INFO" "Certificate already present in container, skipping (use --force to regenerate)"
        log "INFO" "TPM CSR automation completed successfully (certificate already exists)"
        exit 0
    fi
fi

# Process TPM certificate
if ! process_tpm_certificate \
    "container_exec" \
    "container_exec_stdin" \
    "$CONTAINER_NAME"; then
    log "ERROR" "TPM certificate processing failed for container: $CONTAINER_NAME"
    exit 1
fi

log "INFO" "TPM CSR automation completed successfully!"
exit 0