#!/usr/bin/env bash

# Get the directory containing this script (this is src/pash or the installed package dir)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# PASH_TOP is the directory containing compiler/, runtime/, etc.
# Since pa.sh is now in src/pash/, SCRIPT_DIR is already the correct PASH_TOP
if [ -n "${PASH_TOP:-}" ] && [ -d "$PASH_TOP/compiler" ] && [ -d "$PASH_TOP/runtime" ]; then
    # PASH_TOP is already set and valid, keep it
    export PASH_TOP
else
    # Use script location as PASH_TOP
    export PASH_TOP="$SCRIPT_DIR"
fi
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
export RUNTIME_DIR="$PASH_TOP/jit_runtime"
export WRAPPER_LIB_DIR="$RUNTIME_DIR/../wrapper_library/"
export RUNTIME_LIBRARY_DIR="$PASH_TOP/runtime"

## Register the signal handlers
trap kill_all SIGTERM SIGINT

kill_all() {
    kill -s SIGKILL 0
    kill -s SIGKILL "$daemon_pid"
}

## Save the umask to first create some files and then revert it
old_umask=$(umask)
umask u=rwx,g=rx,o=rx

## Get bash version for pash
export PASH_BASH_VERSION="${BASH_VERSINFO[@]:0:3}"

## Create temporary directory and communication FIFOs
export PASH_TMP_PREFIX="$(mktemp -d /tmp/pash_XXXXXXX)/"
export PASH_TIMESTAMP="$(date +"%y-%m-%d-%T")"
export RUNTIME_IN_FIFO="${PASH_TMP_PREFIX}/runtime_in_fifo"
export RUNTIME_OUT_FIFO="${PASH_TMP_PREFIX}/runtime_out_fifo"
rm -f "$RUNTIME_IN_FIFO" "$RUNTIME_OUT_FIFO"
mkfifo "$RUNTIME_IN_FIFO" "$RUNTIME_OUT_FIFO"
export DAEMON_SOCKET="${PASH_TMP_PREFIX}/daemon_socket"
export DSPASH_SOCKET="${PASH_TMP_PREFIX}/dspash_socket"

## Default values for exported flags (used by runtime scripts)
export pash_output_time_flag=1
export pash_execute_flag=1
export pash_dry_run_compiler_flag=0
export pash_assert_compiler_success_flag=0
export pash_assert_all_regions_parallelizable_flag=0
export pash_avoid_pash_runtime_completion_flag=0
export pash_profile_driven_flag=1
export pash_no_parallel_pipelines=0
export pash_daemon_communicates_through_unix_pipes_flag=0
export pash_speculative_flag=0
export distributed_exec=0
export PASH_DEBUG_LEVEL=1
export PASH_REDIR="&2"

###############################################################################
# Argument Parsing
###############################################################################

## Initialize all argument variables with default values
pash_init_arg_defaults() {
    # Script execution args
    input_script=""
    shell_name="pash"
    declare -g -a script_args=()
    command_mode=""
    command_text=""

    # Shell flags
    allexport_flag="+a"
    verbose_flag=""
    xtrace_flag=""

    # Shared arguments (preprocessor + server)
    arg_debug=""
    arg_log_file=""
    arg_speculative=""
    arg_config_path=""
    arg_local_annotations_dir=""

    # Preprocessor-specific arguments
    arg_bash=""

    # Server-specific arguments
    arg_width=""
    arg_no_optimize=""
    arg_dry_run_compiler=""
    arg_assert_compiler_success=""
    arg_assert_all_regions_parallelizable=""
    arg_avoid_pash_runtime_completion=""
    arg_output_optimized=""
    arg_graphviz=""
    arg_graphviz_dir=""
    arg_no_parallel_pipelines=""
    arg_parallel_pipelines_limit=""
    arg_r_split_batch_size=""
    arg_no_eager=""
    arg_profile_driven=""
    arg_termination=""
    arg_daemon_communicates_through_unix_pipes=""
    arg_distributed_exec=""

    # Obsolete args (still supported for backward compatibility)
    arg_no_daemon=""
    arg_parallel_pipelines=""
    arg_r_split=""
    arg_dgsh_tee=""
    arg_speculation=""

    # Help flag
    show_help=0

    # Unknown flag detection
    unknown_flag=""
}

