#!/bin/bash

# AstraLinux SE 1.6, AstraLinux 1.7/1.8 - nvidia binary driver installer
# ======================================================================
# Author: Vladimir Silantev <support@astralinux.ru>
set -ue
. gettext.sh
export TEXTDOMAINDIR=/usr/share/locale
export TEXTDOMAIN="astra-nvidia-install"

# Function to check NVIDIA GPU architecture
check_nvidia_gpu_architecture() {
    echo "$(gettext "Checking NVIDIA GPU...")"

    # Check for lspci
    if ! command -v lspci >/dev/null 2>&1; then
        echo "$(gettext "Installing lspci utility...")"
        apt install --yes pciutils 2>/dev/null || true
    fi

    # Find NVIDIA GPUs via lspci
    local nvidia_cards=$(lspci -nn 2>/dev/null | grep -i nvidia | grep -E "VGA|3D" || true)

    if [ -z "$nvidia_cards" ]; then
        echo "$(gettext "NVIDIA GPU not found in the system.")"
        return 2  # Code 2: GPU not found
    fi

    echo "$(gettext "Found NVIDIA GPUs:")"
    echo "$nvidia_cards"

    # Extract device IDs
    local device_ids=$(echo "$nvidia_cards" | grep -oE '\[10de:[0-9a-f]{4}\]' | sed 's/\[10de://;s/\]//' | tr '[:lower:]' '[:upper:]')

    if [ -z "$device_ids" ]; then
        echo "$(gettext "Failed to determine GPU IDs.")"
        return 3
    fi

    eval_gettext "GPU IDs: \$device_ids"; echo

    # Determine architecture by device IDs
    local has_maxwell_or_pascal=false
    local found_cards=""

    # Device ID lists for Maxwell and Pascal architectures
    # Source: https://devicehunt.com/all-pci-vendors/10DE/device/class/030000
    # Maxwell (series 7xx, 8xx, 9xx, Titan X (Maxwell))
    local maxwell_ids="1380 1381 1382 138A 1390 1391 1392 139A 139B 139C 139D 13B0 13B1 13B2 13B3 13B4 13B6 13B9 13BA 13BB 13BC 13BD 13C0 13C2 13C3 13D7 13D8 13D9 13DA 13DC 13DD 13DE 13DF 1401 1402 1406 1407 140B 1427 1428 1430 1431 1436 1437 1438 143A 143B 143C 143D 143F 1440 1441 1442 1443 1444 1445 1446 1447 144A 144B 144C 144D 144E 144F 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 145A 145B 145C 145D 145E 145F 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 146A 146B 146C 146D 146E 146F 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 147A 147B 147C 147D 147E 147F 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 148A 148B 148C 148D 148E 148F 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 149A 149B 149C 149D 149E 149F 14A0 14A1 14A2 14A3 14A4 14A5 14A6 14A7 14A8 14A9 14AA 14AB 14AC 14AD 14AE 14AF 14B0 14B1 14B2 14B3 14B4 14B5 14B6 14B7 14B8 14B9 14BA 14BB 14BC 14BD 14BE 14BF 14C0 14C1 14C2 14C3 14C4 14C5 14C6 14C7 14C8 14C9 14CA 14CB 14CC 14CD 14CE 14CF 14D0 14D1 14D2 14D3 14D4 14D5 14D6 14D7 14D8 14D9 14DA 14DB 14DC 14DD 14DE 14DF 14E0 14E1 14E2 14E3 14E4 14E5 14E6 14E7 14E8 14E9 14EA 14EB 14EC 14ED 14EE 14EF 14F0 14F1 14F2 14F3 14F4 14F5 14F6 14F7 14F8 14F9 14FA 14FB 14FC 14FD 14FE 14FF"

    # Pascal (series 10xx, Titan X (Pascal), Titan Xp)
    local pascal_ids="15F0 15F1 15F7 15F8 15F9 15FA 15FB 15FC 15FD 1617 1618 1619 161A 161D 161E 161F 1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 162A 162B 162C 162D 162E 162F 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 163A 163B 163C 163D 163E 163F 1640 1641 1642 1643 1644 1645 1646 1647 1648 1649 164A 164B 164C 164D 164E 164F 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 165A 165B 165C 165D 165E 165F 1660 1661 1662 1663 1664 1665 1666 1667 1668 1669 166A 166B 166C 166D 166E 166F 1670 1671 1672 1673 1674 1675 1676 1677 1678 1679 167A 167B 167C 167D 167E 167F 1680 1681 1682 1683 1684 1685 1686 1687 1688 1689 168A 168B 168C 168D 168E 168F 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 169A 169B 169C 169D 169E 169F 16A0 16A1 16A2 16A3 16A4 16A5 16A6 16A7 16A8 16A9 16AA 16AB 16AC 16AD 16AE 16AF 16B0 16B1 16B2 16B3 16B4 16B5 16B6 16B7 16B8 16B9 16BA 16BB 16BC 16BD 16BE 16BF 16C0 16C1 16C2 16C3 16C4 16C5 16C6 16C7 16C8 16C9 16CA 16CB 16CC 16CD 16CE 16CF 16D0 16D1 16D2 16D3 16D4 16D5 16D6 16D7 16D8 16D9 16DA 16DB 16DC 16DD 16DE 16DF 16E0 16E1 16E2 16E3 16E4 16E5 16E6 16E7 16E8 16E9 16EA 16EB 16EC 16ED 16EE 16EF 16F0 16F1 16F2 16F3 16F4 16F5 16F6 16F7 16F8 16F9 16FA 16FB 16FC 16FD 16FE 16FF"

    # Also check by model name
    for card_info in $nvidia_cards; do
        local card_name=$(echo "$card_info" | grep -o "GeForce.*" || true)
        if [ -n "$card_name" ]; then
            echo "  - $card_name"
            found_cards+="$card_name, "

            # Check by model name (heuristic)
            if [[ "$card_name" =~ (GTX\s*[789][0-9]{2}|Titan\s*X\s*(Maxwell)?|GTX\s*10[0-9]{2}|Titan\s*X\s*(Pascal|Xp)) ]]; then
                has_maxwell_or_pascal=true
                echo "    $(gettext "Detected Maxwell or Pascal GPU")"
            fi
        fi
    done

    # Check by device IDs
    for device_id in $device_ids; do
        # Check Maxwell
        if echo "$maxwell_ids" | grep -qw "$device_id"; then
            has_maxwell_or_pascal=true
            eval_gettext "  Detected Maxwell architecture GPU (ID: \$device_id)"; echo
        fi

        # Check Pascal
        if echo "$pascal_ids" | grep -qw "$device_id"; then
            has_maxwell_or_pascal=true
            eval_gettext "  Detected Pascal architecture GPU (ID: \$device_id)"; echo
        fi
    done

    if [ "$has_maxwell_or_pascal" = true ]; then
        echo "$(gettext "WARNING: Detected Maxwell or Pascal architecture GPU!")"
        echo "$(gettext "  - Maxwell: GTX 7xx, 8xx, 9xx, Titan X (Maxwell) series")"
        echo "$(gettext "  - Pascal: GTX 10xx, Titan X (Pascal), Titan Xp series")"
        return 1  # Code 1: Maxwell/Pascal GPU found
    else
        echo "$(gettext "GPU is NOT Maxwell or Pascal architecture.")"
        return 0  # Code 0: GPU is not Maxwell/Pascal
    fi
}

