#!/usr/bin/env bash
set -Eeuo pipefail

readonly THEME_URL='https://raw.githubusercontent.com/JanDeDobbeleer/oh-my-posh/main/themes/atomic.omp.json'
readonly FASTFETCH_RELEASE_URL='https://github.com/fastfetch-cli/fastfetch/releases/latest/download'
readonly FISH_PPA='ppa:fish-shell/release-4'
readonly BLOCK_START='# >>> itzme1on shell setup >>>'
readonly BLOCK_END='# <<< itzme1on shell setup <<<'

log() {
    printf '\033[1;34m[setup]\033[0m %s\n' "$*"
}

warn() {
    printf '\033[1;33m[warning]\033[0m %s\n' "$*" >&2
}

fail() {
    printf '\033[1;31m[error]\033[0m %s\n' "$*" >&2
    exit 1
}

run_as_root() {
    if (( EUID == 0 )); then
        "$@"
    else
        sudo "$@"
    fi
}

if [[ ! -r /etc/os-release ]]; then
    fail 'Cannot identify the operating system.'
fi

# shellcheck source=/dev/null
source /etc/os-release

if [[ "${ID:-}" != 'ubuntu' ]]; then
    fail 'This installer supports Ubuntu only (22.04 LTS and newer).'
fi

case "${VERSION_ID:-}" in
    22.04|24.04) ;;
    *) fail 'This installer supports Ubuntu 22.04 LTS and 24.04 LTS.' ;;
esac

if (( EUID != 0 )) && ! command -v sudo >/dev/null 2>&1; then
    fail 'sudo is required for non-root users.'
fi

user_name="$(id -un)"
user_home="$(getent passwd "$user_name" | cut -d: -f6)"

[[ -n "$user_home" && -d "$user_home" ]] || fail "Home directory not found for $user_name."

if [[ "${HOME:-}" != "$user_home" ]]; then
    warn "Using $user_home as the home directory for $user_name."
    export HOME="$user_home"
fi

readonly TEMP_DIR="$(mktemp -d)"
trap 'rm -rf "$TEMP_DIR"' EXIT

APT_UPDATED=0

apt_update() {
    if (( APT_UPDATED == 0 )); then
        log 'Refreshing APT package lists...'
        run_as_root apt-get update
        APT_UPDATED=1
    fi
}

apt_install() {
    apt_update
    run_as_root env DEBIAN_FRONTEND=noninteractive apt-get install -y "$@"
}

ensure_dependencies() {
    local package
    local missing=()

    for package in ca-certificates curl python3 unzip; do
        if ! dpkg-query -W -f='${Status}' "$package" 2>/dev/null | grep -qx 'install ok installed'; then
            missing+=("$package")
        fi
    done

    if (( ${#missing[@]} > 0 )); then
        log "Installing prerequisites: ${missing[*]}"
        apt_install "${missing[@]}"
    fi
}

install_fish() {
    local version=''

    if command -v fish >/dev/null 2>&1; then
        version="$(fish --version | awk '{print $3}')"

        if [[ "$version" =~ ^[0-9]+\.[0-9]+(\.[0-9]+)?$ ]] && dpkg --compare-versions "$version" ge '4.1.0'; then
            log "Fish $version is already installed."
            return
        fi
    fi

    log 'Installing Fish 4.x from the official fish-shell PPA...'
    apt_install software-properties-common
    run_as_root add-apt-repository --yes "$FISH_PPA"
    APT_UPDATED=0
    apt_install fish

    if ! command -v fish >/dev/null 2>&1; then
        fail 'Fish could not be found after installation.'
    fi

    version="$(fish --version | awk '{print $3}')"
    if ! dpkg --compare-versions "$version" ge '4.1.0'; then
        fail "Fish 4.1 or newer is required for Oh My Posh (found: $version)."
    fi
}

install_fastfetch() {
    local arch asset

    if command -v fastfetch >/dev/null 2>&1; then
        log 'Fastfetch is already installed.'
        return
    fi

    arch="$(dpkg --print-architecture)"
    case "$arch" in
        amd64) asset='fastfetch-linux-amd64.deb' ;;
        arm64) asset='fastfetch-linux-aarch64.deb' ;;
        *) fail "Unsupported architecture for Fastfetch: $arch" ;;
    esac

    log 'Downloading Fastfetch from the official GitHub release...'
    curl -fLsS --retry 3 "$FASTFETCH_RELEASE_URL/$asset" -o "$TEMP_DIR/fastfetch.deb"
    apt_install "$TEMP_DIR/fastfetch.deb"

    command -v fastfetch >/dev/null 2>&1 || fail 'Fastfetch installation failed.'
}

install_oh_my_posh() {
    if command -v oh-my-posh >/dev/null 2>&1 || [[ -x "$HOME/.local/bin/oh-my-posh" ]]; then
        log 'Oh My Posh is already installed.'
        return
    fi

    log 'Installing Oh My Posh to ~/.local/bin...'
    mkdir -p "$HOME/.local/bin"
    curl -fLsS --retry 3 'https://ohmyposh.dev/install.sh' -o "$TEMP_DIR/posh-install.sh"
    bash "$TEMP_DIR/posh-install.sh" -d "$HOME/.local/bin"

    [[ -x "$HOME/.local/bin/oh-my-posh" ]] || fail 'Oh My Posh installation failed.'
}

