#!/usr/bin/env bash
set -u
set -o pipefail

LOG_FILE="${LOG_FILE:-$(pwd)/log2.txt}"
AMDGPU_DEB="/home/adminek/amdgpu-install_6.3.60303-1_all.deb"

exec 3>&1
log(){ echo "[$(date '+%F %T')] $*" | tee -a "$LOG_FILE" >&3; }
run(){ log "RUN: $*"; bash -lc "$*" >>"$LOG_FILE" 2>&1; return $?; }
ask(){ local prompt="$1" ans; read -r -p "$prompt [y/N]: " ans < /dev/tty || ans="n"; [[ "$ans" =~ ^[Yy]$ ]]; }
section(){ log ""; log "=== $* ==="; }

need_root() {
  if [[ ${EUID:-$(id -u)} -ne 0 ]]; then
    log "Uruchom jako root: sudo bash $0"
    exit 1
  fi
}

summary_check() {
  section "Etap 1: diagnostyka"
  log "Sprawdzam pakiety ROCm, ścieżki /opt/rocm*, repozytoria i amdgpu."

  run "uname -a"
  run "cat /etc/os-release"
  run "dpkg -l | egrep 'rocm|hip|rocblas|rocminfo|amdgpu' || true"
  run "find /opt -maxdepth 2 \( -type d -o -type l \) | egrep 'rocm|amdgpu' || true"
  run "ls -l /opt/rocm /opt/rocm-* 2>/dev/null || true"
  run "grep -RIn 'repo.radeon.com' /etc/apt/sources.list /etc/apt/sources.list.d 2>/dev/null || true"
  run "ls -l /etc/apt/preferences.d 2>/dev/null || true"
  run "which rocminfo hipconfig amdgpu-install || true"
  run "ls /opt/rocm/lib/rocblas/library/TensileLibrary*.dat 2>/dev/null || true"

  local installed633 installed_other has_rocm_path has_deb
  installed633=$(dpkg -l | awk '/^ii/ && $2 ~ /(rocm|hip|rocblas|rocminfo|amdgpu)/ && $3 ~ /6\.3\.3|60303/ {c++} END{print c+0}')
  installed_other=$(dpkg -l | awk '/^ii/ && $2 ~ /(rocm|hip|rocblas|rocminfo|amdgpu)/ && $3 ~ /6\.4\.|6\.5\.|6\.6\.|7\./ {c++} END{print c+0}')
  [[ -e /opt/rocm ]] && has_rocm_path=yes || has_rocm_path=no
  [[ -f "$AMDGPU_DEB" ]] && has_deb=yes || has_deb=no

  log "Wynik: pakiety 6.3.3/60303 znalezione: $installed633"
  log "Wynik: pakiety wersji 6.4+/7.x znalezione: $installed_other"
  log "Wynik: ścieżka /opt/rocm istnieje: $has_rocm_path"
  log "Wynik: lokalny instalator istnieje: $has_deb ($AMDGPU_DEB)"

  if [[ "$has_deb" != yes ]]; then
    log "Brak pliku instalatora .deb pod wskazaną ścieżką."
  fi
}

repair() {
  section "Etap 2: odblokowanie pakietów held"
  if ask "Czy odblokować wszystkie pakiety held i kontynuować?"; then
    run "apt-mark showhold | xargs -r apt-mark unhold"
  else
    log "Pominięto odblokowanie pakietów held."
  fi

  section "Etap 3: usunięcie ROCm 6.4.x i starych ścieżek"
  if ask "Czy usunąć wszystkie pakiety ROCm 6.4.x i ścieżki /opt/rocm*?"; then
    run "DEBIAN_FRONTEND=noninteractive apt-get remove --purge -y 'rocm-core' 'rocm-llvm' 'rocm-device-libs' 'rocm-smi-lib' || true"
    run "DEBIAN_FRONTEND=noninteractive apt-get autoremove -y --purge || true"
    run "DEBIAN_FRONTEND=noninteractive apt-get clean"
    run "rm -rf /opt/rocm /opt/rocm-*"
    run "rm -f /etc/apt/sources.list.d/*rocm*.list /etc/apt/preferences.d/*rocm*"
  else
    log "Pominięto czyszczenie ROCm 6.4.x; instalacja może być niekompletna."
  fi

  section "Etap 4: dodanie repozytorium ROCm 6.3.3"
  log "Dodam klucz i repozytorium ROCm 6.3.3 dla Ubuntu 22.04 (jammy)."

  run "mkdir --parents --mode=0755 /etc/apt/keyrings"
  run "wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | tee /etc/apt/keyrings/rocm.gpg > /dev/null"

  run "echo 'deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.3.3 jammy main' | tee /etc/apt/sources.list.d/rocm.list"
  run "echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | tee /etc/apt/preferences.d/rocm-pin-600"

  run "apt-get update"

  section "Etap 5: instalacja ROCm 6.3.3 z repo"
  run "DEBIAN_FRONTEND=noninteractive apt-get install -y rocm-core=6.3.3.60303-* rocblas=6.3.3.60303-* hip-runtime-amd=6.3.42134.60303-* rocminfo=1.0.0.60303-* rocm-hip-runtime=6.3.3.60303-* rocm-opencl-runtime=2.0.0.60303-* rocm-device-libs=1.0.0.60303-* || true"

  section "Etap 6: konfiguracja środowiska"
  run "cat >/etc/profile.d/rocm.sh <<'EOP'
export ROCM_PATH=/opt/rocm
export PATH=\$ROCM_PATH/bin:\$PATH
export LD_LIBRARY_PATH=\$ROCM_PATH/lib:\$ROCM_PATH/lib64:\${LD_LIBRARY_PATH:-}
EOP"
  run "chmod 644 /etc/profile.d/rocm.sh"
  log "Dodano /etc/profile.d/rocm.sh"

  section "Etap 7: weryfikacja końcowa"
  run "apt-cache policy rocm-core rocblas rocminfo hip-runtime-amd rocm-hip-runtime rocm-opencl-runtime"
  run "dpkg -l | egrep 'rocm|hip|rocblas|rocminfo' || true"
  run "ls -l /opt/rocm /opt/rocm-* 2>/dev/null || true"
  run "ls /opt/rocm/lib/rocblas/library/TensileLibrary*.dat 2>/dev/null || true"
  run "which rocminfo hipconfig || true"
  run "rocminfo || true"
  run "hipconfig -l || true"
  run "lsmod | grep amdgpu || true"

  local ok=0
  command -v rocminfo >/dev/null 2>&1 && ok=$((ok+1))
  [[ -d /opt/rocm/lib/rocblas/library ]] && ok=$((ok+1))
  if ls /opt/rocm/lib/rocblas/library/TensileLibrary*gfx906*.dat >/dev/null 2>&1; then ok=$((ok+1)); fi

  if [[ $ok -ge 2 ]]; then
    log "Wynik końcowy: instalacja wygląda sensownie. Jeśli DKMS/kernel był zmieniany, zrób reboot przed testem llama.cpp."
  else
    log "Wynik końcowy: instalacja nadal wygląda niepełnie. Załącz log2.txt do dalszej analizy."
  fi
}

need_root
: > "$LOG_FILE"
log "Start skryptu naprawczego ROCm 6.3.3 v2"
log "Pełny log poleceń będzie zapisany do: $LOG_FILE"
summary_check
repair
log "Koniec."