#!/bin/bash -eu
# SPDX-FileCopyrightText: 2025 Uwe Fechner
# SPDX-License-Identifier: MIT
#
# Install and configure matplotlib for ControlPlots.jl (via PyCall/PyPlot).
#
# Two backends are supported:
#   1) System Python  – uses the matplotlib package installed by Ubuntu/Debian (apt).
#                       Fastest option; shares the system Python install.
#   2) Conda          – installs matplotlib into Julia's own Conda environment.
#                       Self-contained; does not require root / sudo.
#
# Both options configure the Qt (qtagg) backend for interactive plot windows.

print_usage() {
    echo "Usage:"
    echo "  ./bin/install_controlplots [--system | --conda] [-y | --yes] [-h | --help]"
    echo ""
    echo "Options:"
    echo "  --system   Use the system Python and matplotlib (Ubuntu/Debian apt)"
    echo "  --conda    Use Julia's Conda environment to install matplotlib"
    echo "  -y, --yes  Non-interactive; accept defaults"
    echo "  -h, --help Show this help message"
}

_backend=""   # "system" or "conda"
_yes=false

while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            print_usage
            exit 0
            ;;
        --system)
            _backend="system"
            shift
            ;;
        --conda)
            _backend="conda"
            shift
            ;;
        -y|--yes)
            _yes=true
            shift
            ;;
        *)
            echo "Unknown option: $1"
            print_usage
            exit 1
            ;;
    esac
done

# Always run from the repository root.
if [[ "$(basename "$(pwd)")" == "bin" ]]; then
    cd ..
fi

# Source the setup_env helper so the julia() wrapper (with JLL LD_PRELOAD overrides) is available.
export LD_PRELOAD="${LD_PRELOAD:-}"
export LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}"
. ./bin/setup_env

# ── Interactive backend selection ─────────────────────────────────────────────

if [[ -z "$_backend" ]]; then
    echo "Which matplotlib backend do you want to use for ControlPlots?"
    echo ""
    echo "  1) System Python  (uses Ubuntu/Debian python3-matplotlib via apt)"
    echo "     - Requires sudo"
    echo "     - Shares the system-wide Python installation"
    echo ""
    echo "  2) Conda          (installs matplotlib into Julia's own Conda environment)"
    echo "     - No sudo required"
    echo "     - Self-contained; downloads ~200 MB on first use"
    echo ""
    if [[ "$_yes" == true ]]; then
        _choice="1"
        echo "Using default: 1 (system Python)"
    else
        read -rp "Enter 1 or 2 [default: 1]: " _choice
    fi
    case "${_choice:-1}" in
        1) _backend="system" ;;
        2) _backend="conda"  ;;
        *)
            echo "Invalid choice. Please enter 1 or 2."
            exit 1
            ;;
    esac
fi

echo ""
echo "Selected backend: $_backend"
echo ""

# ── Helper: rebuild PyCall and verify matplotlib ──────────────────────────────

_rebuild_pycall_and_verify() {
    # $1 = PYTHON value to pass (empty string → Conda, path → system)
    local _pyval="$1"
    echo "Rebuilding PyCall with PYTHON='${_pyval}'..."
    env -u LD_PRELOAD -u LD_LIBRARY_PATH \
        PYTHON="$_pyval" \
        julia --project=examples -e '
using Pkg
ENV["PYTHON"] = get(ENV, "PYTHON", "")
Pkg.build("PyCall")
println("PyCall rebuilt successfully.")
'
}

_install_matplotlib_conda() {
    echo "Installing matplotlib and pyqt into Julia's Conda environment..."
    env -u LD_PRELOAD -u LD_LIBRARY_PATH \
        PYTHON="" \
        julia --project=examples -e '
using Pkg
# Ensure Conda is available
if Base.find_package("Conda") === nothing
    Pkg.add("Conda")
end
using Conda
Conda.add("matplotlib")
Conda.add("pyqt")
println("matplotlib and pyqt installed in Conda environment.")
'
}

_verify_controlplots() {
    echo "Verifying ControlPlots can be loaded..."
    if env -u LD_PRELOAD -u LD_LIBRARY_PATH \
        julia --project=examples -e '
using ControlPlots
println("ControlPlots loaded successfully.")
'; then
        echo ""
        echo "✓ ControlPlots is working."
    else
        echo ""
        echo "Warning: ControlPlots could not be loaded. Check the error output above."
        echo "You may need to run this script again with the other backend option."
        exit 1
    fi
}

# ── Helper: write MPLBACKEND=qtagg to LocalPreferences.toml ──────────────────

