#!/bin/bash
set -e

# Enable pipefail where supported (bash/ksh); silently skip in POSIX sh
(set -o pipefail) 2>/dev/null && set -o pipefail || true

# Default values (will be set based on detected OS)
export DEBIAN_FRONTEND=noninteractive
ACTION=""
DRIVER_VARIANT=""
UNINSTALL_REMOVED_PACKAGES=0
REPO_NAME="cix-repo-community"
DIST_CODENAME=""
REPO_URL=""
GPG_KEY_URL="https://archive.cixtech.com/ppa-gpg-public-key.asc"
DRIVER_PACKAGE_BASE=""
OPEN_SOURCE_DRIVER_PACKAGE="cix-debian13-k7.0-driver-opensource"
DRIVER_PACKAGE=""
KERNEL_VERSION="6.6"
KERNEL_BOOT_GLOB="/boot/vmlinuz-*-generic"

# Color codes for output
RED=$'\033[0;31m'
GREEN=$'\033[0;32m'
YELLOW=$'\033[1;33m'
NC=$'\033[0m' # No Color

error_exit() {
    echo "${RED}❌ $1${NC}" >&2
    exit 1
}

warn() {
    echo "${YELLOW}⚠ $1${NC}" >&2
}

success() {
    echo "${GREEN}✅ $1${NC}"
}

usage() {
    cat <<EOF
Usage: cix-repo-community.sh [OPTIONS]

Options:
  --install             Install closed-source driver, Linux 6.6 (non-interactive)
  --install-full        Install full closed-source driver, Linux 6.6
                        (Debian 12/bookworm only, non-interactive)
  --install-opensource  Install open-source driver, Linux 7.0 (non-interactive)
  --uninstall           Uninstall CIX packages and configuration (non-interactive)
  -h, --help            Show this help message

Without any of the above action options, the script runs interactively.
EOF
}

parse_args() {
    while [ "$#" -gt 0 ]; do
        case "$1" in
            -h|--help|help)
                usage
                exit 0
                ;;
            --install)
                ACTION="install"
                DRIVER_VARIANT="closed"
                ;;
            --install-full)
                ACTION="install"
                DRIVER_VARIANT="full"
                ;;
            --install-opensource)
                ACTION="install"
                DRIVER_VARIANT="opensource"
                ;;
            --uninstall)
                ACTION="uninstall"
                ;;
            *)
                error_exit "Unknown argument: $1"
                ;;
        esac
        shift
    done
}

set_closed_source_install() {
    ACTION="install"
    case "$DRIVER_VARIANT" in
        full) DRIVER_PACKAGE="${DRIVER_PACKAGE_BASE}-full" ;;
        *)    DRIVER_PACKAGE="${DRIVER_PACKAGE_BASE}" ;;
    esac
    KERNEL_VERSION="6.6"
}

full_install_available() {
    [ "$ID" = "debian" ] && [ "$VERSION_CODENAME" = "bookworm" ]
}

open_source_available() {
    [ "$ID" = "debian" ] && [ "$VERSION_CODENAME" = "trixie" ]
}

closed_source_available() {
    [ -n "$DRIVER_PACKAGE_BASE" ]
}

set_open_source_repository() {
    DIST_CODENAME="trixie"
    REPO_URL="https://archive.cixtech.com/debian"
    REPO_NAME="cix-deb-repo"
}

set_open_source_install() {
    if ! open_source_available; then
        error_exit "Open-source driver is only supported on Debian 13 (trixie)"
    fi

    ACTION="install"
    DRIVER_PACKAGE="${OPEN_SOURCE_DRIVER_PACKAGE}"
    KERNEL_VERSION="7.0"
    set_open_source_repository
}