## Display combined help for PaSh (preprocessor + compiler)
pash_show_help() {
    cat << 'EOF'
Usage: pa.sh [OPTIONS] [-c COMMAND | SCRIPT_FILE] [ARGS...]

PaSh: Parallelizing Shell Scripts

Shell Options:
  -c, --command COMMAND    Execute COMMAND instead of reading from a script file
  -a                       Export all variables (equivalent to bash -a)
  -v                       Verbose mode (equivalent to bash -v)
  -x                       Trace mode (equivalent to bash -x)

Shared Options (Preprocessor + Compiler):
  -d, --debug LEVEL        Debug level (default: 1)
  --log_file FILE          Log file path (default: stderr)
  --speculative            (experimental) Use speculative execution
  --config_path PATH       Path to configuration file

Preprocessor Options:
  --bash                   (experimental) Interpret input as a bash script

Compiler Options:
  -w, --width N            Parallelism width (default: number of CPUs)
  --no_optimize            Disable optimizations
  --dry_run_compiler       Run compiler without executing
  --assert_compiler_success
                           Exit with error if any parallelizable region fails to compile
  --assert_all_regions_parallelizable
                           Exit with error if any region is not parallelizable
  --avoid_pash_runtime_completion
                           Avoid runtime completion
  -p, --output_optimized   Output optimized script
  --graphviz FORMAT        Generate graphviz output (no, dot, svg, pdf, png)
  --graphviz_dir DIR       Directory for graphviz output (default: /tmp)
  --no_parallel_pipelines  Disable parallel pipelines
  --parallel_pipelines_limit N
                           Limit parallel pipelines (default: 2)
  --r_split_batch_size N   Batch size for r_split (default: 1000000)
  --no_eager               Disable eager evaluation
  --profile_driven         (experimental) Enable profile-driven optimization
  --termination MODE       Termination mode: clean_up_graph, drain_stream
  --daemon_communicates_through_unix_pipes
                           (experimental) Use Unix pipes for daemon communication
  --distributed_exec       (experimental) Enable distributed execution

Other Options:
  -h, --help               Show this help message and exit

Examples:
  pa.sh script.sh                    Run script.sh with PaSh
  pa.sh -c "cat file | grep foo"     Run a command with PaSh
  pa.sh -w 4 script.sh               Run with parallelism width 4
  pa.sh --debug 2 script.sh          Run with debug level 2
EOF
}