# Function to check driver version and GPU compatibility
check_driver_compatibility() {
    local driver_version="$1"
    local script_name=$(basename "$0")

    # Check only for versions 590 and above
    if [[ "$driver_version" -ge 590 ]]; then
        echo "========================================"
        eval_gettext "WARNING: Driver version \$driver_version"; echo
        echo "========================================"
        echo ""
        echo "$(gettext "STARTING FROM VERSION 590, NVIDIA DRIVERS STOPPED SUPPORTING:")"
        echo "$(gettext "  • Maxwell architecture GPUs (GTX 7xx, 8xx, 9xx, Titan X series)")"
        echo "$(gettext "  • Pascal architecture GPUs (GTX 10xx, Titan X (Pascal), Titan Xp series)")"
        echo ""
        echo "$(gettext "Installing this driver on these GPUs will lead to:")"
        echo "$(gettext "  • Graphics subsystem failure")"
        echo "$(gettext "  • Inability to load graphical interface")"
        echo "$(gettext "  • System recovery necessity")"
        echo ""

        # Check GPU in the system
        echo "$(gettext "Checking installed GPU...")"
        check_nvidia_gpu_architecture
        local gpu_check_result=$?

        case $gpu_check_result in
            0)
                echo ""
                echo "✓ $(gettext "Your GPU is NOT Maxwell/Pascal.")"
                eval_gettext "  You can safely install driver version \$driver_version."; echo
                ;;
            1)
                echo ""
                echo "❌ $(gettext "DANGER: You have Maxwell or Pascal architecture GPU!")"
                eval_gettext "  Installing driver version \$driver_version will make the system unusable!"; echo
                echo ""
                echo "$(gettext "Recommended to install driver version 535 or lower.")"
                echo ""

                # Request confirmation with enhanced warning
                echo "========================================"
                echo "$(gettext "DO YOU REALLY WANT TO CONTINUE?")"
                echo "========================================"
                echo "$(gettext "This will lead to loss of graphical interface!")"
                echo "$(gettext "Recovery will require booting in recovery mode")"
                echo "$(gettext "and rolling back the driver via command line.")"
                echo ""

                read -p "$(eval_gettext "Continue installing DRIVER VERSION \$driver_version? (yes/NO): ")" confirm_force
                confirm_force=$(echo "$confirm_force" | tr '[:upper:]' '[:lower:]')

                if [ "$confirm_force" != "yes" ]; then
                    echo ""
                    echo "$(gettext "Installation canceled. Recommended actions:")"
                    echo "$(gettext "  1. Choose driver version 535 or lower")"
                    echo "  2. $(gettext "Use option:") $script_name -l $(gettext "to view available drivers")"
                    echo "  3. $(gettext "Install:") $script_name -d nvidia-driver-535"
                    exit 1
                else
                    echo ""
                    echo "⚠️  $(gettext "WARNING: You confirmed installation despite the risk!")"
                    echo "  $(gettext "Don't forget to backup important data.")"
                    echo "  $(gettext "Be ready for system recovery.")"
                    echo ""
                    read -p "$(gettext "Press Enter to continue or Ctrl+C to cancel...")" dummy
                fi
                ;;
            2)
                echo ""
                echo "⚠️  $(gettext "Failed to detect NVIDIA GPU.")"
                echo "  $(gettext "Driver installation will continue, but make sure you don't have Maxwell/Pascal GPU.")"
                echo "  $(gettext "If you have one - the system will become unusable!")"
                echo ""

                read -p "$(gettext "Are you sure you DON'T have GTX 7xx, 8xx, 9xx, 10xx series GPU? (Y/N): ")" confirm_no_card
                if [[ ! "$confirm_no_card" =~ ^[Yy]$ ]]; then
                    echo "$(gettext "Installation canceled.")"
                    exit 1
                fi
                ;;
            *)
                echo ""
                eval_gettext "⚠️  Failed to check GPU (error code: \$gpu_check_result)."; echo
                echo "  $(gettext "Installation will continue at your own risk.")"
                echo ""

                read -p "$(eval_gettext "Continue installing driver version \$driver_version? (Y/N): ")" confirm_risk
                if [[ ! "$confirm_risk" =~ ^[Yy]$ ]]; then
                    echo "$(gettext "Installation canceled.")"
                    exit 1
                fi
                ;;
        esac

        echo ""
        echo "========================================"
    fi
}

# Function for quiet APT update
quiet_apt_update() {
    echo "$(gettext "Updating package lists...")"

    # Execute apt update and filter output
    apt update 2>&1 | \
    while IFS= read -r line; do
        # Skip empty lines
        [ -z "$line" ] && continue

        # Skip WARNING about CLI interface
        [[ "$line" =~ "WARNING: apt does not have a stable CLI interface" ]] && continue

        # Skip repository status lines
        [[ "$line" =~ ^(Ign:|Hit:|Get:|Err:) ]] && continue

        # Skip progress lines
        [[ "$line" =~ (Reading|Building|Fetched|Downloading|Updating) ]] && continue

        # Skip technical messages
        [[ "$line" =~ (Reading|Building) ]] && continue

        # Skip available updates message
        [[ "$line" =~ "can be upgraded" ]] && continue

        # Output the rest (errors, warnings, summaries)
        echo "$line"
    done

    # Return apt update status
    return ${PIPESTATUS[0]}
}

