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

LOG_FILE="${LOG_FILE:-$(pwd)/log.txt}"
ROCM_ROOT="/opt/rocm"

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
}

diagnose() {
  section "Etap 1: diagnostyka hostowego C++"
  run "apt-cache policy g++ g++-11 gcc-11 libstdc++-11-dev build-essential"
  run "dpkg -l | egrep 'g\\+\\+|gcc-11|libstdc\\+\\+-11-dev|build-essential' || true"
  run "which c++ g++ gcc cc || true"
  run "c++ --version || true"
  run "g++ --version || true"
  run "gcc --version || true"
  run "ls -l /usr/include/c++ /usr/include/c++/11 /usr/include/x86_64-linux-gnu/c++/11 2>/dev/null || true"
  run "find /usr/include -name cmath 2>/dev/null | sort | head -n 20 || true"
  run "'$ROCM_ROOT/bin/amdclang++' --version || true"
  run "'$ROCM_ROOT/bin/amdclang++' -v -E -x c++ /dev/null || true"
  run "echo '#include <cmath>' >/tmp/test_cmath.cpp"
  run "'$ROCM_ROOT/bin/amdclang++' --gcc-toolchain=/usr -std=c++17 -E /tmp/test_cmath.cpp >/dev/null || true"
}

repair() {
  section "Etap 2: naprawa pakietów C++"
  if ! ask "Czy doinstalować brakujące pakiety hostowego C++ dla HIP?"; then
    log "Anulowano przez użytkownika."
    exit 0
  fi

  run "apt-get update"
  run "DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential g++ gcc-11 g++-11 libstdc++-11-dev"
}

write_env() {
  section "Etap 3: ustawienie środowiska HIP host compiler"
  cat >/etc/profile.d/rocm-hostcpp.sh <<'EOP'
export ROCM_PATH=/opt/rocm
export HIP_PATH=/opt/rocm
export HIP_ROOT_DIR=/opt/rocm
export HIPCXX=/opt/rocm/bin/amdclang++
export HIP_CLANG_PATH=/opt/rocm/lib/llvm/bin
export CXX=/usr/bin/g++
export CC=/usr/bin/gcc
export CPLUS_INCLUDE_PATH=/usr/include/c++/11:/usr/include/x86_64-linux-gnu/c++/11
EOP
  chmod 644 /etc/profile.d/rocm-hostcpp.sh
  log "Dodano /etc/profile.d/rocm-hostcpp.sh"
}

verify() {
  section "Etap 4: weryfikacja"
  run "source /etc/profile.d/rocm-hostcpp.sh && env | egrep 'ROCM|HIP|CC=|CXX=|CPLUS_INCLUDE_PATH' | sort"
  run "source /etc/profile.d/rocm-hostcpp.sh && '$ROCM_ROOT/bin/amdclang++' --gcc-toolchain=/usr -std=c++17 -E /tmp/test_cmath.cpp >/dev/null"
  run "source /etc/profile.d/rocm-hostcpp.sh && hipcc --version || true"

  if bash -lc "source /etc/profile.d/rocm-hostcpp.sh && '$ROCM_ROOT/bin/amdclang++' --gcc-toolchain=/usr -std=c++17 -E /tmp/test_cmath.cpp >/dev/null" >>"$LOG_FILE" 2>&1; then
    log "Naprawa zakończona powodzeniem. Hostowy C++ jest widoczny dla HIP."
    log "Następny krok: ponownie uruchomić build_llama_hip.sh."
  else
    log "Naprawa niepełna. Załącz log.txt."
    exit 1
  fi
}

need_root
: > "$LOG_FILE"
log "Start naprawy hostowego toolchaina C++ dla HIP"
log "Log: $LOG_FILE"

diagnose
repair
write_env
verify

log "Koniec."