#!/bin/sh

printHelp() {
    echo "Mule Troubleshooting Tool"
    echo "========================="
    echo
    echo "Usage: ./diag [options] [command] [command-options]"
    echo
    echo "Commands:"
    echo "  diaf                    Generate a complete Mule diagnostic dump (default)"
    echo "  help                    Show this help message"
    echo "  <operation-name>        Execute a specific troubleshooting operation"
    echo
    echo "Global Options:"
    echo "  --stdout                Output the diagnostic dump to standard output"
    echo "  --output <path>         Specify custom output directory or file path"
    echo "  --extended              Enable extended mode (includes heap dump)"
    echo "  --debug                 Enable debug mode with remote debugging on port 5005"
    echo
    echo "Examples:"
    echo "  ./diag                          # Generate diagnostic dump to logs directory"
    echo "  ./diag --stdout                 # Output diagnostic dump to stdout"
    echo "  ./diag --extended               # Include heap dump in diagnostic"
    echo "  ./diag --output /tmp/mule       # Save to specific file (see Archive Format below)"
    echo "  ./diag --output /tmp/           # Save to specific directory"
    echo "  ./diag <operation-name>         # Execute specific operation"
    echo
    echo "Output:"
    echo "  The tool creates an archive file containing:"
    echo "  - mule_dump_<timestamp>.diaf    # Diagnostic information"
    echo "  - thread_dump_<timestamp>.txt   # Thread dump"
    echo "  - heap_dump_<timestamp>.hprof   # Heap dump (if --extended is used)"
    echo "  - tls_support/                  # TLS introspection results (if available)"
    echo
    echo "  Archive Format:"
    echo "  - Uses ZIP format when 'zip' command is available"
    echo "  - Falls back to TAR.GZ format when only 'tar' is available"
}

# Check for help option
if [ "$1" = "--help" ] || [ "$1" = "-h" ] || [ "$1" = "help" ]; then
    printHelp
    exit 0
fi

#===============================  IMPORTANT  ===============================
# Avoid echoing anything to stdout, as it must remain clean for diag output.
# ==========================================================================

TEMP_ARGS="$@"

# Parse command line arguments to find --output parameter
OUTPUT_DIR=""
USE_STDOUT=false
while [ $# -gt 0 ]; do
    case "$1" in
        --stdout)
            USE_STDOUT=true
            shift
            break
            ;;
        --output)
            shift
            OUTPUT_DIR="$1"
            shift
            break
            ;;
        --output=*)
            OUTPUT_DIR="${1#--output=}"
            break
            ;;
    esac
    shift
done

# If no output directory specified, use default logs directory
if [ -z "$OUTPUT_DIR" ]; then
    # Use the same default as Java code
    SCRIPT_DIR="$(dirname "$0")"
    OUTPUT_DIR="$SCRIPT_DIR/../logs"
fi

# Determine output strategy based on intended output type
OUTPUT_IS_ZIP=false
ACTUAL_OUTPUT_DIR=""
TARGET_ZIP_PATH=""

if [ "$USE_STDOUT" = "true" ]; then
    # For stdout, always use temporary directory and zip to stdout
    OUTPUT_IS_ZIP=true
    ACTUAL_OUTPUT_DIR=$(mktemp -d)
    echo "Using temporary directory for stdout mode: $ACTUAL_OUTPUT_DIR" >&2
elif [ "${OUTPUT_DIR%/}" != "$OUTPUT_DIR" ]; then
    # OUTPUT_DIR ends with /, treat as directory - use directly
    ACTUAL_OUTPUT_DIR="${OUTPUT_DIR%/}"
    OUTPUT_IS_ZIP=false
    echo "Using directory output mode: $ACTUAL_OUTPUT_DIR" >&2
elif [ -d "$OUTPUT_DIR" ]; then
    # OUTPUT_DIR is an existing directory - use directly
    ACTUAL_OUTPUT_DIR="$OUTPUT_DIR"
    OUTPUT_IS_ZIP=false
    echo "Using existing directory: $ACTUAL_OUTPUT_DIR" >&2