# Function to show help
show_help() {
    local script_name=$(basename "$0")

    echo "$(gettext "Usage:") $script_name $(gettext "[OPTIONS] [DRIVER]")"
    echo ""
    echo "$(gettext "Install NVIDIA drivers in Astra Linux")"
    echo ""
    echo "$(gettext "IMPORTANT: Starting from version 590, NVIDIA drivers don't support")"
    echo "$(gettext "           Maxwell (GTX 7xx, 8xx, 9xx) and Pascal (GTX 10xx) GPUs!")"
    echo ""
    echo "$(gettext "Options:")"
    echo "  -h, --help     $(gettext "Show this help and exit")"
    echo "  -l, --list     $(gettext "Show list of available drivers and exit")"
    echo "  -v, --version  $(gettext "Show script version")"
    echo "  -d, --driver   $(gettext "Install specified driver (examples below)")"
    echo ""
    echo "$(gettext "Operation modes:")"
    echo "  $(gettext "1. Without options: interactive selection from available drivers")"
    echo "  $(gettext "2. With -d option: install specified driver")"
    echo ""
    echo "$(gettext "Examples with -d option:")"
    echo "  $script_name -d nvidia-driver-560       $(gettext "# Install nvidia-driver-560")"
    echo "  $script_name -d nvidia-driver-560-open  $(gettext "# Install open variant")"
    echo "  $script_name -d nvidia-no-dkms-340      $(gettext "# Install without DKMS (old format)")"
    echo "  $script_name -d nvidia-340              $(gettext "# Install old format driver")"
    echo ""
    echo "$(gettext "If specified driver is not found, available drivers will be shown.")"
    exit 0
}

# Function to show version
show_version() {
    echo "astra-nvidia-install version 2.0.2"
    echo "$(gettext "Added compatibility check for driver version 590+")"
    echo "$(gettext "For Astra Linux SE 1.6, Astra Linux 1.7/1.8 and above")"
    echo "$(gettext "Support: Vladimir Silantev <support@astralinux.ru>")"
    exit 0
}