select_driver_package() {
    local driver_choice=""

    # Non-interactive mode: action pre-set via command-line flags
    if [ -n "$ACTION" ]; then
        case "$DRIVER_VARIANT" in
            closed)
                if ! closed_source_available; then
                    error_exit "Closed-source driver is not supported on $ID ($VERSION_CODENAME)"
                fi
                set_closed_source_install
                ;;
            full)
                if ! full_install_available; then
                    error_exit "--install-full is only supported on Debian 12 (bookworm)"
                fi
                set_closed_source_install
                ;;
            opensource)
                set_open_source_install
                ;;
        esac
        return 0
    fi

    if [ ! -r /dev/tty ]; then
        if full_install_available; then
            DRIVER_VARIANT="full"
            set_closed_source_install
        elif open_source_available; then
            set_open_source_install
        else
            DRIVER_VARIANT="closed"
            set_closed_source_install
        fi
        warn "No interactive terminal detected. Defaulting to ${DRIVER_PACKAGE}."
        return 0
    fi

    echo ""
    echo "Please select the driver package version to install:"
    if full_install_available; then
        echo "  1) ${DRIVER_PACKAGE_BASE}  (Closed-source solution, Linux 6.6)"
        echo "  2) ${DRIVER_PACKAGE_BASE}-full  (Full closed-source solution, Linux 6.6)"
        echo "  3) Uninstall CIX packages and configuration"
        echo ""
        while true; do
            read -r -p "Enter your choice [1-3, default 2]: " driver_choice < /dev/tty
            driver_choice=${driver_choice:-2}
            case "$driver_choice" in
                1)
                    DRIVER_VARIANT="closed"
                    set_closed_source_install
                    break
                    ;;
                2)
                    DRIVER_VARIANT="full"
                    set_closed_source_install
                    break
                    ;;
                3)
                    ACTION="uninstall"
                    break
                    ;;
                *)
                    warn "Invalid option, please try again"
                    ;;
            esac
        done
        return 0
    fi

    if open_source_available; then
        echo "  1) ${DRIVER_PACKAGE_BASE}  (Closed-source solution, Linux 6.6)"
        echo "  2) ${OPEN_SOURCE_DRIVER_PACKAGE}  (Open-source solution, Linux 7.0)"
        echo "  3) Uninstall CIX packages and configuration"
        echo ""
        while true; do
            read -r -p "Enter your choice [1-3, default 2]: " driver_choice < /dev/tty
            driver_choice=${driver_choice:-2}
            case "$driver_choice" in
                1)
                    DRIVER_VARIANT="closed"
                    set_closed_source_install
                    break
                    ;;
                2)
                    set_open_source_install
                    break
                    ;;
                3)
                    ACTION="uninstall"
                    break
                    ;;
                *)
                    warn "Invalid option, please try again"
                    ;;
            esac
        done
        return 0
    fi

    if ! closed_source_available; then
        error_exit "No supported CIX driver package is available for $ID ($VERSION_CODENAME)"
    fi

    echo "  1) ${DRIVER_PACKAGE_BASE}  (Closed-source solution, Linux 6.6)"
    echo "  2) Uninstall CIX packages and configuration"
    echo ""
    while true; do
        read -r -p "Enter your choice [1-2, default 1]: " driver_choice < /dev/tty
        driver_choice=${driver_choice:-1}
        case "$driver_choice" in
            1)
                DRIVER_VARIANT="closed"
                set_closed_source_install
                break
                ;;
            2)
                ACTION="uninstall"
                break
                ;;
            *)
                warn "Invalid option, please try again"
                ;;
        esac
    done
}

find_latest_target_kernel() {
    find /boot -maxdepth 1 -type f -name "$(basename "$KERNEL_BOOT_GLOB")" -printf '%p\n' 2>/dev/null \
        | awk '$0 !~ /-rc[0-9]+-/' \
        | sort -V \
        | tail -n 1
}

set_driver_package_base() {
    case "$ID" in
        debian)
            case "$VERSION_CODENAME" in
                trixie)
                    DRIVER_PACKAGE_BASE="cix-debian13-k6.6.89-driver"
                    ;;
                bookworm)
                    DRIVER_PACKAGE_BASE="cix-debian12-k6.6.89-driver"
                    ;;
            esac
            ;;
        ubuntu)
            case "$VERSION_CODENAME" in
                noble)
                    DRIVER_PACKAGE_BASE="cix-ubuntu2404-driver"
                    ;;
                plucky)
                    DRIVER_PACKAGE_BASE="cix-ubuntu2504-driver"
                    ;;
                resolute)
                    DRIVER_PACKAGE_BASE="cix-ubuntu2604-driver"
                    ;;
            esac
            ;;
    esac
}