install_atomic_theme() {
    local config_dir="$HOME/.config/oh-my-posh"
    local theme_path="$config_dir/atomic.omp.json"

    log 'Downloading the atomic theme for offline use...'
    mkdir -p "$config_dir"
    curl -fLsS --retry 3 "$THEME_URL" -o "$TEMP_DIR/atomic.omp.json"
    python3 -m json.tool "$TEMP_DIR/atomic.omp.json" >/dev/null || fail 'The downloaded theme is not valid JSON.'

    if [[ -f "$theme_path" ]] && ! cmp -s "$theme_path" "$TEMP_DIR/atomic.omp.json"; then
        cp -p "$theme_path" "$theme_path.bak.$(date +%Y%m%d-%H%M%S)"
    fi

    install -m 644 "$TEMP_DIR/atomic.omp.json" "$theme_path"
}

configure_fish() {
    local config_dir="$HOME/.config/fish"
    local config_path="$config_dir/config.fish"
    local clean_path="$TEMP_DIR/config.clean.fish"
    local new_path="$TEMP_DIR/config.new.fish"
    local show_fastfetch=1

    log 'Configuring SSH Fastfetch and the atomic prompt...'
    mkdir -p "$config_dir"
    touch "$config_path"

    # Replace only our managed block and simple, pre-existing Oh My Posh init lines.
    # All other user configuration is preserved, with a backup on change.
    awk -v start="$BLOCK_START" -v stop="$BLOCK_END" '
        $0 == start { managed = 1; next }
        $0 == stop { managed = 0; next }
        managed { next }
        /^[[:space:]]*oh-my-posh[[:space:]]+init[[:space:]]+fish([[:space:]]|$)/ {
            lines[++count] = "# Disabled by itzme1on shell setup: " $0
            next
        }
        { lines[++count] = $0 }
        END {
            while (count > 0 && lines[count] ~ /^[[:space:]]*$/) count--
            for (i = 1; i <= count; i++) print lines[i]
        }
    ' "$config_path" > "$clean_path"

    if grep -Eq '^[[:space:]]*fastfetch([[:space:]]|$)' "$clean_path"; then
        show_fastfetch=0
        log 'An existing Fastfetch startup command was found; keeping it.'
    fi

    cat "$clean_path" > "$new_path"
    {
        printf '\n%s\n' "$BLOCK_START"
        printf 'if status is-interactive\n'

        if (( show_fastfetch )); then
            printf '    if set -q SSH_CONNECTION; and type -q fastfetch\n'
            printf '        fastfetch\n'
            printf '    end\n\n'
        fi

        cat <<'FISH'
    if type -q oh-my-posh
        oh-my-posh init fish --config "$HOME/.config/oh-my-posh/atomic.omp.json" | source
    else if test -x "$HOME/.local/bin/oh-my-posh"
        "$HOME/.local/bin/oh-my-posh" init fish --config "$HOME/.config/oh-my-posh/atomic.omp.json" | source
    end
end
FISH
        printf '%s\n' "$BLOCK_END"
    } >> "$new_path"

    fish -n "$new_path" || fail 'The generated Fish configuration did not pass syntax validation.'

    if ! cmp -s "$config_path" "$new_path"; then
        if [[ -s "$config_path" ]]; then
            cp -p "$config_path" "$config_path.bak.$(date +%Y%m%d-%H%M%S)"
        fi
        install -m 600 "$new_path" "$config_path"
    fi

    if grep -Eq '^[[:space:]]*([^#[:space:]].*)?oh-my-posh[[:space:]]+init[[:space:]]+fish' "$clean_path"; then
        warn 'A custom Oh My Posh initialization may remain in config.fish; check for duplicates.'
    fi

    if [[ -d "$config_dir/conf.d" ]] && grep -REq --include='*.fish' '^[[:space:]]*([^#[:space:]].*)?oh-my-posh[[:space:]]+init[[:space:]]+fish' "$config_dir/conf.d"; then
        warn 'Another Oh My Posh initializer exists in fish/conf.d; review it to avoid duplicate prompts.'
    fi
}

configure_login() {
    local fish_path current_shell user_name

    touch "$HOME/.hushlogin"

    fish_path="$(command -v fish)"
    user_name="$(id -un)"
    current_shell="$(getent passwd "$user_name" | cut -d: -f7)"

    if ! grep -Fxq "$fish_path" /etc/shells; then
        printf '%s\n' "$fish_path" | run_as_root tee -a /etc/shells >/dev/null
    fi

    if [[ "$current_shell" != "$fish_path" ]]; then
        if [[ "$user_name" == 'root' ]]; then
            warn 'Changing the root login shell to Fish may affect SSH commands that assume Bash.'
        fi

        log "Changing the login shell for $user_name to $fish_path..."
        run_as_root usermod --shell "$fish_path" "$user_name"
    fi
}

main() {
    log 'Setting up your Ubuntu SSH shell...'
    ensure_dependencies
    install_fish
    install_fastfetch
    install_oh_my_posh
    install_atomic_theme
    configure_fish
    configure_login

    log 'Done! Reconnect over SSH to see Fastfetch and the atomic theme.'
    log 'If symbols look wrong, select a Nerd Font in your SSH terminal.'
}

main "$@"