# Function to check Extended repository compatibility
check_extended_repo_compatibility() {
    echo "$(gettext "Checking Extended repository...")"

    # Determine Astra Linux version
    if [ ! -f /etc/astra_version ]; then
        echo "$(gettext "File /etc/astra_version not found.")"
        return 1
    fi

    local astra_version=$(cat /etc/astra_version)
    eval_gettext "System version: \$astra_version"; echo

    # Extract major version (first two digits: 1.7, 1.8, etc.)
    local major_version=$(echo "$astra_version" | grep -o '^[0-9]\+\.[0-9]\+')

    # For Astra 1.8 and above Extended repository is built-in, skip check
    if [[ "$major_version" > "1.7" ]] || [[ "$major_version" == "1.8" ]]; then
        eval_gettext "For version \$astra_version Extended repository is built-in."; echo
        return 0
    fi

    # Only for 1.7 check Extended repository
    if [[ "$major_version" != "1.7" ]]; then
        eval_gettext "Version \$astra_version doesn't require Extended repository check."; echo
        return 0
    fi

    # Find ALL Extended repositories in the system
    local extended_repos=($(grep -r "extended-repository" /etc/apt/sources.list /etc/apt/sources.list.d/ 2>/dev/null | \
                          grep -v "^Binary" | awk '{print $2}' | sort -u))

    if [ ${#extended_repos[@]} -eq 0 ]; then
        echo "$(gettext "Extended repository not found in the system.")"
        echo "$(gettext "No drivers available.")"
        return 1
    fi

    eval_gettext "Found Extended repositories: \${#extended_repos[@]}"; echo

    local selected_repo=""
    local selected_version=""
    local max_version=""
    local found_valid_repo=0

    # Check each repository and select the newest one
    for repo_url in "${extended_repos[@]}"; do
        local repo_name=$(basename "$repo_url")
        eval_gettext "Checking repository: \$repo_name"; echo

        # Try to get version via version file (new way)
        local version_url=""
        if [[ "$repo_url" == *"/dists/1.7_x86-64/"* ]]; then
            # Replace path part to version file
            version_url="${repo_url%/dists/1.7_x86-64/*}/version"
        else
            # Try alternative path
            version_url="${repo_url%/}/../version"
        fi

        local repo_version_info=""

        if command -v curl >/dev/null 2>&1; then
            repo_version_info=$(curl -s --max-time 10 "$version_url" 2>/dev/null | head -1)
        elif command -v wget >/dev/null 2>&1; then
            repo_version_info=$(wget -qO- --timeout=10 "$version_url" 2>/dev/null | head -1)
        fi

        # If version file didn't work, try old way via InRelease
        if [ -z "$repo_version_info" ] || [[ ! "$repo_version_info" =~ ^[0-9] ]]; then
            local inrelease_url="${repo_url%/}/dists/1.7_x86-64/InRelease"
            if command -v curl >/dev/null 2>&1; then
                repo_version_info=$(curl -s --max-time 10 "$inrelease_url" 2>/dev/null | grep -i "^Version:" | head -1)
            elif command -v wget >/dev/null 2>&1; then
                repo_version_info=$(wget -qO- --timeout=10 "$inrelease_url" 2>/dev/null | grep -i "^Version:" | head -1)
            fi

            if [ -n "$repo_version_info" ]; then
                # Extract version from InRelease
                repo_version_info=$(echo "$repo_version_info" | sed 's/Version:\s*//i' | tr -d '[:space:]')
            fi
        fi

        if [ -z "$repo_version_info" ] || [[ ! "$repo_version_info" =~ ^[0-9] ]]; then
            echo "  $(gettext "Failed to get repository version.")"
            continue
        fi

        # Extract base version from repo_version (e.g., 1.7.8 from 1.7.8.EXT1.4)
        local repo_base_version=$(echo "$repo_version_info" | grep -o '^[0-9]\+\.[0-9]\+\.[0-9]\+')

        if [ -z "$repo_base_version" ]; then
            echo "  $(gettext "Failed to determine base version.")"
            continue
        fi

        eval_gettext "  Repository version: \$repo_version_info (base: \$repo_base_version)"; echo

        # Check if base version matches system
        if [ "$astra_version" != "$repo_base_version" ]; then
            eval_gettext "  Doesn't match system version (\$astra_version)."; echo
            continue
        fi

        found_valid_repo=1

        # Compare full versions to select maximum
        if [ -z "$max_version" ] || [[ "$repo_version_info" > "$max_version" ]]; then
            max_version="$repo_version_info"
            selected_repo="$repo_url"
            selected_version="$repo_base_version"
        fi
    done

    if [ $found_valid_repo -eq 0 ]; then
        eval_gettext "No Extended repository found matching system version (\$astra_version)."; echo
        echo "$(gettext "No drivers available.")"
        return 1
    fi

    if [ -z "$selected_repo" ]; then
        echo "$(gettext "Failed to select suitable Extended repository.")"
        echo "$(gettext "No drivers available.")"
        return 1
    fi

    eval_gettext "Selected newest Extended repository: \$max_version"; echo
    eval_gettext "✓ Extended repository matches system version (\$astra_version)."; echo
    return 0
}

# Function to show driver list
# Function to show driver list
show_driver_list() {
    local script_name=$(basename "$0")

    echo "$(gettext "List of available NVIDIA drivers in Astra Linux repositories")"
    echo "============================================================"
    echo ""
    echo "$(gettext "IMPORTANT: Starting from version 590, NVIDIA drivers don't support")"
    echo "$(gettext "           Maxwell (GTX 7xx, 8xx, 9xx) and Pascal (GTX 10xx) GPUs!")"
    echo ""

    # For -l mode check Extended repository
    if ! check_extended_repo_compatibility; then
        exit 1
    fi

    # Update package lists
    quiet_apt_update

    # Find all NVIDIA metapackages
    echo ""
    echo "$(gettext "Searching for NVIDIA driver metapackages...")"
    local meta_packages=$(apt-cache search --names-only "^nvidia-" | \
                         cut -d' ' -f1 | sort -u -V -r)

    if [ -z "$meta_packages" ]; then
        echo "$(gettext "No NVIDIA metapackages found!")"
        echo "$(gettext "Possible reasons:")"
        echo "$(gettext "  1. Non-free repository missing")"
        echo "$(gettext "  2. Extended repository doesn't match system version")"
        echo "$(gettext "  3. No internet connection")"
        echo "$(gettext "  4. 'apt update' not executed")"
        exit 1
    fi

    # Group and filter packages
    declare -A packages_by_version

    echo "$(gettext "Filtering packages...")"

    while read -r pkg; do
        # Skip xserver-xorg-video-nvidia-* packages
        if [[ "$pkg" == xserver-xorg-video-nvidia-* ]]; then
            continue
        fi

        # Get description
        local full_desc=$(apt-cache show "$pkg" 2>/dev/null | \
                         grep "Description:" | \
                         head -1 | \
                         cut -d: -f2- | \
                         sed 's/^[[:space:]]*//;s/[[:space:]]*$//')

        [ -z "$full_desc" ] && continue

        # Skip transitional and development packages
        if [[ "$full_desc" == *"Transitional"* ]] || [[ "$full_desc" == *"transitional"* ]] || \
           [[ "$full_desc" == *"Xorg driver development"* ]] || [[ "$full_desc" == *"Development"* ]]; then
            continue
        fi

        # Extract version
        local version=""
        local pkg_clean="${pkg%%_*}"

        if [[ "$pkg_clean" =~ nvidia-driver-no-dkms-([0-9]+) ]]; then
            version="${BASH_REMATCH[1]}"
        elif [[ "$pkg_clean" =~ nvidia-driver-([0-9]+) ]]; then
            version="${BASH_REMATCH[1]}"
        elif [[ "$pkg_clean" =~ nvidia-no-dkms-([0-9]+) ]]; then
            version="${BASH_REMATCH[1]}"
        elif [[ "$pkg_clean" =~ nvidia-([0-9]+) ]]; then
            version="${BASH_REMATCH[1]}"
        elif [[ "$pkg_clean" =~ nvidia-driver-([0-9]+)-open ]]; then
            version="${BASH_REMATCH[1]}"
        fi

        if [ -z "$version" ]; then
            continue
        fi

        # Add to filtered list
        packages_by_version["$version"]+="$pkg "
    done <<< "$meta_packages"

    # Sort versions ASCENDING (old at top, new at bottom)
    local sorted_versions=$(for v in "${!packages_by_version[@]}"; do echo "$v"; done | sort -n)

    echo ""
    echo "$(gettext "Available drivers grouped by version (from old to new):")"
    echo "---------------------------------------------"

    # Исправленный подсчет
    local total_filtered=0

    for version in $sorted_versions; do
        echo ""
        if [[ "$version" -ge 590 ]]; then
            eval_gettext "Version \$version: ⚠️  NOT FOR MAXWELL/PASCAL!"; echo
        else
            eval_gettext "Version \$version:"; echo
        fi

        # Получаем пакеты для этой версии
        local packages_str="${packages_by_version["$version"]}"
        # Разбиваем строку на массив, удаляя пустые элементы
        local packages=()
        for pkg in $packages_str; do
            if [ -n "$pkg" ]; then
                packages+=("$pkg")
            fi
        done

        # Считаем пакеты для этой версии
        local version_count=${#packages[@]}

        # Sort packages inside version: main first, then variants
        local sorted_packages=()
        for pkg in "${packages[@]}"; do
            if [ -n "$pkg" ]; then
                if [[ "$pkg" == nvidia-driver-$version ]] || [[ "$pkg" == nvidia-driver-$version* ]]; then
                    sorted_packages=("$pkg" "${sorted_packages[@]}")
                elif [[ "$pkg" == nvidia-$version ]] || [[ "$pkg" == nvidia-$version* ]]; then
                    sorted_packages=("$pkg" "${sorted_packages[@]}")
                elif [[ "$pkg" == nvidia-no-dkms-$version ]] || [[ "$pkg" == nvidia-driver-no-dkms-$version ]]; then
                    sorted_packages=("$pkg" "${sorted_packages[@]}")
                else
                    sorted_packages+=("$pkg")
                fi
            fi
        done

        # Выводим пакеты для этой версии
        for pkg in "${sorted_packages[@]}"; do
            if [ -n "$pkg" ]; then
                local desc=""
                if apt-cache show "$pkg" >/dev/null 2>&1; then
                    desc=$(apt-cache show "$pkg" 2>/dev/null | \
                          grep -m1 "Description:" | \
                          cut -d: -f2- | \
                          sed 's/^[[:space:]]*//;s/[[:space:]]*$//' | \
                          head -c 80)
                fi

                # Display package with type indicator
                local type=""
                case "$pkg" in
                    *-open) type=" $(gettext "(open variant)")" ;;
                    *-no-dkms*)
                        if [[ "$pkg" == nvidia-driver-no-dkms-* ]]; then
                            type=" $(gettext "(new format without DKMS)")"
                        else
                            type=" $(gettext "(without DKMS)")"
                        fi
                        ;;
                    nvidia-[0-9]*) type=" $(gettext "(old format)")" ;;
                    nvidia-driver-[0-9]*) type=" $(gettext "(main)")" ;;
                esac

                echo "  • $pkg$type"
                if [ -n "$desc" ] && [[ ! "$desc" =~ [Tt]ransitional ]]; then
                    echo "      $desc..."
                fi
            fi
        done

        # Добавляем количество пакетов этой версии к общему счету
        total_filtered=$((total_filtered + version_count))
    done

    echo ""
    eval_gettext "Total metapackages found: \$total_filtered"; echo

    echo ""
    echo "$(gettext "WARNING:")"
    echo "  • $(gettext "Driver versions 590+ DON'T WORK on GPUs:")"
    echo "    - $(gettext "GTX 7xx, 8xx, 9xx, Titan X (Maxwell) - Maxwell architecture")"
    echo "    - $(gettext "GTX 10xx, Titan X (Pascal) - Pascal architecture")"
    echo ""
    echo "  • $(gettext "For these GPUs use driver version 535 or lower")"
    echo ""

    echo "$(gettext "Examples of specific driver installation:")"
    echo "  sudo $script_name -d nvidia-driver-535      $(gettext "# For Maxwell/Pascal GPUs")"
    echo "  sudo $script_name -d nvidia-driver-560      $(gettext "# Newer version")"
    echo "  sudo $script_name -d nvidia-driver-560-open $(gettext "# Open variant")"
    echo "  sudo $script_name -d nvidia-no-dkms-340     $(gettext "# For very old GPUs")"
    echo "  sudo $script_name -d nvidia-driver-no-dkms-560 $(gettext "# New format without DKMS")"
    exit 0
}

