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

LOG_FILE="${LOG_FILE:-$(pwd)/log.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 system, pakiety, repozytoria, pliki ROCm i sterownik AMDGPU."

  run "uname -a"
  run "cat /etc/os-release"
  run "apt-mark showhold"
  run "dpkg --print-architecture"
  run "dpkg --print-foreign-architectures"
  run \"apt-cache policy rocm-core rocblas rocminfo hip-runtime-amd rocm-hip-runtime rocm-opencl-runtime amdgpu-dkms amdgpu-install\"
  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 "modinfo amdgpu | head -n 30 || true"
  run "lsmod | grep amdgpu || true"
  run "lspci -nnk | grep -A3 -E 'VGA|Display|3D' || true"
  run "dkms status || 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\.[0-2]\.|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 innych wersji ROCm/AMDGPU 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ą. Przerwano."
    exit 1
  fi
}

repair() {
  section "Etap 2: przygotowanie naprawy"

  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

  if ask "Czy usunąć wszystkie pakiety ROCm/hip/rocblas/rocminfo/amdgpu, aby zrobić czystą instalację 6.3.3?"; then
    run "DEBIAN_FRONTEND=noninteractive apt-get remove --purge -y 'rocm*' 'hip*' 'rocblas*' 'rocminfo*' 'rocsolver*' 'rocrand*' 'rocfft*' 'rocprim*' 'rccl*' 'miopen*' 'comgr*' 'hsa-*' 'hsakmt-*' 'amdgpu*' || true"
    run "DEBIAN_FRONTEND=noninteractive apt-get autoremove -y --purge || true"
    run "DEBIAN_FRONTEND=noninteractive apt-get clean"
  else
    log "Wybrano naprawę bez pełnego purge pakietów."
  fi

  section "Etap 3: czyszczenie pozostałości"
  if ask "Czy usunąć pozostałe katalogi /opt/rocm i stare listy repozytoriów AMD/ROCm?"; then
    run "rm -rf /opt/rocm /opt/rocm-*"
    run "rm -f /etc/apt/sources.list.d/*rocm*.list /etc/apt/sources.list.d/*amdgpu*.list"
    run "rm -f /etc/apt/preferences.d/*rocm* /etc/apt/preferences.d/*amdgpu* /etc/apt/preferences.d/*radeon*"
  else
    log "Pominięto usuwanie pozostałości."
  fi

  section "Etap 4: przygotowanie APT"
  if ask "Czy włączyć architekturę i386? Jest potrzebna tylko do pakietów 32-bit; poprzedni błąd wskazywał na brak i386."; then
    run "dpkg --add-architecture i386"
  else
    log "Pominięto dodanie i386; instalacja pójdzie bez komponentów 32-bit."
  fi

  run "apt-get update"
  run "DEBIAN_FRONTEND=noninteractive apt-get install -y dkms linux-headers-\$(uname -r) curl gnupg2 ca-certificates"

  section "Etap 5: instalacja amdgpu-install 6.3.3"
  run "DEBIAN_FRONTEND=noninteractive apt-get install -y '$AMDGPU_DEB'"
  run "which amdgpu-install"
  run "amdgpu-install --help | head -n 80"

  section "Etap 6: instalacja sterownika i ROCm 6.3.3"
  log "Najpierw próbuję instalacji minimalnej bez lib32, żeby ominąć konflikt i386."
  if run "DEBIAN_FRONTEND=noninteractive amdgpu-install --usecase=dkms,rocm --no-32 --accept-eula -y"; then
    log "Instalacja minimalna zakończona powodzeniem."
  else
    log "Instalacja minimalna nie powiodła się. Próbuję wariantu bez dkms."
    run "DEBIAN_FRONTEND=noninteractive amdgpu-install --usecase=rocm --no-dkms --no-32 --accept-eula -y"
  fi

  section "Etap 7: doinstalowanie brakujących pakietów ROCm"
  run "apt-get update"
  run "DEBIAN_FRONTEND=noninteractive apt-get install -y rocm-core rocm-hip-runtime rocm-opencl-runtime rocblas hip-runtime-amd rocminfo || true"

  section "Etap 8: konfiguracja środowiska"
  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
  chmod 644 /etc/profile.d/rocm.sh
  log "Dodano /etc/profile.d/rocm.sh"

  section "Etap 9: weryfikacja końcowa"
  run "apt-cache policy rocm-core rocblas rocminfo hip-runtime-amd rocm-hip-runtime rocm-opencl-runtime amdgpu-dkms"
  run "dpkg -l | egrep 'rocm|hip|rocblas|rocminfo|amdgpu' || 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 amdgpu-install || true"
  run "rocminfo || true"
  run "hipconfig -l || true"
  run "lsmod | grep amdgpu || true"
  run "dkms status || 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 log.txt do dalszej analizy."
  fi
}

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