## Parse all command-line arguments into variables and set exported flags
pash_parse_args() {
    local i=1
    local arg next_i next_arg
    local positional_only=0

    while [ $i -le $# ]; do
        arg="${!i}"
        next_i=$((i+1))
        next_arg="${!next_i}"

        # Once in positional-only mode, all args go to script_args
        if [ "$positional_only" -eq 1 ]; then
            script_args+=("$arg")
            i=$((i+1))
            continue
        fi

        case "$arg" in
            # Help flag (must be handled before other flags)
            -h|--help)
                show_help=1
                ;;

            # Shell flags
            -c|--command)
                command_mode="-c"
                command_text="$next_arg"
                i=$next_i
                positional_only=1
                ;;
            -a)
                allexport_flag="-a"
                ;;
            +a)
                allexport_flag="+a"
                ;;
            -v)
                verbose_flag="-v"
                ;;
            -x)
                xtrace_flag="-x"
                ;;

            # Shared arguments (preprocessor + server)
            -d|--debug)
                arg_debug="$next_arg"
                export PASH_DEBUG_LEVEL="$next_arg"
                i=$next_i
                ;;
            --log_file)
                arg_log_file="$next_arg"
                export PASH_REDIR="$next_arg"
                i=$next_i
                ;;
            --speculative)
                arg_speculative="--speculative"
                export pash_speculative_flag=1
                ;;
            --config_path)
                arg_config_path="$next_arg"
                i=$next_i
                ;;
            --local-annotations-dir)
                arg_local_annotations_dir="$next_arg"
                export ANNOTATIONS_PATH=$(realpath "$next_arg")
                export PYTHONPATH="$ANNOTATIONS_PATH:$PYTHONPATH"
                i=$next_i
                ;;

            # Preprocessor-specific
            --bash)
                arg_bash="--bash"
                export pash_shell="bash"
                ;;

            # Server-specific arguments
            -w|--width)
                arg_width="$next_arg"
                i=$next_i
                ;;
            --no_optimize)
                arg_no_optimize="--no_optimize"
                ;;
            --dry_run_compiler)
                arg_dry_run_compiler="--dry_run_compiler"
                export pash_dry_run_compiler_flag=1
                ;;
            --assert_compiler_success)
                arg_assert_compiler_success="--assert_compiler_success"
                export pash_assert_compiler_success_flag=1
                ;;
            --assert_all_regions_parallelizable)
                arg_assert_all_regions_parallelizable="--assert_all_regions_parallelizable"
                export pash_assert_all_regions_parallelizable_flag=1
                ;;
            --avoid_pash_runtime_completion)
                arg_avoid_pash_runtime_completion="--avoid_pash_runtime_completion"
                export pash_avoid_pash_runtime_completion_flag=1
                ;;
            -p|--output_optimized)
                arg_output_optimized="--output_optimized"
                ;;
            --graphviz)
                arg_graphviz="$next_arg"
                i=$next_i
                ;;
            --graphviz_dir)
                arg_graphviz_dir="$next_arg"
                i=$next_i
                ;;
            --no_parallel_pipelines)
                arg_no_parallel_pipelines="--no_parallel_pipelines"
                export pash_no_parallel_pipelines=1
                ;;
            --parallel_pipelines_limit)
                arg_parallel_pipelines_limit="$next_arg"
                i=$next_i
                ;;
            --r_split_batch_size)
                arg_r_split_batch_size="$next_arg"
                i=$next_i
                ;;
            --no_eager)
                arg_no_eager="--no_eager"
                ;;
            --profile_driven)
                arg_profile_driven="--profile_driven"
                export pash_profile_driven_flag=1
                ;;
            --termination)
                arg_termination="$next_arg"
                i=$next_i
                ;;
            --daemon_communicates_through_unix_pipes)
                arg_daemon_communicates_through_unix_pipes="--daemon_communicates_through_unix_pipes"
                export pash_daemon_communicates_through_unix_pipes_flag=1
                ;;
            --distributed_exec)
                arg_distributed_exec="--distributed_exec"
                export distributed_exec=1
                ;;

            # Obsolete arguments (still accept for backward compatibility)
            --no_daemon)
                arg_no_daemon="--no_daemon"
                ;;
            --parallel_pipelines)
                arg_parallel_pipelines="--parallel_pipelines"
                ;;
            --r_split)
                arg_r_split="--r_split"
                ;;
            --dgsh_tee)
                arg_dgsh_tee="--dgsh_tee"
                ;;
            --speculation)
                arg_speculation="$next_arg"
                i=$next_i
                ;;
            -t|--output_time)
                # Obsolete, time is always logged now
                ;;

            # Positional arguments (input script and script args)
            *)
                # Check if this looks like an unknown flag
                if [[ "$arg" == -* ]] && [ -z "$input_script" ] && [ -z "$command_mode" ]; then
                    unknown_flag="$arg"
                elif [ -z "$input_script" ] && [ -z "$command_mode" ]; then
                    input_script="$arg"
                    shell_name="$arg"
                    positional_only=1
                else
                    script_args+=("$arg")
                fi
                ;;
        esac
        i=$((i+1))
    done

    # In -c mode, first positional arg becomes shell_name ($0)
    if [ -n "$command_mode" ] && [ ${#script_args[@]} -gt 0 ]; then
        shell_name="${script_args[0]}"
        script_args=("${script_args[@]:1}")
    fi
}

## Handle -c command mode: write command text to a temp file
pash_handle_command_mode() {
    if [ -n "$command_text" ]; then
        local command_file
        command_file=$(mktemp "${PASH_TMP_PREFIX}/command_XXXXXX.sh")
        printf '%s' "$command_text" > "$command_file"
        input_script="$command_file"
    fi
}

## Build the preprocessor arguments array
pash_build_preprocessor_args() {
    declare -g -a preprocessor_args=()
    preprocessor_args+=("--output" "$preprocessed_output")
    [ -n "$arg_debug" ] && preprocessor_args+=("-d" "$arg_debug")
    [ -n "$arg_log_file" ] && preprocessor_args+=("--log_file" "$arg_log_file")
    [ -n "$arg_bash" ] && preprocessor_args+=("$arg_bash")
    [ -n "$arg_speculative" ] && preprocessor_args+=("$arg_speculative")
    preprocessor_args+=("$input_script")
}

## Build the compilation server arguments array
pash_build_server_args() {
    declare -g -a server_args=()

    # Shared arguments
    [ -n "$arg_debug" ] && server_args+=("-d" "$arg_debug")
    [ -n "$arg_log_file" ] && server_args+=("--log_file" "$arg_log_file")
    [ -n "$arg_speculative" ] && server_args+=("$arg_speculative")
    [ -n "$arg_config_path" ] && server_args+=("--config_path" "$arg_config_path")

    # Server-specific arguments
    [ -n "$arg_width" ] && server_args+=("-w" "$arg_width")
    [ -n "$arg_no_optimize" ] && server_args+=("$arg_no_optimize")
    [ -n "$arg_dry_run_compiler" ] && server_args+=("$arg_dry_run_compiler")
    [ -n "$arg_assert_compiler_success" ] && server_args+=("$arg_assert_compiler_success")
    [ -n "$arg_assert_all_regions_parallelizable" ] && server_args+=("$arg_assert_all_regions_parallelizable")
    [ -n "$arg_avoid_pash_runtime_completion" ] && server_args+=("$arg_avoid_pash_runtime_completion")
    [ -n "$arg_output_optimized" ] && server_args+=("$arg_output_optimized")
    [ -n "$arg_graphviz" ] && server_args+=("--graphviz" "$arg_graphviz")
    [ -n "$arg_graphviz_dir" ] && server_args+=("--graphviz_dir" "$arg_graphviz_dir")
    [ -n "$arg_no_parallel_pipelines" ] && server_args+=("$arg_no_parallel_pipelines")
    [ -n "$arg_parallel_pipelines_limit" ] && server_args+=("--parallel_pipelines_limit" "$arg_parallel_pipelines_limit")
    [ -n "$arg_r_split_batch_size" ] && server_args+=("--r_split_batch_size" "$arg_r_split_batch_size")
    [ -n "$arg_no_eager" ] && server_args+=("$arg_no_eager")
    [ -n "$arg_profile_driven" ] && server_args+=("$arg_profile_driven")
    [ -n "$arg_termination" ] && server_args+=("--termination" "$arg_termination")
    [ -n "$arg_daemon_communicates_through_unix_pipes" ] && server_args+=("$arg_daemon_communicates_through_unix_pipes")
    [ -n "$arg_distributed_exec" ] && server_args+=("$arg_distributed_exec")

    # Obsolete arguments
    [ -n "$arg_no_daemon" ] && server_args+=("$arg_no_daemon")
    [ -n "$arg_parallel_pipelines" ] && server_args+=("$arg_parallel_pipelines")
    [ -n "$arg_r_split" ] && server_args+=("$arg_r_split")
    [ -n "$arg_dgsh_tee" ] && server_args+=("$arg_dgsh_tee")
    [ -n "$arg_speculation" ] && server_args+=("--speculation" "$arg_speculation")
}

###############################################################################
# Logging Functions
###############################################################################

pash_setup_logging() {
    if [ "$PASH_DEBUG_LEVEL" -le 1 ]; then
        pash_redir_output() {
            :
        }

        pash_redir_all_output() {
            :
        }

        pash_redir_all_output_always_execute() {
            > /dev/null 2>&1 "$@"
        }
    else
        if [ "$PASH_REDIR" == '&2' ]; then
            pash_redir_output() {
                >&2 "$@"
            }

            pash_redir_all_output() {
                >&2 "$@"
            }

            pash_redir_all_output_always_execute() {
                >&2 "$@"
            }
        else
            pash_redir_output() {
                >>"$PASH_REDIR" "$@"
            }

            pash_redir_all_output() {
                >>"$PASH_REDIR" 2>&1 "$@"
            }

            pash_redir_all_output_always_execute() {
                >>"$PASH_REDIR" 2>&1 "$@"
            }
        fi
    fi

    export -f pash_redir_output
    export -f pash_redir_all_output
    export -f pash_redir_all_output_always_execute

    pash_declare_vars() {
        local vars_file="${1?File not given}"
        declare -p > "$vars_file"
    }
    export -f pash_declare_vars
}

###############################################################################
# Communication Functions
###############################################################################

pash_setup_communication() {
    ## Socket communication utilities
    pash_wait_until_unix_socket_listening() {
        local server_name=$1
        local socket=$2
        local i=0
        local maximum_retries=1000
        until echo "Daemon Start" 2> /dev/null | nc -U "$socket" >/dev/null 2>&1; do
            sleep 0.01
            i=$((i+1))
            if [ $i -eq $maximum_retries ]; then
                echo "Error: Maximum retries: $maximum_retries exceeded when waiting for server: ${server_name} to bind to socket: ${socket}!" 1>&2
                echo "Exiting..." 1>&2
                exit 1
            fi
        done
    }

    pash_communicate_unix_socket() {
        local server_name=$1
        local socket=$2
        local message=$3
        pash_redir_output echo "Sending msg to ${server_name}: $message"
        daemon_response=$(echo "$message" | nc -U "${socket}")
        pash_redir_output echo "Got response from ${server_name}: $daemon_response"
        echo "$daemon_response"
    }

    export -f pash_wait_until_unix_socket_listening
    export -f pash_communicate_unix_socket

    ## Daemon communication (conditional on unix pipes flag)
    if [ "$pash_daemon_communicates_through_unix_pipes_flag" -eq 1 ]; then
        pash_communicate_daemon() {
            local message=$1
            pash_redir_output echo "Sending msg to daemon: $message"
            echo "$message" > "$RUNTIME_IN_FIFO"
            daemon_response=$(cat "$RUNTIME_OUT_FIFO")
            pash_redir_output echo "Got response from daemon: $daemon_response"
            echo "$daemon_response"
        }

        pash_communicate_daemon_just_send() {
            local message=$1
            pash_redir_output echo "Sending msg to daemon: $message"
            echo "$message" > "$RUNTIME_IN_FIFO"
        }

        pash_wait_until_daemon_listening() {
            :
        }
    else
        pash_communicate_daemon() {
            local message=$1
            pash_communicate_unix_socket "compilation-server" "${DAEMON_SOCKET}" "${message}"
        }

        pash_communicate_daemon_just_send() {
            pash_communicate_daemon "$1"
        }

        pash_wait_until_daemon_listening() {
            pash_wait_until_unix_socket_listening "compilation-server" "${DAEMON_SOCKET}"
        }
    fi

    export -f pash_communicate_daemon
    export -f pash_communicate_daemon_just_send
    export -f pash_wait_until_daemon_listening

    ## Worker manager communication (for distributed execution)
    if [ "$distributed_exec" -eq 1 ]; then
        pash_communicate_worker_manager() {
            local message=$1
            pash_redir_output echo "Sending msg to worker manager: $message"
            manager_response=$(echo "$message" | nc -U "$DSPASH_SOCKET")
            pash_redir_output echo "Got response from worker manager: $manager_response"
            echo "$manager_response"
        }
        export -f pash_communicate_worker_manager
    fi
}

###############################################################################
# Server Functions
###############################################################################

## Determine Python executable - use venv if available, otherwise system Python
get_pash_python() {
    # First check if there's a local venv (for development/source installs)
    # PASH_TOP points to src/pash, so python_pkgs is at ../../python_pkgs
    local venv_python="${PASH_TOP}/../../python_pkgs/bin/python"
    if [ -x "$venv_python" ]; then
        echo "$venv_python"
    else
        # Fallback to system Python (for pip installs)
        echo "python3"
    fi
}
export PASH_PYTHON=$(get_pash_python)

pash_setup_server_functions() {
    if [ "${pash_speculative_flag}" -eq 1 ]; then
        ## Speculative mode - source the speculative setup
        source "$PASH_TOP/jit_runtime/speculative/pash_spec_init_setup.sh"
    else
        ## Normal PaSh mode
        start_server() {
            "$PASH_PYTHON" "$PASH_TOP/compiler/compilation_server.py" "$@" &
            export daemon_pid=$!
            pash_wait_until_daemon_listening
        }

        cleanup_server() {
            local daemon_pid=$1
            if ps -p "$daemon_pid" > /dev/null 2>&1; then
                msg="Done"
                daemon_response=$(pash_communicate_daemon "$msg")
                if [ "$distributed_exec" -eq 1 ]; then
                    manager_response=$(pash_communicate_worker_manager "$msg")
                fi
                ## Check if assertion failed
                if [[ "$daemon_response" == *"ASSERT_FAILED"* ]]; then
                    pash_assert_failed=1
                fi
                wait 2> /dev/null 1>&2
            fi
        }

        export -f start_server
        export -f cleanup_server
    fi
}

###############################################################################
# Main Execution
###############################################################################

## 1. Parse arguments (also sets exported flags)
pash_init_arg_defaults
pash_parse_args "$@"

## Show help and exit if requested
if [ "$show_help" -eq 1 ]; then
    pash_show_help
    exit 0
fi

## Check for unknown flags
if [ -n "$unknown_flag" ]; then
    echo "pa.sh: unknown option: $unknown_flag" >&2
    echo "" >&2
    pash_show_help >&2
    exit 1
fi

## 2. Setup functions based on parsed arguments
pash_setup_logging
pash_setup_communication
pash_setup_server_functions

## 3. Handle -c command mode
pash_handle_command_mode

## 4. Create temporary file for preprocessed output
preprocessed_output=$(mktemp "${PASH_TMP_PREFIX}/preprocessed_XXXXXX.sh")

## 5. Build argument arrays for preprocessor and server
pash_build_preprocessor_args
pash_build_server_args

## 6. Start the compilation server
start_server "${server_args[@]}"

## 7. Restore umask before executing user scripts
umask "$old_umask"

## 8. Run the preprocessor
PYTHONPATH="$PASH_TOP/preprocessor:$PASH_TOP/compiler:$PYTHONPATH" \
    PASH_FROM_SH="PaSh preprocessor" "$PASH_PYTHON" \
    "$PASH_TOP/preprocessor/pash_preprocessor.py" "${preprocessor_args[@]}"
pash_exit_code=$?

## 9. If preprocessing succeeded, execute the preprocessed script
if [ "$pash_exit_code" -eq 0 ]; then
    bash_flags="$allexport_flag $verbose_flag $xtrace_flag"
    # shellcheck disable=SC2086
    bash $bash_flags -c "source $preprocessed_output" "$shell_name" "${script_args[@]}"
    pash_exit_code=$?
fi

## 10. Cleanup
cleanup_server "${daemon_pid}"

## Check if assertion failed (--assert_all_regions_parallelizable or --assert_compiler_success)
if [ "${pash_assert_failed:-0}" -eq 1 ]; then
    pash_exit_code=1
fi

if [ "$PASH_DEBUG_LEVEL" -le 1 ]; then
    rm -rf "${PASH_TMP_PREFIX}"
fi

(exit "$pash_exit_code")