# Check if this is Astra Linux
check_astra() {
    if [ ! -f /etc/astra_version ] && ! grep -q "Astra" /etc/os-release 2>/dev/null; then
        echo "$(gettext "This script is for Astra Linux only!")"
        exit 1
    fi
}

# Check root privileges
check_root() {
    if [ x"$EUID" != "x0" ]; then
        gettext "Launch program as administrator (sudo) (root) .."; echo
        exit 1
    fi
}

# Function to check driver availability in repositories
check_driver_available() {
    local driver_name="$1"

    eval_gettext "Checking driver availability: \$driver_name"; echo

    # Check non-free repository
    if ! grep -r "non-free" /etc/apt/sources.list /etc/apt/sources.list.d/ 2>/dev/null | grep -v "^Binary" > /dev/null; then
        echo "$(gettext "ERROR: Non-free repository not found!")"
        return 1
    fi

    # Update package lists
    quiet_apt_update

    # Check driver availability
    if apt-cache show "$driver_name" >/dev/null 2>&1; then
        eval_gettext "✓ Driver \$driver_name available in repositories"; echo
        return 0
    else
        eval_gettext "✗ Driver \$driver_name not found in repositories"; echo
        return 1
    fi
}

# Function to install specific driver
install_specific_driver() {
    local driver_name="$1"

    echo "$(gettext "Specific driver installation mode:") $driver_name"
    echo "================================================="

    # Check driver availability
    if ! check_driver_available "$driver_name"; then
        echo ""
        eval_gettext "Driver '\$driver_name' not found in repositories."; echo
        echo "$(gettext "Please choose one of available drivers:")"
        echo ""
        # Call show_driver_list to display list
        show_driver_list
        exit 1
    fi

    # Extract version from driver name
    local driver_version=""
    if [[ "$driver_name" =~ nvidia-(driver-)?([0-9]+) ]] || [[ "$driver_name" =~ nvidia-(no-dkms-)?([0-9]+) ]]; then
        driver_version="${BASH_REMATCH[2]}"
        eval_gettext "Driver version: \$driver_version"; echo
    fi

    # Check compatibility for versions 590+
    if [[ -n "$driver_version" ]] && [[ "$driver_version" -ge 590 ]]; then
        check_driver_compatibility "$driver_version"
    fi

    # Check for additional packages
    echo ""
    echo "$(gettext "Searching for additional packages...")"

    # nvidia-settings
    if apt-cache show nvidia-settings >/dev/null 2>&1; then
        NVIDIA_SETTINGS="nvidia-settings"
        echo "✓ $(gettext "Found nvidia-settings")"
    else
        NVIDIA_SETTINGS=""
        echo "✗ $(gettext "nvidia-settings not found")"
    fi

    # xserver-xorg-video-nvidia-* (only for new drivers)
    if [ -n "$driver_version" ] && [[ "$driver_name" != nvidia-340 ]] && [[ "$driver_name" != nvidia-no-dkms-340 ]]; then
        XORG_PACKAGE="xserver-xorg-video-nvidia-$driver_version"
        if apt-cache show "$XORG_PACKAGE" >/dev/null 2>&1; then
            eval_gettext "✓ Found \$XORG_PACKAGE"; echo
        else
            XORG_PACKAGE=""
            eval_gettext "✗ \$XORG_PACKAGE not found"; echo
        fi
    else
        XORG_PACKAGE=""
    fi

    # Installation confirmation
    echo ""
    echo "========================================"
    echo "$(gettext "The following packages will be installed:")"
    echo "  • $driver_name"
    [ -n "$NVIDIA_SETTINGS" ] && echo "  • $NVIDIA_SETTINGS"
    [ -n "$XORG_PACKAGE" ] && echo "  • $XORG_PACKAGE"
    echo "========================================"

    read -p "$(gettext "Continue installation? (Y/n): ")" confirm_install
    if [[ "$confirm_install" =~ ^[Nn]$ ]]; then
        echo "$(gettext "Installation canceled.")"
        exit 0
    fi

    # Install packages
    echo ""
    echo "$(gettext "Installing driver...")"
    PACKAGES_TO_INSTALL="$driver_name"
    [ -n "$NVIDIA_SETTINGS" ] && PACKAGES_TO_INSTALL="$PACKAGES_TO_INSTALL $NVIDIA_SETTINGS"
    [ -n "$XORG_PACKAGE" ] && PACKAGES_TO_INSTALL="$PACKAGES_TO_INSTALL $XORG_PACKAGE"

    apt install --yes $PACKAGES_TO_INSTALL

    # Post-installation actions
    post_install_configuration
}