configure_repository_for_detected_os() {
    case "$ID" in
        debian)
            case "$VERSION_CODENAME" in
                trixie)
                    DIST_CODENAME="trixie"
                    REPO_URL="https://archive.cixtech.com/debian"
                    REPO_NAME="cix-deb-repo"
                    success "Debian 13 (trixie) detected"
                    ;;
                bookworm)
                    DIST_CODENAME="bookworm"
                    REPO_URL="https://archive.cixtech.com/debian"
                    REPO_NAME="cix-deb-repo"
                    success "Debian 12 (bookworm) detected"
                    ;;
                *)
                    error_exit "Unsupported Debian version: $VERSION_CODENAME"
                    ;;
            esac
            ;;
        ubuntu)
            case "$VERSION_CODENAME" in
                noble)
                    DIST_CODENAME="noble"
                    REPO_URL="https://archive.cixtech.com/ubuntu"
                    REPO_NAME="cix-ubuntu-repo"
                    success "Ubuntu 24.04 LTS (noble) detected"
                    ;;
                plucky)
                    DIST_CODENAME="plucky"
                    REPO_URL="https://archive.cixtech.com/ubuntu"
                    REPO_NAME="cix-ubuntu-repo"
                    success "Ubuntu 25.04 (plucky) detected"
                    ;;
                resolute)
                    DIST_CODENAME="resolute"
                    REPO_URL="https://archive.cixtech.com/ubuntu"
                    REPO_NAME="cix-ubuntu-repo"
                    success "Ubuntu 26.04 (resolute) detected"
                    ;;
                *)
                    error_exit "Unsupported Ubuntu version: $VERSION_CODENAME (VERSION_ID=$VERSION_ID)"
                    ;;
            esac
            ;;
        *)
            error_exit "Unsupported distribution: $ID (Only Debian 12/13 and Ubuntu 24.04/25.04/26.04 are supported)"
            ;;
    esac
}

validate_requested_driver_variant() {
    if [ "$ACTION" = "install" ] && [ "$DRIVER_VARIANT" = "opensource" ] && ! open_source_available; then
        error_exit "Open-source driver is only supported on Debian 13 (trixie)"
    fi
}

set_repo_paths() {
    KEYRING_PATH="/usr/share/keyrings/${REPO_NAME}.gpg"
    LIST_PATH="/etc/apt/sources.list.d/${REPO_NAME}.list"
    PREFERENCES_PATH="/etc/apt/preferences.d/${REPO_NAME}"
}

collect_named_cix_packages() {
    dpkg-query -W -f='${db:Status-Abbrev}\t${binary:Package}\t${Version}\n' 2>/dev/null | awk '
        $1 != "ii" { next }
        $2 ~ /cix/ { print $2 }
    ' | sort -u
}

collect_packages_requiring_repo_reinstall() {
    dpkg-query -W -f='${db:Status-Abbrev}\t${binary:Package}\t${Version}\n' 2>/dev/null | awk '
        $1 != "ii" { next }
        $2 ~ /cix/ { next }
        $3 ~ /(^|[+~.-])cix([+~.-]|$)/ { print $2 }
    ' | sort -u
}

remove_named_cix_packages() {
    local package_list
    package_list=$(collect_named_cix_packages)

    if [ -n "$package_list" ]; then
        echo "➡ Removing installed packages with 'cix' in the name: $package_list"
        if ! echo "$package_list" | xargs apt purge -y; then
            warn "Failed to remove one or more installed CIX-named packages."
            return 1
        fi
        UNINSTALL_REMOVED_PACKAGES=1
        success "Installed CIX-named packages removed successfully"
    else
        warn "No installed packages with 'cix' in the name were found."
    fi

    return 0
}

