#!/usr/bin/env bash
# =============================================================================
# setup_custom_tls.sh - Configure custom TLS certificates for Caddy
# =============================================================================
# Generates caddy-addon/custom-tls.conf for using corporate/internal certificates
# instead of Let's Encrypt.
#
# Usage:
#   bash scripts/setup_custom_tls.sh                    # Interactive mode
#   bash scripts/setup_custom_tls.sh cert.crt key.key   # Non-interactive mode
#
# Prerequisites:
#   - Place certificate files in ./certs/ directory
#   - Certificate paths are relative to container (/etc/caddy/certs/)
# =============================================================================

set -euo pipefail

source "$(dirname "$0")/utils.sh" && init_paths

ADDON_FILE="$PROJECT_ROOT/caddy-addon/custom-tls.conf"
CERTS_DIR="$PROJECT_ROOT/certs"

# =============================================================================
# FUNCTIONS
# =============================================================================

show_help() {
    cat << EOF
Setup Custom TLS Certificates for Caddy

Usage: $(basename "$0") [OPTIONS] [CERT_FILE] [KEY_FILE]

Options:
  -h, --help     Show this help message
  --remove       Remove custom TLS configuration

Arguments:
  CERT_FILE      Path to certificate file (relative to ./certs/)
  KEY_FILE       Path to private key file (relative to ./certs/)

Examples:
  $(basename "$0")                           # Interactive mode
  $(basename "$0") wildcard.crt wildcard.key # Use specific files
  $(basename "$0") --remove                  # Remove custom TLS config

The script will:
1. Detect certificate files in ./certs/
2. Read active services from .env
3. Generate caddy-addon/custom-tls.conf
4. Optionally restart Caddy

EOF
}

find_certificates() {
    # Find certificate files in certs directory
    local certs=()
    if [[ -d "$CERTS_DIR" ]]; then
        while IFS= read -r -d '' file; do
            certs+=("$(basename "$file")")
        done < <(find "$CERTS_DIR" -maxdepth 1 -type f \( -name "*.crt" -o -name "*.pem" -o -name "*.cer" \) -print0 2>/dev/null)
    fi
    echo "${certs[*]:-}"
}

find_keys() {
    # Find key files in certs directory
    local keys=()
    if [[ -d "$CERTS_DIR" ]]; then
        while IFS= read -r -d '' file; do
            keys+=("$(basename "$file")")
        done < <(find "$CERTS_DIR" -maxdepth 1 -type f \( -name "*.key" -o -name "*-key.pem" \) -print0 2>/dev/null)
    fi
    echo "${keys[*]:-}"
}

get_active_services() {
    # Get list of services with their hostnames from .env
    load_env
    local services=()

    # Map of service names to their hostname variables
    declare -A service_map=(
        ["n8n"]="N8N_HOSTNAME"
        ["flowise"]="FLOWISE_HOSTNAME"
        ["webui"]="WEBUI_HOSTNAME"
        ["grafana"]="GRAFANA_HOSTNAME"
        ["prometheus"]="PROMETHEUS_HOSTNAME"
        ["portainer"]="PORTAINER_HOSTNAME"
        ["langfuse"]="LANGFUSE_HOSTNAME"
        ["supabase"]="SUPABASE_HOSTNAME"
        ["dify"]="DIFY_HOSTNAME"
        ["nocodb"]="NOCODB_HOSTNAME"
        ["ragapp"]="RAGAPP_HOSTNAME"
        ["ragflow"]="RAGFLOW_HOSTNAME"
        ["waha"]="WAHA_HOSTNAME"
        ["searxng"]="SEARXNG_HOSTNAME"
        ["comfyui"]="COMFYUI_HOSTNAME"
        ["welcome"]="WELCOME_HOSTNAME"
        ["databasus"]="DATABASUS_HOSTNAME"
        ["letta"]="LETTA_HOSTNAME"
        ["lightrag"]="LIGHTRAG_HOSTNAME"
        ["weaviate"]="WEAVIATE_HOSTNAME"
        ["qdrant"]="QDRANT_HOSTNAME"
        ["neo4j"]="NEO4J_HOSTNAME"
        ["postiz"]="POSTIZ_HOSTNAME"
        ["libretranslate"]="LT_HOSTNAME"
        ["paddleocr"]="PADDLEOCR_HOSTNAME"
        ["docling"]="DOCLING_HOSTNAME"
    )

    for service in "${!service_map[@]}"; do
        local hostname_var="${service_map[$service]}"
        local hostname="${!hostname_var:-}"
        if [[ -n "$hostname" && "$hostname" != *"yourdomain.com" ]]; then
            services+=("$service:$hostname")
        fi
    done

    echo "${services[*]:-}"
}

