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

LOG_FILE="${LOG_FILE:-$(pwd)/log.txt}"
ROCM_ROOT="/opt/rocm"
REQ_PKGS=(
  hipcc
  hip-dev
  rocm-llvm
  hipblas
  hipblas-dev
)

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
}

show_status() {
  section "Etap 1: diagnostyka pakietów HIP/LLVM"
  run "apt-cache policy hipcc hip-dev rocm-llvm hipblas hipblas-dev"
  run "dpkg -l | egrep 'hipcc|hip-dev|rocm-llvm|hipblas|hipblas-dev|hip-runtime-amd|rocm-hip-runtime|rocm-core|rocminfo' || true"
  run "ls -l $ROCM_ROOT/bin/hipcc $ROCM_ROOT/bin/hipconfig $ROCM_ROOT/llvm/bin/clang++ $ROCM_ROOT/bin/amdclang++ 2>/dev/null || true"
  run "ls -l $ROCM_ROOT/lib/cmake/hip-lang/hip-lang-config.cmake 2>/dev/null || true"

  local missing=0
  for p in "${REQ_PKGS[@]}"; do
    if ! dpkg -s "$p" >/dev/null 2>&1; then
      log "Brakuje pakietu: $p"
      missing=$((missing + 1))
    else
      log "Pakiet już jest: $p"
    fi
  done

  if [[ $missing -eq 0 ]]; then
    log "Wszystkie wymagane pakiety developerskie są już zainstalowane."
  else
    log "Liczba brakujących pakietów developerskich: $missing"
  fi
}

install_missing() {
  section "Etap 2: instalacja brakujących pakietów"
  if ! ask "Czy doinstalować brakujące pakiety HIP/LLVM dla ROCm 6.3.3?"; then
    log "Anulowano przez użytkownika."
    exit 0
  fi

  run "apt-get update"
  run "DEBIAN_FRONTEND=noninteractive apt-get install -y hipcc hip-dev rocm-llvm hipblas hipblas-dev"

  section "Etap 3: odświeżenie linkera i środowiska"
  run "ldconfig"

  if [[ ! -f /etc/profile.d/rocm.sh ]]; then
    cat >/etc/profile.d/rocm.sh <<'EOP'
export ROCM_PATH=/opt/rocm
export PATH=$ROCM_PATH/bin:$ROCM_PATH/llvm/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"
  else
    log "Plik /etc/profile.d/rocm.sh już istnieje, nie nadpisuję."
  fi
}

verify() {
  section "Etap 4: weryfikacja po instalacji"
  run "apt-cache policy hipcc hip-dev rocm-llvm hipblas hipblas-dev"
  run "dpkg -l | egrep 'hipcc|hip-dev|rocm-llvm|hipblas|hipblas-dev|hip-runtime-amd|rocm-hip-runtime|rocm-core|rocminfo' || true"
  run "which hipcc hipconfig clang++ amdclang++ || true"
  run "ls -l $ROCM_ROOT/bin/hipcc $ROCM_ROOT/bin/hipconfig $ROCM_ROOT/llvm/bin/clang++ $ROCM_ROOT/bin/amdclang++ 2>/dev/null || true"
  run "ls -l $ROCM_ROOT/lib/cmake/hip-lang/hip-lang-config.cmake 2>/dev/null || true"
  run "hipcc --version || true"
  run "hipconfig --full || true"
  run "$ROCM_ROOT/bin/rocminfo | head -n 80 || true"

  local ok=0
  [[ -x "$ROCM_ROOT/bin/hipcc" ]] && ok=$((ok+1))
  [[ -x "$ROCM_ROOT/bin/hipconfig" ]] && ok=$((ok+1))
  [[ -x "$ROCM_ROOT/llvm/bin/clang++" || -x "$ROCM_ROOT/bin/amdclang++" ]] && ok=$((ok+1))
  [[ -f "$ROCM_ROOT/lib/cmake/hip-lang/hip-lang-config.cmake" ]] && ok=$((ok+1))

  section "Etap 5: wynik"
  if [[ $ok -ge 4 ]]; then
    log "Naprawa zakończona powodzeniem."
    log "Masz już komplet narzędzi HIP do CMake."
    log "Następny krok: uruchomić skrypt przebudowy llama.cpp albo ręcznie skonfigurować build."
    log "Przykładowa konfiguracja:"
    log "cmake -S /home/adminek/llama.cpp -B /home/adminek/llama.cpp/build -DCMAKE_BUILD_TYPE=Release -DCMAKE_PREFIX_PATH=/opt/rocm -DGGML_HIP=ON -DAMDGPU_TARGETS=gfx906"
  else
    log "Naprawa jest niepełna. Załącz log.txt do dalszej analizy."
    exit 1
  fi
}

need_root
: > "$LOG_FILE"
log "Start skryptu naprawczego HIP/LLVM dla ROCm 6.3.3"
log "Pełny log będzie zapisany do: $LOG_FILE"

show_status
install_missing
verify

log "Koniec."