else
    # OUTPUT_DIR is a zip file path - use temporary directory
    OUTPUT_IS_ZIP=true
    ACTUAL_OUTPUT_DIR=$(mktemp -d)
    TARGET_ZIP_PATH="$OUTPUT_DIR"
    echo "Using temporary directory for zip output: $ACTUAL_OUTPUT_DIR" >&2
    echo "Target zip file: $TARGET_ZIP_PATH" >&2
fi

# Create actual output directory if it doesn't exist
if [ ! -d "$ACTUAL_OUTPUT_DIR" ]; then
    echo "Creating output directory: $ACTUAL_OUTPUT_DIR" >&2
    mkdir -p "$ACTUAL_OUTPUT_DIR" 2>/dev/null || {
        echo "Error: Failed to create output directory $ACTUAL_OUTPUT_DIR. Skipping diag." >&2
        return 0
    }
fi

# Also ensure parent directory of zip file exists if needed
if [ "$OUTPUT_IS_ZIP" = "true" ] && [ -n "$TARGET_ZIP_PATH" ]; then
    TARGET_ZIP_DIR="$(dirname "$TARGET_ZIP_PATH")"
    if [ ! -d "$TARGET_ZIP_DIR" ]; then
        echo "Creating target zip directory: $TARGET_ZIP_DIR" >&2
        mkdir -p "$TARGET_ZIP_DIR" 2>/dev/null || {
            echo "Error: Failed to create target zip directory $TARGET_ZIP_DIR. Skipping diag." >&2
            return 0
        }
    fi
fi

# Set default MULE_BASE if not provided
if [ -z "$MULE_BASE" ]; then
    MULE_BASE="$(dirname "$0")/.."
    echo "MULE_BASE not set, using default: $MULE_BASE" >&2
fi
    