remove_open_source_kernel_packages() {
    local package_list
    package_list=$(dpkg-query -W -f='${db:Status-Abbrev}\t${binary:Package}\n' 2>/dev/null | awk '
        $1 != "ii" { next }
        $2 == "linux-image-7.0.0-rc5-generic" { print $2 }
        $2 == "linux-headers-7.0.0-rc5-generic" { print $2 }
        $2 == "linux-image-7.0.0-generic" { print $2 }
        $2 == "linux-headers-7.0.0-generic" { print $2 }
        $2 ~ /^linux-image-.*-cix-generic$/ { print $2 }
        $2 ~ /^linux-headers-.*-cix-generic$/ { print $2 }
    ' | sort -u)

    if [ -z "$package_list" ]; then
        return 0
    fi

    echo "➡ Removing open-source kernel packages: $package_list"
    if ! echo "$package_list" | xargs apt purge -y; then
        warn "Failed to remove one or more open-source kernel packages."
        return 1
    fi

    UNINSTALL_REMOVED_PACKAGES=1
    success "Open-source kernel packages removed successfully"
    return 0
}

refresh_package_index() {
    echo "➡ Updating package index"
    if ! apt update; then
        warn "Failed to update package index after repository cleanup."
        return 1
    fi
    success "Package index updated successfully"
    return 0
}

reinstall_repo_replaced_packages() {
    local candidate_version=""
    local install_specs=""
    local pkg_list=""
    local pkg=""

    pkg_list=$(collect_packages_requiring_repo_reinstall)

    if [ -z "$pkg_list" ]; then
        return 0
    fi

    for pkg in $pkg_list; do
        candidate_version=$(apt-cache policy "$pkg" 2>/dev/null | awk '
            /^[[:space:]]*Version table:/ {
                in_table=1
                next
            }
            !in_table {
                next
            }
            {
                line=$0
                gsub(/^[[:space:]]+/, "", line)
                sub(/^\*\*\*[[:space:]]+/, "", line)
                split(line, fields, /[[:space:]]+/)
                version=fields[1]
                priority=fields[2]
                if (version == "" || priority !~ /^[0-9]+$/) {
                    next
                }
                if (version !~ /(^|[+~.-])cix([+~.-]|$)/) {
                    print version
                    exit
                }
            }
        ')
        if [ -z "$candidate_version" ] || [ "$candidate_version" = "(none)" ]; then
            warn "No non-CIX repository version found for $pkg after removing the repository."
            return 1
        fi
        install_specs="$install_specs ${pkg}=${candidate_version}"
    done

    echo "➡ Reinstalling packages that were replaced by CIX repository versions:$install_specs"
    if ! echo "$install_specs" | xargs apt install -y --allow-downgrades --allow-change-held-packages; then
        warn "Failed to reinstall one or more packages from the default repository."
        return 1
    fi

    success "Default repository packages reinstalled successfully"
    return 0
}

cleanup_unused_dependencies() {
    if [ "$UNINSTALL_REMOVED_PACKAGES" -ne 1 ]; then
        return 0
    fi

    echo "➡ Removing automatically installed dependencies"
    if ! apt autoremove --purge -y; then
        warn "Automatic dependency cleanup encountered issues."
    else
        success "Automatic dependency cleanup completed successfully"
    fi

    return 0
}

remove_repo_configuration() {
    local removed_any=0

    if [ -f "$LIST_PATH" ]; then
        echo "➡ Removing APT source list $LIST_PATH"
        if ! rm -f "$LIST_PATH"; then
            warn "Failed to remove $LIST_PATH."
            return 1
        fi
        removed_any=1
    fi

    if [ -f "$PREFERENCES_PATH" ]; then
        echo "➡ Removing APT pin file $PREFERENCES_PATH"
        if ! rm -f "$PREFERENCES_PATH"; then
            warn "Failed to remove $PREFERENCES_PATH."
            return 1
        fi
        removed_any=1
    fi

    if [ -f "$KEYRING_PATH" ]; then
        echo "➡ Removing keyring $KEYRING_PATH"
        if ! rm -f "$KEYRING_PATH"; then
            warn "Failed to remove $KEYRING_PATH."
            return 1
        fi
        removed_any=1
    fi

    if [ "$removed_any" -eq 1 ]; then
        success "Repository configuration removed successfully"
    else
        warn "No repository configuration added by cix-repo-community was found."
    fi

    return 0
}

configure_apt_pin_priority() {
    local repo_host=""
    local pin_priority=""

    case "$REPO_URL:$DIST_CODENAME" in
        https://archive.cixtech.com/debian:trixie)
            pin_priority="100"
            ;;
        https://archive.cixtech.com/debian:bookworm)
            pin_priority="1001"
            ;;
        *)
            return 0
            ;;
    esac

    if [ ! -d "/etc/apt/preferences.d" ]; then
        mkdir -p /etc/apt/preferences.d || error_exit "Failed to create APT preferences directory"
    fi

    repo_host="${REPO_URL#*://}"
    repo_host="${repo_host%%/*}"

    echo "➡ Configuring APT pin priority"
    cat > "$PREFERENCES_PATH" <<EOF
Package: *
Pin: origin "$repo_host"
Pin-Priority: $pin_priority
EOF

    chmod 644 "$PREFERENCES_PATH"
    success "APT pin priority configured successfully"
}

uninstall_cix_repo() {
    echo "➡ Uninstalling CIX packages and configuration"
    UNINSTALL_REMOVED_PACKAGES=0

    remove_named_cix_packages || return 1
    remove_open_source_kernel_packages || return 1

    remove_repo_configuration || return 1
    refresh_package_index || return 1
    reinstall_repo_replaced_packages || return 1
    cleanup_unused_dependencies || true

    # cix-grub-config is already purged by now; fall back to the system tools
    # so the resulting grub.cfg reflects the removed CIX kernels regardless of
    # dpkg ordering or whether zz-update-grub ran.
    echo ""
    echo "➡ Refreshing GRUB configuration"
    if command -v update-grub >/dev/null 2>&1; then
        update-grub || warn "Failed to refresh GRUB configuration; reboot and verify manually"
    elif command -v grub-mkconfig >/dev/null 2>&1; then
        grub-mkconfig -o /boot/grub/grub.cfg || warn "Failed to refresh GRUB configuration; reboot and verify manually"
    else
        warn "Neither update-grub nor grub-mkconfig was found; skipping GRUB refresh"
    fi

    echo ""
    echo "${YELLOW}═══════════════════════════════════════════════════${NC}"
    echo "${GREEN}🔄 Reboot to complete driver removal${NC}"
    echo "${YELLOW}═══════════════════════════════════════════════════${NC}"
    echo ""
}

parse_args "$@"

if [ "$(id -u)" -ne 0 ]; then
    error_exit "Please run this script as root"
fi

if [ ! -f /etc/os-release ]; then
    error_exit "Unable to detect operating system"
fi

. /etc/os-release
validate_requested_driver_variant
configure_repository_for_detected_os

set_driver_package_base
select_driver_package
set_repo_paths

if [ "$ACTION" = "uninstall" ]; then
    uninstall_cix_repo
    exit 0
fi

for cmd in curl gpg; do
    if ! command -v "$cmd" >/dev/null 2>&1; then
        error_exit "Required command '$cmd' not found. Please install it first."
    fi
done

if [ -f "$LIST_PATH" ] && [ -f "$KEYRING_PATH" ]; then
    warn "Repository already configured. Reconfiguring..."
fi

if [ ! -d "/usr/share/keyrings" ]; then
    mkdir -p /usr/share/keyrings || error_exit "Failed to create keyring directory"
fi

echo "➡ Downloading and installing GPG key"
if ! curl -fsSL "$GPG_KEY_URL" | gpg --dearmor > "$KEYRING_PATH"; then
    error_exit "Failed to download or install GPG key from $GPG_KEY_URL"
fi

if [ ! -f "$KEYRING_PATH" ] || [ ! -s "$KEYRING_PATH" ]; then
    error_exit "GPG keyring file is empty or was not created"
fi

chmod 644 "$KEYRING_PATH"
success "GPG key installed successfully"

echo "➡ Configuring APT repository"
cat > "$LIST_PATH" <<EOF
deb [signed-by=$KEYRING_PATH] $REPO_URL $DIST_CODENAME main
EOF

success "Repository configured successfully"

configure_apt_pin_priority

echo "➡ Updating package index"
if ! apt update; then
    error_exit "Failed to update package index. Please check your network connection and repository configuration."
fi

success "Package index updated successfully"

echo ""
echo "${GREEN}🎉 Repository setup completed successfully!${NC}"
echo ""

echo "➡ Installing CIX driver package"

if [ -n "$DRIVER_PACKAGE" ]; then
    echo "Installing $DRIVER_PACKAGE..."
    if ! apt install -y --allow-downgrades "$DRIVER_PACKAGE"; then
        error_exit "Failed to install $DRIVER_PACKAGE"
    fi
    success "$DRIVER_PACKAGE installed successfully"

    grub_refreshed=0
    echo ""
    if [ -x /usr/libexec/cix-grub-config-refresh ]; then
        echo "➡ Refreshing GRUB configuration"
        if /usr/libexec/cix-grub-config-refresh; then
            success "GRUB configuration refreshed"
            grub_refreshed=1
        else
            warn "Failed to refresh GRUB configuration; reboot and verify manually"
        fi
    fi

    # Detect the target kernel only after refresh so the printed name matches
    # what 09_cix_grub_config actually wrote to grub.cfg.
    GRUB_CIX_KERNEL=$(find_latest_target_kernel)

    echo ""
    echo "${YELLOW}═══════════════════════════════════════════════════${NC}"
    echo "${GREEN}🔄 Reboot to complete driver installation${NC}"
    echo "${YELLOW}═══════════════════════════════════════════════════${NC}"
    echo ""
    if [ "$grub_refreshed" -eq 1 ] && [ -n "$GRUB_CIX_KERNEL" ]; then
        echo "CIX kernel ${GREEN}$(basename "$GRUB_CIX_KERNEL")${NC} is now the first GRUB boot option"
    elif [ -n "$GRUB_CIX_KERNEL" ]; then
        echo "Detected target kernel ${GREEN}$(basename "$GRUB_CIX_KERNEL")${NC}; reboot and verify it is the selected GRUB kernel"
    else
        echo "Please reboot and verify ${GREEN}Linux $KERNEL_VERSION${NC} is the selected GRUB kernel"
    fi
    echo "${YELLOW}═══════════════════════════════════════════════════${NC}"
    echo ""
else
    warn "No driver package found for $ID $VERSION_CODENAME"
    echo ""
    echo "You can now install packages from the CIX repository:"
    echo "  apt search <package-name>"
    echo "  apt install <package-name>"
fi