generate_config() {
    local cert_file="$1"
    local key_file="$2"
    local services=("${@:3}")

    cat > "$ADDON_FILE" << 'HEADER'
# Custom TLS Configuration
# Generated by setup_custom_tls.sh
#
# This file overrides default Let's Encrypt certificates with custom ones.
# Regenerate with: make setup-tls

# Reusable TLS snippet
(custom_tls) {
HEADER

    echo "    tls /etc/caddy/certs/$cert_file /etc/caddy/certs/$key_file" >> "$ADDON_FILE"
    echo "}" >> "$ADDON_FILE"
    echo "" >> "$ADDON_FILE"

    # Service-specific reverse proxy mappings
    declare -A proxy_map=(
        ["n8n"]="n8n:5678"
        ["flowise"]="flowise:3001"
        ["webui"]="open-webui:8080"
        ["grafana"]="grafana:3000"
        ["prometheus"]="prometheus:9090"
        ["portainer"]="portainer:9000"
        ["langfuse"]="langfuse-web:3000"
        ["supabase"]="kong:8000"
        ["dify"]="nginx:80"
        ["nocodb"]="nocodb:8080"
        ["ragapp"]="ragapp:8000"
        ["ragflow"]="ragflow:80"
        ["waha"]="waha:3000"
        ["searxng"]="searxng:8080"
        ["comfyui"]="comfyui:8188"
        ["welcome"]="file_server"
        ["databasus"]="databasus:4005"
        ["letta"]="letta:8283"
        ["lightrag"]="lightrag:9621"
        ["weaviate"]="weaviate:8080"
        ["qdrant"]="qdrant:6333"
        ["neo4j"]="neo4j:7474"
        ["postiz"]="postiz:5000"
        ["libretranslate"]="libretranslate:5000"
        ["paddleocr"]="paddleocr:8080"
        ["docling"]="docling:5001"
    )

    # Services that need basic auth (format: USERNAME_VAR:PASSWORD_HASH_VAR)
    declare -A auth_services=(
        ["prometheus"]="PROMETHEUS_USERNAME:PROMETHEUS_PASSWORD_HASH"
        ["ragapp"]="RAGAPP_USERNAME:RAGAPP_PASSWORD_HASH"
        ["comfyui"]="COMFYUI_USERNAME:COMFYUI_PASSWORD_HASH"
        ["welcome"]="WELCOME_USERNAME:WELCOME_PASSWORD_HASH"
        ["libretranslate"]="LT_USERNAME:LT_PASSWORD_HASH"
        ["paddleocr"]="PADDLEOCR_USERNAME:PADDLEOCR_PASSWORD_HASH"
        ["docling"]="DOCLING_USERNAME:DOCLING_PASSWORD_HASH"
    )

    for service_entry in "${services[@]}"; do
        local service="${service_entry%%:*}"
        local hostname="${service_entry#*:}"
        local proxy="${proxy_map[$service]:-}"

        [[ -z "$proxy" ]] && continue

        echo "# $service" >> "$ADDON_FILE"
        echo "$hostname {" >> "$ADDON_FILE"
        echo "    import custom_tls" >> "$ADDON_FILE"

        # Add basic auth if needed
        if [[ -n "${auth_services[$service]:-}" ]]; then
            local auth_config="${auth_services[$service]}"
            local username_var="${auth_config%%:*}"
            local password_hash_var="${auth_config#*:}"
            echo "    basic_auth {" >> "$ADDON_FILE"
            echo "        {\$${username_var}} {\$${password_hash_var}}" >> "$ADDON_FILE"
            echo "    }" >> "$ADDON_FILE"
        fi

        # Add reverse proxy or file server
        if [[ "$proxy" == "file_server" ]]; then
            echo "    root * /srv/welcome" >> "$ADDON_FILE"
            echo "    file_server" >> "$ADDON_FILE"
            echo "    try_files {path} /index.html" >> "$ADDON_FILE"
        else
            echo "    reverse_proxy $proxy" >> "$ADDON_FILE"
        fi

        echo "}" >> "$ADDON_FILE"
        echo "" >> "$ADDON_FILE"
    done

    log_success "Generated $ADDON_FILE"
}

remove_config() {
    if [[ -f "$ADDON_FILE" ]]; then
        rm -f "$ADDON_FILE"
        log_success "Removed custom TLS configuration"
    else
        log_info "No custom TLS configuration found"
    fi
}

restart_caddy() {
    if wt_yesno "Restart Caddy" "Do you want to restart Caddy to apply the new configuration?" "yes"; then
        log_info "Restarting Caddy..."
        docker compose -p localai restart caddy
        log_success "Caddy restarted"
    else
        log_info "Skipped Caddy restart. Run manually: docker compose -p localai restart caddy"
    fi
}

# =============================================================================
# MAIN
# =============================================================================

main() {
    # Handle arguments
    case "${1:-}" in
        -h|--help)
            show_help
            exit 0
            ;;
        --remove)
            remove_config
            restart_caddy
            exit 0
            ;;
    esac

    # Ensure certs directory exists
    mkdir -p "$CERTS_DIR"

    local cert_file=""
    local key_file=""

    # Non-interactive mode
    if [[ $# -ge 2 ]]; then
        cert_file="$1"
        key_file="$2"

        if [[ ! -f "$CERTS_DIR/$cert_file" ]]; then
            log_error "Certificate not found: $CERTS_DIR/$cert_file"
            exit 1
        fi
        if [[ ! -f "$CERTS_DIR/$key_file" ]]; then
            log_error "Key not found: $CERTS_DIR/$key_file"
            exit 1
        fi
    else
        # Interactive mode
        require_whiptail

        # Find available certificates
        local certs_arr
        IFS=' ' read -ra certs_arr <<< "$(find_certificates)"

        if [[ ${#certs_arr[@]} -eq 0 ]]; then
            wt_msg "No Certificates Found" "No certificate files found in ./certs/\n\nPlease place your certificate (.crt, .pem, .cer) and key (.key) files in the certs/ directory first."
            exit 1
        fi

        # Build menu items for certificates
        local cert_items=()
        for cert in "${certs_arr[@]}"; do
            cert_items+=("$cert" "")
        done

        cert_file=$(wt_menu "Select Certificate" "Choose your TLS certificate file:" "${cert_items[@]}")
        [[ -z "$cert_file" ]] && exit 1

        # Find available keys
        local keys_arr
        IFS=' ' read -ra keys_arr <<< "$(find_keys)"

        if [[ ${#keys_arr[@]} -eq 0 ]]; then
            wt_msg "No Keys Found" "No key files found in ./certs/\n\nPlease place your private key (.key) file in the certs/ directory."
            exit 1
        fi

        # Build menu items for keys
        local key_items=()
        for key in "${keys_arr[@]}"; do
            key_items+=("$key" "")
        done

        key_file=$(wt_menu "Select Private Key" "Choose your TLS private key file:" "${key_items[@]}")
        [[ -z "$key_file" ]] && exit 1
    fi

    log_info "Using certificate: $cert_file"
    log_info "Using key: $key_file"

    # Get active services
    local services_arr
    IFS=' ' read -ra services_arr <<< "$(get_active_services)"

    if [[ ${#services_arr[@]} -eq 0 ]]; then
        log_warning "No services with configured hostnames found in .env"
        log_info "Make sure to update *_HOSTNAME variables in .env with your domain"
        exit 1
    fi

    log_info "Found ${#services_arr[@]} services with configured hostnames"

    # Generate configuration
    generate_config "$cert_file" "$key_file" "${services_arr[@]}"

    # Show summary
    echo ""
    log_info "Configuration generated for the following services:"
    for service_entry in "${services_arr[@]}"; do
        local service="${service_entry%%:*}"
        local hostname="${service_entry#*:}"
        echo "  - $service: $hostname"
    done
    echo ""

    # Restart Caddy
    restart_caddy
}

main "$@"