processTlsIntrospectionFiles() {
    # Extract operation name (first non-option argument)
    operationName=""
    for arg in "$@"; do
        case "$arg" in
            --*)
                # Skip options
                ;;
            *)
                # This is the operation name
                operationName="$arg"
                break
                ;;
        esac
    done
    
    # Only execute if operation name is "tls" or not set
    if [ -n "$operationName" ] && [ "$operationName" != "tls" ]; then
        return 0
    fi
    
    # Define the introspection directory
    INTROSPECTION_DIR="$MULE_BASE/.mule/.introspection"
    
    # Check if introspection directory exists
    if [ ! -d "$INTROSPECTION_DIR" ]; then
        echo "Warning: Introspection directory $INTROSPECTION_DIR does not exist. Skipping TLS processing." >&2
        return 0
    fi
    
    # Get the script directory and TLS script path
    SCRIPT_DIR="$(dirname "$0")"
    TLS_SCRIPT="$SCRIPT_DIR/tls/testTlsVersionSupport.sh"
    
    # Check if TLS testing script exists
    if [ ! -f "$TLS_SCRIPT" ]; then
        echo "Warning: TLS testing script not found at $TLS_SCRIPT. Skipping TLS processing." >&2
        return 0
    fi
    
    # Make sure TLS script is executable
    chmod +x "$TLS_SCRIPT" 2>/dev/null || true
    
    # Check if there are CSV files to process before proceeding
    CSV_FILES_EXIST=false
    for csv_check in "$INTROSPECTION_DIR"/*.csv; do
        if [ -f "$csv_check" ]; then
            CSV_FILES_EXIST=true
            break
        fi
    done
    
    if [ "$CSV_FILES_EXIST" = false ]; then
        echo "No CSV files found in $INTROSPECTION_DIR. Skipping TLS processing." >&2
        return 0
    fi
    
    # Create tls_support subdirectory in actual output folder
    TLS_OUTPUT_DIR="$ACTUAL_OUTPUT_DIR/tls_support"
    if [ ! -d "$TLS_OUTPUT_DIR" ]; then
        echo "Creating TLS output directory: $TLS_OUTPUT_DIR" >&2
        mkdir -p "$TLS_OUTPUT_DIR" 2>/dev/null || {
            echo "Error: Failed to create TLS output directory $TLS_OUTPUT_DIR. Skipping TLS processing." >&2
            return 0
        }
    fi
    
    # Find and process CSV files in the introspection directory
    echo "Processing TLS introspection files from $INTROSPECTION_DIR..." >&2
    echo "TLS results will be saved to: $TLS_OUTPUT_DIR" >&2
    
    for csv_file in "$INTROSPECTION_DIR"/*.csv; do
        # Check if the glob didn't match any files (should not happen due to pre-check)
        if [ ! -f "$csv_file" ]; then
            continue
        fi
        
        # Extract filename without path and extension for output naming
        csv_basename=$(basename "$csv_file" .csv)
        output_file="$TLS_OUTPUT_DIR/${csv_basename}_tls_results.csv"
        
        echo "Processing TLS testing for: $csv_file" >&2
        echo "Output will be saved to: $output_file" >&2
        
        # Call the TLS testing script with input and output files
        if "$TLS_SCRIPT" "$csv_file" "$output_file" >&2; then
            echo "Successfully processed TLS testing for $csv_file" >&2
        else
            echo "Error: Failed to process TLS testing for $csv_file" >&2
        fi
        
        echo "" >&2
    done
    
    echo "Completed processing all TLS introspection files." >&2
}

muleTroubleshooting() {
    local actual_output_dir="$1"
    shift  # Remove the first argument (ACTUAL_OUTPUT_DIR) from $@

    # Use jvm from JAVA_HOME if present

    if [ -z "$JAVA_HOME" ]; then
        # Determine it from $PATH
        JAVA_PATH=`command -v java 2>/dev/null`
        if [ -n "$JAVA_PATH" ]; then
            if [ `echo "$JAVA_PATH" | grep "/bin/java$"` ]; then
                JAVA_HOME=$(readlink -f $(command -v java) | sed 's|/bin/java||')
                export JAVA_HOME
            fi
        fi
        
        # OSX always places Java in the same location so we can reliably set JAVA_HOME
        if [ "$DIST_OS" = "macosx" ]
        then
            if [ -x /usr/libexec/java_home ]; then
                JAVA_HOME=`/usr/libexec/java_home`; export JAVA_HOME
            else
                JAVA_HOME="/Library/Java/Home"; export JAVA_HOME
            fi
        fi
    fi

    if [ -z "$JAVA_HOME" ]
    then
        echo "**********************************************************************"
        echo "* ERROR: java could not be found. Please install the JRE or JDK.     *"
        echo "* If it is already installed, configure the system variables PATH or *"
        echo "* JAVA_HOME appropriately.                                           *"
        echo "**********************************************************************"
        exit 1
    fi

    echo "Found JAVA_HOME: $JAVA_HOME" >&2

    JAVA_CMD="$JAVA_HOME/bin/java"
    SCRIPT_DIR="$(dirname "$0")"
    JAR_PATH="$SCRIPT_DIR/mule-troubleshooting-plugin-client-1.0.1-jar-with-dependencies.jar"
    MAIN_CLASS="com.mulesoft.mule.troubleshooting.client.Client"
    LOGS_DIR_SYSTEM_PROPERTY="-Ddiag.logs.dir=$SCRIPT_DIR/../logs"

    # Rebuild arguments for Java application, filtering out --output from original args
    FILTERED_ARGS=""
    set -- $TEMP_ARGS
    while [ $# -gt 0 ]; do
        case "$1" in
            --stdout)
                # Skip --stdout (already processed)
                shift
                ;;
            --output)
                # Skip --output and its value (already processed)
                shift
                shift
                ;;
            --output=*)
                # Skip --output=value (already processed)
                shift
                ;;
            *)
                # Keep all other arguments
                FILTERED_ARGS="$FILTERED_ARGS $1"
                shift
                ;;
        esac
    done

    JAVA_ARGS="--output $actual_output_dir $FILTERED_ARGS"

    echo "Calling mule-troubleshooting-plugin-client with args: $JAVA_ARGS" >&2

    # Check for debug option in processed arguments
    if echo "$JAVA_ARGS" | grep -q '\--debug'; then
        # Remove --debug from JAVA_ARGS for execution
        JAVA_ARGS=$(echo "$JAVA_ARGS" | sed 's/--debug[ ]*//g')
        eval "$JAVA_CMD \"$LOGS_DIR_SYSTEM_PROPERTY\" -agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=5005 -cp \"$JAR_PATH\" \"$MAIN_CLASS\" \"$SCRIPT_DIR\" $JAVA_ARGS"
    else
        eval "$JAVA_CMD \"$LOGS_DIR_SYSTEM_PROPERTY\" -cp \"$JAR_PATH\" \"$MAIN_CLASS\" \"$SCRIPT_DIR\" $JAVA_ARGS"
    fi
}