# Post-installation configuration function
post_install_configuration() {
    local timestamp="$(date +%s)"

    echo ""
    echo "$(gettext "Configuring system after installation...")"

    # Fix for Steam
    if [ -e /usr/lib32/i386-linux-gnu/ ] ; then
        if [ ! -e /usr/lib/i386-linux-gnu ] ; then
            ln -s /usr/lib32/i386-linux-gnu/ /usr/lib/i386-linux-gnu
            echo "✓ $(gettext "Fixed libraries for Steam")"
        fi
    fi

    # Block nouveau
    echo "$(gettext "Configuring nouveau blocking...")"
    if ! grep -qE "^\s*blacklist\s+nouveau\s*$" /etc/modprobe.d/* 2>/dev/null; then
        if [ -f /etc/modprobe.d/blacklist.conf ]; then
            cp /etc/modprobe.d/blacklist.conf "/etc/modprobe.d/blacklist.conf.$timestamp"
        fi
        echo "blacklist nouveau" >> /etc/modprobe.d/blacklist.conf
        echo "✓ $(gettext "Added nouveau blacklist")"
    fi

    if ! grep -qE "^\s*options\s+nouveau\s+modeset\s*=\s*0\s*$" /etc/modprobe.d/* 2>/dev/null; then
        if [ -f /etc/modprobe.d/blacklist.conf ]; then
            echo "options nouveau modeset=0" >> /etc/modprobe.d/blacklist.conf
            echo "✓ $(gettext "Added modeset=0 for nouveau")"
        fi
    fi

    # Remove modeset=1 if present
    if grep -qE "^\s*nouveau\s+modeset\s*=\s*1\s*$" /etc/initramfs-tools/modules 2>/dev/null; then
        sed -i."$timestamp" -e "/^\s*nouveau\s\+modeset\s*=\s*1\s*$/d" /etc/initramfs-tools/modules
        echo "✓ $(gettext "Removed modeset=1 for nouveau from initramfs")"
    fi

    # Update initramfs
    echo "$(gettext "Updating initramfs...")"
    update-initramfs -u -k all

    # Final check
    echo ""
    echo "========================================"
    echo "$(gettext "Installation check:")"
    if command -v nvidia-smi >/dev/null 2>&1; then
        nvidia-smi --query-gpu=driver_version,name --format=csv
        echo ""
        gettext "Reboot computer to apply changes"; echo
    else
        echo "$(gettext "nvidia-smi not found. Reboot might be needed.")"
        echo "$(gettext "After reboot check: nvidia-smi")"
    fi
    echo "========================================"
}

# Check non-free repositories for Astra Linux
check_astra_repos() {
    echo "$(gettext "Checking Astra Linux repositories...")"

    # Check for non-free in Astra sources
    if ! grep -r "non-free" /etc/apt/sources.list /etc/apt/sources.list.d/ 2>/dev/null | grep -v "^Binary" > /dev/null; then
        echo "$(gettext "WARNING: Non-free repository not found!")"
        echo "$(gettext "In Astra Linux, non-free repository is needed for NVIDIA drivers installation.")"
        echo ""

        # Determine Astra Linux version
        if [ -f /etc/os-release ]; then
            . /etc/os-release
            eval_gettext "Distribution: \$PRETTY_NAME"; echo

            case $VERSION_ID in
                "1.6"|"1.7")
                    eval_gettext "For Astra SE \$VERSION_ID use repository:"; echo
                    echo "deb http://update.astra.in.ua/repository/astra/stable/1.6_x86-64/ non-free"
                    ;;
                "2.12"|"2.14")
                    eval_gettext "For Astra CE \$VERSION_ID use repository:"; echo
                    echo "deb http://dl.astralinux.ru/astra/stable/2.12_x86-64/repository-main/ non-free"
                    ;;
                *)
                    eval_gettext "Version: \$VERSION_ID"; echo
                    echo "$(gettext "Refer to Astra Linux documentation for non-free repository setup.")"
                    ;;
            esac
        fi

        echo ""
        read -p "$(gettext "Add non-free repository automatically? (y/N): ")" add_repo
        if [[ "$add_repo" =~ ^[Yy]$ ]]; then
            backup_sources="/etc/apt/sources.list.backup.$(date +%s)"
            cp /etc/apt/sources.list "$backup_sources"
            eval_gettext "Backup created: \$backup_sources"; echo

            sed -i 's/^deb \(.*\) main$/deb \1 main contrib non-free/;
                    s/^deb \(.*\) main contrib$/deb \1 main contrib non-free/;
                    s/^deb \(.*\) main$/\0 contrib non-free/' /etc/apt/sources.list

            quiet_apt_update
            if [ $? -ne 0 ]; then
                echo "$(gettext "Error updating repositories. Restoring original sources.list...")"
                cp "$backup_sources" /etc/apt/sources.list
                exit 1
            fi
            echo "$(gettext "Non-free repository added.")"
        else
            echo "$(gettext "Cannot continue without non-free repository.")"
            exit 1
        fi
    else
        echo "$(gettext "Non-free repository found. Updating package lists...")"
        quiet_apt_update
    fi
}

# Function to find ALL NVIDIA driver metapackages (interactive version)
find_all_nvidia_metapackages() {
    echo "$(gettext "Searching for all available NVIDIA driver metapackages...")"
    echo "====================================================="

    # Use the same approach as in show_driver_list
    echo "$(gettext "Scanning repositories...")"

    # Get all packages with nvidia-
    local all_packages=$(apt-cache search --names-only "^nvidia-" 2>/dev/null | \
                         cut -d' ' -f1 | sort -u -V -r)

    if [ -z "$all_packages" ]; then
        echo "$(gettext "No NVIDIA packages found!")"
        echo "$(gettext "Trying extended search...")"

        all_packages=$(apt list 2>/dev/null | grep -E "^nvidia-(driver-)?[0-9]" | cut -d'/' -f1 | sort -u -V -r)

        if [ -z "$all_packages" ]; then
            echo "$(gettext "ERROR: No NVIDIA drivers found in repositories!")"
            exit 1
        fi
    fi

    # Create arrays for selection
    echo ""
    echo "$(gettext "Available NVIDIA driver metapackages:")"
    echo "======================================"

    declare -A packages_by_version
    declare -A package_descriptions
    declare -A package_types

    local i=1
    declare -A selection_map

    # Filter and group packages
    while read -r pkg; do
        # Skip xserver-xorg-video-nvidia-* packages
        if [[ "$pkg" == xserver-xorg-video-nvidia-* ]]; then
            continue
        fi

        # Get full description
        local full_desc=$(apt-cache show "$pkg" 2>/dev/null | \
                         grep "Description:" | \
                         head -1 | \
                         cut -d: -f2- | \
                         sed 's/^[[:space:]]*//;s/[[:space:]]*$//')

        [ -z "$full_desc" ] && continue

        # Skip transitional packages
        if [[ "$full_desc" == *"Transitional"* ]] || [[ "$full_desc" == *"transitional"* ]] || \
           [[ "$full_desc" == *"Xorg driver development"* ]] || [[ "$full_desc" == *"Development"* ]]; then
            continue
        fi

        # Extract version from package name
        local version=""
        local pkg_clean="${pkg%%_*}"

        # Support all formats: nvidia-driver-no-dkms-XXX, nvidia-driver-XXX, nvidia-no-dkms-XXX, nvidia-XXX
        if [[ "$pkg_clean" =~ nvidia-driver-no-dkms-([0-9]+) ]]; then
            version="${BASH_REMATCH[1]}"
        elif [[ "$pkg_clean" =~ nvidia-driver-([0-9]+) ]]; then
            version="${BASH_REMATCH[1]}"
        elif [[ "$pkg_clean" =~ nvidia-no-dkms-([0-9]+) ]]; then
            version="${BASH_REMATCH[1]}"
        elif [[ "$pkg_clean" =~ nvidia-([0-9]+) ]]; then
            version="${BASH_REMATCH[1]}"
        elif [[ "$pkg_clean" =~ nvidia-driver-([0-9]+)-open ]]; then
            version="${BASH_REMATCH[1]}"
        fi

        if [ -z "$version" ]; then
            continue
        fi

        packages_by_version["$version"]+="$pkg "

        # Save short description
        local short_desc=$(echo "$full_desc" | head -c 60)
        package_descriptions["$pkg"]="$short_desc"

        # Determine package type
        local type=""
        case "$pkg" in
            *-open) type=" $(gettext "(open variant)")" ;;
            *-no-dkms*)
                if [[ "$pkg" == nvidia-driver-no-dkms-* ]]; then
                    type=" $(gettext "(new format without DKMS)")"
                else
                    type=" $(gettext "(without DKMS)")"
                fi
                ;;
            nvidia-[0-9]*) type=" $(gettext "(old format)")" ;;
            nvidia-driver-[0-9]*) type=" $(gettext "(main)")" ;;
        esac
        package_types["$pkg"]="$type"
    done <<< "$all_packages"

    # Sort versions ASCENDING (old at top, new at bottom)
    local sorted_versions=$(for v in "${!packages_by_version[@]}"; do echo "$v"; done | sort -n)

    # Display versions from old to new
    for version in $sorted_versions; do
        local packages=(${packages_by_version["$version"]})

        if [ ${#packages[@]} -eq 0 ]; then
            continue
        fi

        echo ""
        if [[ "$version" -ge 590 ]]; then
            eval_gettext "Version \$version: ⚠️  NOT FOR MAXWELL/PASCAL!"; echo
            echo "--------------"
        else
            eval_gettext "Version \$version:"; echo
            echo "--------------"
        fi

        # Sort packages inside version: main first, then variants
        local sorted_packages=()
        for pkg in "${packages[@]}"; do
            if [[ "$pkg" == nvidia-driver-$version ]] || [[ "$pkg" == nvidia-driver-$version* ]]; then
                sorted_packages=("$pkg" "${sorted_packages[@]}")
            elif [[ "$pkg" == nvidia-$version ]] || [[ "$pkg" == nvidia-$version* ]]; then
                sorted_packages=("$pkg" "${sorted_packages[@]}")
            else
                sorted_packages+=("$pkg")
            fi
        done

        for pkg in "${sorted_packages[@]}"; do
            if [ -n "$pkg" ]; then
                selection_map[$i]="$pkg"

                local desc="${package_descriptions["$pkg"]}"
                local type="${package_types["$pkg"]}"

                printf "%3d. %-40s%s\n" "$i" "$pkg" "$type"

                if [ -n "$desc" ]; then
                    printf "     %s\n" "$desc"
                fi

                ((i++))
            fi
        done
    done

    local total_packages=$((i-1))

    if [ $total_packages -eq 0 ]; then
        echo "$(gettext "No packages available for selection.")"
        exit 1
    fi

    # Determine last version (newest)
    local last_version=$(echo "$sorted_versions" | tail -1)
    echo ""
    echo "======================================"
    echo -n "$(eval_gettext "Select package [1-\$total_packages] (default \$total_packages - version \$last_version): ")"
    read -r choice

    if [ -z "$choice" ]; then
        choice=$total_packages
    fi

    if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le $total_packages ]; then
        SELECTED_METAPACKAGE="${selection_map[$choice]}"
        eval_gettext "Selected package: \$SELECTED_METAPACKAGE"; echo
    else
        echo "$(gettext "Invalid choice. Using newest version by default.")"
        SELECTED_METAPACKAGE="${selection_map[$total_packages]}"
    fi

    # Extract version
    local selected_clean="${SELECTED_METAPACKAGE%%_*}"
    DRIVER_VERSION=""

    # Use the same regular expressions to extract version
    if [[ "$selected_clean" =~ nvidia-driver-no-dkms-([0-9]+) ]]; then
        DRIVER_VERSION="${BASH_REMATCH[1]}"
    elif [[ "$selected_clean" =~ nvidia-driver-([0-9]+) ]]; then
        DRIVER_VERSION="${BASH_REMATCH[1]}"
    elif [[ "$selected_clean" =~ nvidia-no-dkms-([0-9]+) ]]; then
        DRIVER_VERSION="${BASH_REMATCH[1]}"
    elif [[ "$selected_clean" =~ nvidia-([0-9]+) ]]; then
        DRIVER_VERSION="${BASH_REMATCH[1]}"
    elif [[ "$selected_clean" =~ nvidia-driver-([0-9]+)-open ]]; then
        DRIVER_VERSION="${BASH_REMATCH[1]}"
    fi

    if [ -n "$DRIVER_VERSION" ]; then
        eval_gettext "Driver version: \$DRIVER_VERSION"; echo
    else
        DRIVER_VERSION=""
    fi

    # Check compatibility for versions 590+
    if [[ -n "$DRIVER_VERSION" ]] && [[ "$DRIVER_VERSION" -ge 590 ]]; then
        check_driver_compatibility "$DRIVER_VERSION"
    fi

    # Check for additional packages
    echo ""
    if [ -n "$DRIVER_VERSION" ]; then
        eval_gettext "Searching for additional packages for version \$DRIVER_VERSION..."; echo
    else
        echo "$(gettext "Searching for additional packages...")"
    fi

    # nvidia-settings
    if apt-cache show nvidia-settings >/dev/null 2>&1; then
        NVIDIA_SETTINGS="nvidia-settings"
        echo "✓ $(gettext "Found nvidia-settings")"
    else
        NVIDIA_SETTINGS=""
        echo "✗ $(gettext "nvidia-settings not found")"
    fi

    # xserver-xorg-video-nvidia-* (only for new drivers)
    if [ -n "$DRIVER_VERSION" ] && [[ "$selected_clean" != nvidia-340 ]] && [[ "$selected_clean" != nvidia-no-dkms-340 ]]; then
        XORG_PACKAGE="xserver-xorg-video-nvidia-$DRIVER_VERSION"
        if apt-cache show "$XORG_PACKAGE" >/dev/null 2>&1; then
            eval_gettext "✓ Found \$XORG_PACKAGE"; echo
        else
            XORG_PACKAGE=""
            eval_gettext "✗ \$XORG_PACKAGE not found"; echo
        fi
    else
        XORG_PACKAGE=""
    fi

    # Installation
    echo ""
    echo "========================================"
    echo "$(gettext "Installing NVIDIA driver")"
    eval_gettext "Package: \$SELECTED_METAPACKAGE"; echo

    PACKAGES_TO_INSTALL="$SELECTED_METAPACKAGE"
    [ -n "$NVIDIA_SETTINGS" ] && PACKAGES_TO_INSTALL="$PACKAGES_TO_INSTALL $NVIDIA_SETTINGS"
    [ -n "$XORG_PACKAGE" ] && PACKAGES_TO_INSTALL="$PACKAGES_TO_INSTALL $XORG_PACKAGE"

    eval_gettext "Will be installed: \$PACKAGES_TO_INSTALL"; echo
    echo "========================================"

    read -p "$(gettext "Start installation? (Y/n): ")" confirm_install
    if [[ "$confirm_install" =~ ^[Nn]$ ]]; then
        echo "$(gettext "Installation canceled.")"
        exit 0
    fi

    apt install --yes $PACKAGES_TO_INSTALL

    # Post-installation configuration
    post_install_configuration
}