_set_mplbackend_qtagg() {
    local _prefs_file="examples/LocalPreferences.toml"
    # Remove any existing [ENV] section managed by this script.
    if [[ -f "$_prefs_file" ]]; then
        _tmp=$(mktemp)
        awk '
            /^\[ENV\]/ { in_env=1; next }
            in_env && /^\[/ { in_env=0 }
            !in_env { print }
        ' "$_prefs_file" > "$_tmp"
        mv "$_tmp" "$_prefs_file"
    fi
    {
        echo ""
        echo "[ENV]"
        echo "MPLBACKEND = \"qtagg\""
    } >> "$_prefs_file"
    echo "Set MPLBACKEND=qtagg in $_prefs_file."
}

# ── System Python backend ─────────────────────────────────────────────────────

if [[ "$_backend" == "system" ]]; then
    # Detect package manager and install python3-matplotlib if missing.
    if [[ "$(uname -s)" == "Linux" ]]; then
        if grep -qiE "ubuntu|debian" /etc/os-release 2>/dev/null; then
            if ! dpkg -s python3-matplotlib &>/dev/null 2>&1 || ! dpkg -s python3-pyqt5 &>/dev/null 2>&1; then
                echo "Installing python3-matplotlib and python3-pyqt5 via apt..."
                sudo apt install -y python3-matplotlib python3-pyqt5
            else
                echo "python3-matplotlib and python3-pyqt5 are already installed."
            fi
        elif grep -qi "fedora" /etc/os-release 2>/dev/null; then
            if ! rpm -q python3-matplotlib &>/dev/null 2>&1 || ! rpm -q python3-pyqt5 &>/dev/null 2>&1; then
                echo "Installing python3-matplotlib and python3-qt5 via dnf..."
                sudo dnf install -y python3-matplotlib python3-pyqt5
            else
                echo "python3-matplotlib and python3-pyqt5 are already installed."
            fi
        else
            echo "Warning: Could not detect Ubuntu/Debian or Fedora."
            echo "Please ensure python3-matplotlib is installed manually before continuing."
            if [[ "$_yes" == false ]]; then
                read -rp "Continue anyway? (y/n) [default: y]: " _cont
                case "${_cont:-y}" in
                    n|N) echo "Aborted."; exit 1 ;;
                esac
            fi
        fi
    else
        echo "Warning: System Python backend is intended for Ubuntu/Debian/Fedora Linux."
        echo "On this OS ($(uname -s)) you may need to install matplotlib manually."
        if [[ "$_yes" == false ]]; then
            read -rp "Continue anyway? (y/n) [default: y]: " _cont
            case "${_cont:-y}" in
                n|N) echo "Aborted."; exit 1 ;;
            esac
        fi
    fi

    # Locate the system python3 executable.
    _syspython=""
    for _candidate in python3 /usr/bin/python3; do
        if command -v "$_candidate" &>/dev/null; then
            _syspython=$(command -v "$_candidate")
            break
        fi
    done
    if [[ -z "$_syspython" ]]; then
        echo "Error: python3 not found on PATH. Please install python3."
        exit 1
    fi
    echo "Found system Python: $_syspython"

    # Point PyCall at the system Python and rebuild.
    env -u LD_PRELOAD -u LD_LIBRARY_PATH \
        PYTHON="$_syspython" \
        julia --project=examples -e '
using Pkg
ENV["PYTHON"] = ENV["PYTHON"]   # already set by the shell
Pkg.build("PyCall")
println("PyCall rebuilt to use system Python: " * ENV["PYTHON"])
'

    echo ""
    echo "Saving PYTHON=$_syspython to examples/LocalPreferences.toml..."
    # Write (or update) LocalPreferences.toml in the examples directory so that
    # Julia's PyCall always uses the system Python when running examples.
    _prefs_file="examples/LocalPreferences.toml"
    if [[ -f "$_prefs_file" ]]; then
        # Remove any existing [PyCall] section (naive sed approach).
        _tmp=$(mktemp)
        awk '
            /^\[PyCall\]/ { in_pycall=1; next }
            in_pycall && /^\[/ { in_pycall=0 }
            !in_pycall { print }
        ' "$_prefs_file" > "$_tmp"
        mv "$_tmp" "$_prefs_file"
    fi
    {
        echo ""
        echo "[PyCall]"
        echo "python = \"$_syspython\""
    } >> "$_prefs_file"
    echo "Written to $_prefs_file."

    _set_mplbackend_qtagg
    _verify_controlplots

    echo ""
    echo "Done. ControlPlots will use the system Python matplotlib with the Qt (qtagg) backend."

# ── Conda backend ─────────────────────────────────────────────────────────────

elif [[ "$_backend" == "conda" ]]; then
    # Remove any LocalPreferences.toml [PyCall] section so PyCall uses Conda.
    _prefs_file="examples/LocalPreferences.toml"
    if [[ -f "$_prefs_file" ]]; then
        _tmp=$(mktemp)
        awk '
            /^\[PyCall\]/ { in_pycall=1; next }
            in_pycall && /^\[/ { in_pycall=0 }
            !in_pycall { print }
        ' "$_prefs_file" > "$_tmp"
        # If PyCall section existed, update; otherwise leave unchanged.
        if ! diff -q "$_prefs_file" "$_tmp" &>/dev/null; then
            mv "$_tmp" "$_prefs_file"
            echo "Removed [PyCall] section from $_prefs_file (switching to Conda)."
        else
            rm -f "$_tmp"
        fi
    fi

    # Rebuild PyCall with PYTHON="" to trigger Conda mode.
    env -u LD_PRELOAD -u LD_LIBRARY_PATH \
        PYTHON="" \
        julia --project=examples -e '
using Pkg
ENV["PYTHON"] = ""
Pkg.build("PyCall")
println("PyCall rebuilt to use Conda-managed Python.")
'

    _install_matplotlib_conda

    _set_mplbackend_qtagg
    _verify_controlplots

    echo ""
    echo "Done. ControlPlots will use the Conda-managed matplotlib with the Qt (qtagg) backend."
fi