# Process TLS introspection files before running main diagnostics
processTlsIntrospectionFiles "$@"

# Call muleTroubleshooting function with actual output directory and original arguments
muleTroubleshooting "$ACTUAL_OUTPUT_DIR" "$@"

# Handle zipping if needed
if [ "$OUTPUT_IS_ZIP" = "true" ]; then
    echo "Creating archive from generated files..." >&2
    
    # Create archive file from all files in the output directory
    createArchiveFromDirectory() {
        local source_dir="$1"
        local archive_path="$2"
        local use_stdout="$3"
        
        # Check if zip is available, otherwise use tar
        if command -v zip >/dev/null 2>&1; then
            # Use zip
            if [ "$use_stdout" = "true" ]; then
                # Create zip and output to stdout
                (cd "$source_dir" && zip -r - . 2>/dev/null)
            else
                # Create zip file at specified location
                # Use absolute path for zip file to avoid relative path issues
                local abs_zip_path
                if [ "${archive_path#/}" = "$archive_path" ]; then
                    # Path is relative, make it absolute
                    abs_zip_path="$(pwd)/$archive_path"
                else
                    # Path is already absolute
                    abs_zip_path="$archive_path"
                fi
                (cd "$source_dir" && zip -r "$abs_zip_path" . > /dev/null 2>&1)
                if [ $? -eq 0 ]; then
                    echo "Diagnostic zip file created: $abs_zip_path" >&2
                else
                    echo "Error: Failed to create zip file $abs_zip_path" >&2
                    return 1
                fi
            fi
        elif command -v tar >/dev/null 2>&1; then
            # Use tar as fallback
            if [ "$use_stdout" = "true" ]; then
                # Create tar.gz and output to stdout
                (cd "$source_dir" && tar -czf - . 2>/dev/null)
            else
                # Create tar.gz file at specified location
                local tar_path="$archive_path"
                
                # Use absolute path for tar file to avoid relative path issues
                local abs_tar_path
                if [ "${tar_path#/}" = "$tar_path" ]; then
                    # Path is relative, make it absolute
                    abs_tar_path="$(pwd)/$tar_path"
                else
                    # Path is already absolute
                    abs_tar_path="$tar_path"
                fi
                (cd "$source_dir" && tar -czf "$abs_tar_path" . 2>/dev/null)
                if [ $? -eq 0 ]; then
                    echo "Diagnostic tar.gz file created: $abs_tar_path" >&2
                else
                    echo "Error: Failed to create tar.gz file $abs_tar_path" >&2
                    return 1
                fi
            fi
        else
            echo "Error: Neither zip nor tar command is available for creating archive" >&2
            return 1
        fi
    }
    
    if [ "$USE_STDOUT" = "true" ]; then
        createArchiveFromDirectory "$ACTUAL_OUTPUT_DIR" "" true
        echo "Diagnostic archive generated through standard output" >&2
    else
        createArchiveFromDirectory "$ACTUAL_OUTPUT_DIR" "$TARGET_ZIP_PATH" false
    fi
    
    # Clean up temporary directory if it was used
    if [ "$USE_STDOUT" = "true" ] || [ -n "$TARGET_ZIP_PATH" ]; then
        echo "Cleaning up temporary directory: $ACTUAL_OUTPUT_DIR" >&2
        rm -rf "$ACTUAL_OUTPUT_DIR"
    fi
else
    echo "Files generated in directory: $ACTUAL_OUTPUT_DIR" >&2
fi