# Main interactive installation function
interactive_installation() {
    local timestamp="$(date +%s)"

    check_astra
    check_root

    # Check Extended repository
    if ! check_extended_repo_compatibility; then
        exit 1
    fi

    # 1. Check and configure repositories
    check_astra_repos

    # 2. Remove old xorg.conf
    gettext "Removing xorg.conf .."; echo
    if [ -f /etc/X11/xorg.conf ]; then
        mv -f /etc/X11/xorg.conf "/etc/X11/backup-$timestamp-xorg.conf-backup"
        eval_gettext '/etc/X11/xorg.conf removed, backup version: /etc/X11/backup-$timestamp-xorg.conf-backup' ; echo
    fi

    # 3. Reconfigure Xorg
    dpkg-reconfigure xserver-xorg

    # 4. Find and select metapackage
    find_all_nvidia_metapackages
}

# Command line argument parsing
# -----------------------------------
MODE="interactive"  # Default interactive mode
DRIVER_NAME=""
SCRIPT_NAME=$(basename "$0")

# Process arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            show_help
            ;;
        -l|--list)
            show_driver_list
            ;;
        -v|--version)
            show_version
            ;;
        -d|--driver)
            if [ -n "$2" ] && [[ ! "$2" =~ ^- ]]; then
                DRIVER_NAME="$2"
                MODE="specific"
                shift 2
                continue
            else
                echo "$(gettext "ERROR: Option -d requires driver name")"
                echo "$(gettext "Example:") $SCRIPT_NAME -d nvidia-driver-560"
                exit 1
            fi
            ;;
        -*)
            eval_gettext "Unknown option: \$1"; echo
            echo "$(gettext "Use") $SCRIPT_NAME -h $(gettext "for help")"
            exit 1
            ;;
        *)
            # If something passed without option, consider it driver name
            DRIVER_NAME="$1"
            MODE="specific"
            shift
            ;;
    esac
    shift
done

# Execute based on mode
case $MODE in
    "interactive")
        interactive_installation
        ;;
    "specific")
        check_astra
        check_root
        # Check Extended repository for -d mode
        if ! check_extended_repo_compatibility; then
            exit 1
        fi
        install_specific_driver "$DRIVER_NAME"
        ;;
esac
