#!/bin/bash

# Common helpers as requirement
if [[ -f /etc/nestri/common.sh ]]; then
    source /etc/nestri/common.sh
else
    echo "Error: Common script not found at /etc/nestri/common.sh" >&2
    exit 1
fi

# Parses resolution string
parse_resolution() {
    local resolution="$1"
    if [[ -z "$resolution" ]]; then
        log "Error: No resolution provided"
        return 1
    fi

    IFS='x' read -r width height <<< "$resolution"
    if ! [[ "$width" =~ ^[0-9]+$ ]] || ! [[ "$height" =~ ^[0-9]+$ ]]; then
        log "Error: Invalid resolution format. Expected: WIDTHxHEIGHT (e.g., 1920x1080), got: $resolution"
        return 1
    fi

    export WIDTH="$width"
    export HEIGHT="$height"
    return 0
}

# Configuration
MAX_RETRIES=3
RETRY_COUNT=0
WAYLAND_READY_DELAY=3
ENTCMD_PREFIX=""
PRELOAD_SHIM_64=/usr/lib64/libvimputti_shim.so
PRELOAD_SHIM_32=/usr/lib32/libvimputti_shim.so

# Kills process if running
kill_if_running() {
    local pid="$1"
    local name="$2"
    if [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null; then
        log "Killing existing $name process (PID: $pid)..."
        kill "$pid"
        wait "$pid" 2>/dev/null || true
    fi
}

# Starts nestri-server
start_nestri_server() {
    kill_if_running "${NESTRI_PID:-}" "nestri-server"

    WAYLAND_SOCKET="${XDG_RUNTIME_DIR}/wayland-1"

    # Make sure to remove old socket if exists (along with .lock file)
    if [[ -e "$WAYLAND_SOCKET" ]]; then
        log "Removing stale Wayland socket $WAYLAND_SOCKET"
        rm -f "$WAYLAND_SOCKET" 2>/dev/null || {
            log "Error: Failed to remove stale Wayland socket $WAYLAND_SOCKET"
            exit 1
        }
        # Ignore error if .lock file doesn't exist
        rm -f "${WAYLAND_SOCKET}.lock" 2>/dev/null || true
    fi

    # Also if gstreamer cache exists, remove it to avoid previous errors from persisting
    local gst_cache="${NESTRI_HOME}/.cache/gstreamer-1.0/"
    if [[ -d "$gst_cache" ]]; then
        log "Removing gstreamer cache at $gst_cache"
        rm -rf "${gst_cache}" 2>/dev/null || {
            log "Warning: Failed to remove gstreamer cache at $gst_cache"
        }
    fi

    # Start nestri-server
    log "Starting nestri-server.."
    # Try with realtime scheduling first (chrt -f 80), if fails, launch normally
    if $ENTCMD_PREFIX chrt -f 80 true 2>/dev/null; then
        $ENTCMD_PREFIX chrt -f 80 nestri-server $NESTRI_PARAMS &
        NESTRI_PID=$!
        log "Started nestri-server with realtime scheduling"
    else
        $ENTCMD_PREFIX nestri-server $NESTRI_PARAMS &
        NESTRI_PID=$!
        log "Started nestri-server"
    fi

    log "Waiting for Wayland display $WAYLAND_SOCKET.."
    for ((i=1; i<=15; i++)); do
        if [[ -e "$WAYLAND_SOCKET" ]]; then
            log "Wayland display $WAYLAND_SOCKET ready"
            sleep "${WAYLAND_READY_DELAY:-3}"
            start_compositor
            return
        fi
        sleep 1
    done

    log "Error: Wayland display $WAYLAND_SOCKET not available"
    increment_retry "nestri-server"
    restart_chain
}

# Starts compositor with optional application
start_compositor() {
    kill_if_running "${COMPOSITOR_PID:-}" "compositor"
    kill_if_running "${APP_PID:-}" "application"

    # Set default compositor if unset
    if [[ -z "${NESTRI_LAUNCH_COMPOSITOR+x}" ]]; then
        NESTRI_LAUNCH_COMPOSITOR="gamescope --backend wayland -g -f --rt -W ${WIDTH} -H ${HEIGHT} -r ${FRAMERATE:-60}"
    fi

    # If PRELOAD_SHIM_arch's are set and exist, set LD_PRELOAD for 32/64-bit apps
    local do_ld_preload=false
    if [[ -f "$PRELOAD_SHIM_64" ]] || [[ -f "$PRELOAD_SHIM_32" ]]; then
        do_ld_preload=true
        log "Using LD_PRELOAD shim(s)"
    fi

    # Configure launch cmd with dbus if set
    local launch_cmd=""
    if [[ -n "${NESTRI_LAUNCH_CMD+x}" ]]; then
        if $do_ld_preload; then
            launch_cmd="LD_PRELOAD='/usr/\$LIB/libvimputti_shim.so' dbus-launch $NESTRI_LAUNCH_CMD"
        else
            launch_cmd="dbus-launch $NESTRI_LAUNCH_CMD"
        fi
    fi

    # Launch compositor if configured
    if [[ -n "${NESTRI_LAUNCH_COMPOSITOR}" ]]; then
        local compositor_cmd="$NESTRI_LAUNCH_COMPOSITOR"
        local is_gamescope=false

        # Check if this is a gamescope command
        if [[ "$compositor_cmd" == *"gamescope"* ]]; then
            is_gamescope=true
            if [[ -n "$launch_cmd" ]] && [[ "$compositor_cmd" != *" -- "* ]]; then
                # If steam in launch command, enable gamescope integration via -e and enable mangohud
                if [[ "$launch_cmd" == *"steam"* ]]; then
                    compositor_cmd+=" --mangoapp -e"
                fi
                compositor_cmd+=" -- bash -c $(printf %q "$launch_cmd")"
            fi
        fi

        # Get appropriate socket based on compositor type
        if $is_gamescope; then
            COMPOSITOR_SOCKET="${XDG_RUNTIME_DIR}/gamescope-0"
        else
            COMPOSITOR_SOCKET="${XDG_RUNTIME_DIR}/wayland-0"
        fi

        # Clean up old socket if exists
        if [[ -e "$COMPOSITOR_SOCKET" ]]; then
            log "Removing stale compositor socket $COMPOSITOR_SOCKET"
            rm -f "$COMPOSITOR_SOCKET" 2>/dev/null || {
                log "Error: Failed to remove stale compositor socket $COMPOSITOR_SOCKET"
                exit 1
            }
            # Ignore error if .lock file doesn't exist
            rm -f "${COMPOSITOR_SOCKET}.lock" 2>/dev/null || true
        fi

        log "Starting compositor: $compositor_cmd"
        WAYLAND_DISPLAY=wayland-1 /bin/bash -c "$compositor_cmd" &
        COMPOSITOR_PID=$!

        # If ld_preload is true, export LD_PRELOAD
        if $do_ld_preload; then
            export LD_PRELOAD='/usr/$LIB/libvimputti_shim.so'
        fi

        log "Waiting for compositor socket $COMPOSITOR_SOCKET.."
        for ((i=1; i<=15; i++)); do
            if [[ -e "$COMPOSITOR_SOCKET" ]]; then
                log "Compositor socket ready $COMPOSITOR_SOCKET"
                # Patch resolution with wlr-randr for non-gamescope compositors
                if ! $is_gamescope; then
                    local OUTPUT_NAME
                    OUTPUT_NAME=$(WAYLAND_DISPLAY=wayland-0 wlr-randr --json | jq -r '.[] | select(.enabled == true) | .name' | head -n 1)
                    if [ -z "$OUTPUT_NAME" ]; then
                        log "Warning: No enabled outputs detected. Skipping wlr-randr resolution patch"
                        return
                    fi
                    WAYLAND_DISPLAY=wayland-0 wlr-randr --output "$OUTPUT_NAME" --custom-mode "$WIDTH"x"$HEIGHT"
                    log "Patched resolution with wlr-randr"

                    if [[ -n "$launch_cmd" ]]; then
                        log "Starting application: $launch_cmd"
                        WAYLAND_DISPLAY="$COMPOSITOR_SOCKET" bash -c "$launch_cmd" &
                        APP_PID=$!
                    fi
                else
                    log "Gamescope detected, skipping wlr-randr resolution patch"
                fi
                return
            fi
            sleep 1
        done
        log "Warning: Compositor socket not found after 15 seconds ($COMPOSITOR_SOCKET)"
    else
        # Launch standalone application if no compositor
        if [[ -n "$launch_cmd" ]]; then
            log "Starting standalone application: $launch_cmd"
            WAYLAND_DISPLAY=wayland-1 bash -c "$launch_cmd" &
            APP_PID=$!
        else
            log "No compositor or application configured"
        fi
    fi
}

# Increments retry counter
increment_retry() {
    local component="$1"
    ((RETRY_COUNT++))
    if [[ "$RETRY_COUNT" -ge "$MAX_RETRIES" ]]; then
        log "Error: Max retries reached for $component"
        exit 1
    fi
}

# Restarts the chain
restart_chain() {
    log "Restarting nestri-server and compositor..."
    start_nestri_server
}

# Cleans up processes
cleanup() {
    local exit_code=$?
    log "Terminating processes..."
    kill_if_running "${NESTRI_PID:-}" "nestri-server"
    kill_if_running "${COMPOSITOR_PID:-}" "compositor"
    kill_if_running "${APP_PID:-}" "application"
    exit $exit_code
}

# Monitor processes for unexpected exits
main_loop() {
    trap cleanup SIGINT SIGTERM EXIT

    while true; do
        sleep 1
        # Check nestri-server
        if [[ -n "${NESTRI_PID:-}" ]] && ! kill -0 "${NESTRI_PID}" 2>/dev/null; then
            log "nestri-server died"
            increment_retry "nestri-server"
            restart_chain
        # Check compositor
        elif [[ -n "${COMPOSITOR_PID:-}" ]] && ! kill -0 "${COMPOSITOR_PID}" 2>/dev/null; then
            log "compositor died"
            increment_retry "compositor"
            start_compositor
        # Check application
        elif [[ -n "${APP_PID:-}" ]] && ! kill -0 "${APP_PID}" 2>/dev/null; then
            log "application died"
            increment_retry "application"
            start_compositor
        fi
    done
}

main() {
    parse_resolution "${RESOLUTION:-1920x1080}" || exit 1
    get_container_info || {
        log "Warning: Failed to detect container information."
    }

    if [[ "$container_runtime" != "apptainer" ]]; then
        ENTCMD_PREFIX="sudo -E -u ${NESTRI_USER}"
    fi

    restart_chain
    main_loop
}

main