#!/usr/bin/env bash
# Fedora ↔ PyTorch local checkup with Conda, GCC & Clang builds, and runtime probes.
# Outputs a complete report folder to ~/pytorch-fedora-report

set -euo pipefail

# -------- Config (override via env vars before running) ----------
#: "${PYTORCH_REPO_DIR:=${HOME}/pytorch}"
: "${PYTORCH_REPO_DIR:=$(pwd)/pytorch}"
: "${CONDA_ENV_NAME:=torch-fedora}"
: "${REPORT_DIR:=$(pwd)/pytorch-fedora-report}"
: "${PY_VER:=3.12}"
# -----------------------------------------------------------------

mkdir -p "${REPORT_DIR}"

log() { echo -e "\033[1;36m[INFO]\033[0m $*"; }
warn() { echo -e "\033[1;33m[WARN]\033[0m $*" >&2; }
err() { echo -e "\033[1;31m[ERR]\033[0m  $*" >&2; }

# 0) Preflight
log "Fedora: $(cat /etc/fedora-release 2>/dev/null || echo 'unknown')" | tee "${REPORT_DIR}/preflight.txt"
log "Kernel: $(uname -r)" | tee -a "${REPORT_DIR}/preflight.txt"

# 1) Ensure conda is available
if ! command -v conda >/dev/null 2>&1; then
  err "Conda not found. Install Miniconda first and re-run."
  exit 1
fi
# shellcheck disable=SC1091
source "$(conda info --base)/etc/profile.d/conda.sh"

# 2) Create/activate env
if ! conda env list | awk '{print $1}' | grep -qx "${CONDA_ENV_NAME}"; then
  log "Creating conda env ${CONDA_ENV_NAME} (python=${PY_VER}) ..."
  conda create -n "${CONDA_ENV_NAME}" "python=${PY_VER}" -y
fi
conda activate "${CONDA_ENV_NAME}"

python -V | tee -a "${REPORT_DIR}/preflight.txt"

# 3) System deps (sudo dnf/dnf5)
DNF_BIN="$(command -v dnf5 || command -v dnf)"

echo "[INFO] Using package manager: $DNF_BIN"

# Install the dev tools group (dnf5 uses 'group install', classic dnf uses 'groupinstall')
if "$DNF_BIN" --version 2>/dev/null | grep -qi '^DNF5'; then
  sudo "$DNF_BIN" group install -y "Development Tools" | tee "${REPORT_DIR}/dnf_groupinstall.log" || true
  # Fallback: some setups prefer @development-tools
  sudo "$DNF_BIN" install -y @development-tools | tee -a "${REPORT_DIR}/dnf_groupinstall.log" || true
else
  sudo "$DNF_BIN" -y groupinstall "Development Tools" | tee "${REPORT_DIR}/dnf_groupinstall.log" || true
fi

sudo "$DNF_BIN" install -y \
  git cmake ninja-build pkgconf-pkg-config \
  gcc gcc-c++ clang libomp \
  openblas-devel libjpeg-turbo-devel libpng-devel zlib-devel \
  firewalld \
  | tee "${REPORT_DIR}/dnf_install.log" || true


# 4) Clone PyTorch (or update)
log "PYTORCH_REPO_DIR is set to: ${PYTORCH_REPO_DIR}"

# Check if target directory exists and is not empty (but not a git repo)
if [ -d "${PYTORCH_REPO_DIR}" ] && [ ! -d "${PYTORCH_REPO_DIR}/.git" ]; then
  if [ "$(ls -A "${PYTORCH_REPO_DIR}" 2>/dev/null)" ]; then
    err "Directory ${PYTORCH_REPO_DIR} already exists and is not empty, but is not a git repository"
    err "Please remove the directory or choose a different PYTORCH_REPO_DIR"
    exit 1
  fi
fi

if [ ! -d "${PYTORCH_REPO_DIR}/.git" ]; then
  log "Cloning PyTorch into ${PYTORCH_REPO_DIR} ..."
  git clone --recursive https://github.com/pytorch/pytorch.git "${PYTORCH_REPO_DIR}"
else
  log "PyTorch repository already exists at ${PYTORCH_REPO_DIR}, skipping clone..."
fi

# Check if we have a valid PyTorch repository
if [ ! -f "${PYTORCH_REPO_DIR}/setup.py" ] || [ ! -d "${PYTORCH_REPO_DIR}/torch" ]; then
  err "Directory ${PYTORCH_REPO_DIR} does not appear to be a valid PyTorch repository"
  exit 1
fi

cd "${PYTORCH_REPO_DIR}"

# Only update if it's a git repository and we want to fetch updates
if [ -d ".git" ]; then
  log "Updating PyTorch repository..."
  git fetch origin
  git submodule sync --recursive
  git submodule update --init --recursive
else
  log "Not a git repository, skipping git operations..."
fi

# 5) Common build env
export MAX_JOBS="$(nproc)"
export CMAKE_PREFIX_PATH="${CONDA_PREFIX}"
# Make sure pip tools are fresh in the env
#python -m pip install -U pip setuptools wheel lintrunner usort black ruff pytest >/dev/null
#changes
# --- Python deps for building/using PyTorch from source (with verification) ---
# Note: 'six' is required for NNPACK/PeachPy assembly generation
echo "[INFO] Installing/upgrading Python build deps and PyTorch prerequisites..."
python -m pip install -U \
  pip setuptools wheel \
  typing_extensions pyyaml sympy filelock networkx jinja2 fsspec \
  six \
  cmake ninja \
  lintrunner usort black ruff pytest

# (Optional) Retry once if the first attempt hit a transient resolver/cache issue
if [ $? -ne 0 ]; then
  echo "[WARN] pip install failed once; retrying with --no-cache-dir ..."
  python -m pip install -U --no-cache-dir \
    pip setuptools wheel \
    typing_extensions pyyaml sympy filelock networkx jinja2 fsspec \
    six \
    cmake ninja \
    lintrunner usort black ruff pytest
fi

echo "[INFO] Verifying Python deps are importable..."
python - <<PY
import sys
mods = [
    "typing_extensions", "yaml", "sympy", "filelock", "networkx",
    "jinja2", "fsspec", "six"
]
missing = []
for m in mods:
    try:
        __import__(m)
    except Exception as e:
        missing.append(f"{m} ({e})")
if missing:
    sys.exit("Missing/failed imports: " + ", ".join(missing))
print("All required modules import OK.")
PY



# 6) Build with GCC (CPU)
log "Building (GCC, CPU-only) ..."
( USE_CUDA=0 USE_ROCM=0 python setup.py develop \
  2>&1 | tee "${REPORT_DIR}/build_gcc.log" ) || true

log "Extracting top GCC warnings ..."
grep -i 'warning:' "${REPORT_DIR}/build_gcc.log" | sed 's/^[^:]*://g' \
  | sort | uniq -c | sort -nr > "${REPORT_DIR}/warnings_gcc.txt" || true

# 7) Build with Clang + OpenMP (CPU)
log "Building (Clang+OpenMP, CPU-only) ..."
( CC=clang CXX=clang++ USE_OPENMP=1 USE_CUDA=0 USE_ROCM=0 \
  python setup.py develop \
  2>&1 | tee "${REPORT_DIR}/build_clang.log" ) || true

log "Extracting top Clang warnings ..."
grep -i 'warning:' "${REPORT_DIR}/build_clang.log" | sed 's/^[^:]*://g' \
  | sort | uniq -c | sort -nr > "${REPORT_DIR}/warnings_clang.txt" || true

# 8) Smoke test (import + matmul + OpenMP)
log "Running smoke test ..."
python - <<'PY' | tee "${REPORT_DIR}/smoke.txt"
import torch, os, platform
print("Torch:", torch.__version__)
print("Platform:", platform.platform())
print("OpenMP available:", torch.backends.openmp.is_available())
x = torch.randn(256, 256)
y = x @ x.t()
print("Matmul OK:", bool((y.abs().sum().item() > 0)))
PY

# 9) /dev/shm & DataLoader stress
log "Checking /dev/shm & DataLoader ..."
df -h /dev/shm | tee "${REPORT_DIR}/devshm.txt" || true
python - <<'PY' 2>&1 | tee "${REPORT_DIR}/dataloader_stress.txt" || true
import os, torch
from torch.utils.data import DataLoader, TensorDataset
import warnings, sys

# Report /dev/shm size precisely
try:
    st = os.statvfs("/dev/shm")
    shm_mb = (st.f_frsize * st.f_blocks) // (1024*1024)
    print(f"/dev/shm total ~{shm_mb} MB")
except Exception as e:
    print("statvfs failed:", e)

# Stress run
ds = TensorDataset(torch.arange(300_000))
dl = DataLoader(ds, batch_size=2048, num_workers=4)
c = 0
for _ in dl:
    c += 1
    if c % 50 == 0:
        print("Batches processed:", c)
print("DataLoader stress: done, total batches:", c)
PY

# 10) torch.distributed TCPStore under firewalld
log "Testing torch.distributed TCPStore (firewalld) ..."
# ensure firewalld is active (ignore error if not supported)
sudo systemctl enable --now firewalld >/dev/null 2>&1 || true
python - <<'PY' 2>&1 | tee "${REPORT_DIR}/distributed_tcpstore.txt" || true
import torch.distributed as dist, time
from contextlib import suppress
try:
    dist.init_process_group(
        backend="gloo",
        init_method="tcp://127.0.0.1:29500",
        rank=0, world_size=1, timeout=None
    )
    print("TCPStore OK (init_process_group)")
    dist.destroy_process_group()
except Exception as e:
    print("TCPStore failed:", e)
PY

# 11) Optional CUDA check (only if NVIDIA stack present)
log "CUDA availability probe (optional) ..."
( ldconfig -p | grep -E 'cuda|nvidia|cudart' || true ) | tee "${REPORT_DIR}/ldconfig_cuda.txt"
python - <<'PY' 2>&1 | tee "${REPORT_DIR}/cuda_probe.txt" || true
import shutil, torch
print("nvidia-smi:", shutil.which("nvidia-smi"))
print("torch.version.cuda:", torch.version.cuda)
print("cuda.is_available():", torch.cuda.is_available())
PY

# 12) Generate SUMMARY with suggested PRs
log "Creating SUMMARY.txt ..."
{
  echo "=== Fedora ↔ PyTorch Local Check Summary ==="
  echo "Repo: ${PYTORCH_REPO_DIR}"
  echo "Conda env: ${CONDA_ENV_NAME}"
  echo "Report dir: ${REPORT_DIR}"
  echo

  echo "---- Toolchain ----"
  echo "Python: $(python -V 2>&1)"
  echo "GCC: $(gcc --version 2>/dev/null | head -n1)"
  echo "Clang: $(clang --version 2>/dev/null | head -n1)"
  echo

  echo "---- Build warnings (top lines) ----"
  echo "GCC:"
  head -n 15 "${REPORT_DIR}/warnings_gcc.txt" 2>/dev/null || echo "  (no warnings captured)"
  echo
  echo "Clang:"
  head -n 15 "${REPORT_DIR}/warnings_clang.txt" 2>/dev/null || echo "  (no warnings captured)"
  echo

  echo "---- Smoke test ----"
  cat "${REPORT_DIR}/smoke.txt" 2>/dev/null || true
  echo

  echo "---- /dev/shm & DataLoader ----"
  cat "${REPORT_DIR}/devshm.txt" 2>/dev/null || true
  echo
  awk 'NR<=80{print} NR==81{print "... (truncated)"}' "${REPORT_DIR}/dataloader_stress.txt" 2>/dev/null || true
  echo

  echo "---- Distributed TCPStore ----"
  cat "${REPORT_DIR}/distributed_tcpstore.txt" 2>/dev/null || true
  echo

  echo "---- CUDA probe ----"
  awk 'NR<=40{print} NR==41{print "... (truncated)"}' "${REPORT_DIR}/cuda_probe.txt" 2>/dev/null || true
  echo

  echo "==== Suggested PRs (based on findings) ===="
  # Heuristics:
  SHM_MB=$(awk '/\/dev\/shm/ && $2 ~ /[0-9]/ {print $2}' "${REPORT_DIR}/devshm.txt" 2>/dev/null | tail -n1)
  # Not robust to all locales; rely on dataloader_stress summary instead:
  if grep -qi "Bus error" "${REPORT_DIR}/dataloader_stress.txt" 2>/dev/null; then
    echo "- DataLoader: Add proactive warning when /dev/shm is small (observed Bus error)."
  elif awk '/total ~[0-9]+ MB/{if($3<256) f=1} END{exit !f}' "${REPORT_DIR}/dataloader_stress.txt" 2>/dev/null; then
    echo "- DataLoader: Add proactive warning when /dev/shm < 256 MB (observed small shared memory)."
  fi

  if grep -qi "libomp.so.*cannot open" "${REPORT_DIR}/build_clang.log" 2>/dev/null; then
    echo "- CMake: Add Fedora-friendly rpath or link hint for Clang OpenMP (libomp) on torch_cpu."
  fi

  if grep -qi "ModuleNotFoundError.*six" "${REPORT_DIR}/build_gcc.log" 2>/dev/null; then
    echo "- Dependencies: Add 'six' module to PyTorch build requirements (required for NNPACK/PeachPy)."
  fi

  if grep -qi "TCPStore failed" "${REPORT_DIR}/distributed_tcpstore.txt" 2>/dev/null; then
    echo "- Distributed: Improve TCPStore diagnostics & firewall hint (firewalld) and small retry on bind."
  fi

  if [ -s "${REPORT_DIR}/warnings_gcc.txt" ] || [ -s "${REPORT_DIR}/warnings_clang.txt" ]; then
    echo "- Warnings: Submit focused cleanups for top GCC/Clang warnings (correctness-preserving casts, sign-compare)."
  fi

  if grep -qi "cuda.is_available(): False" "${REPORT_DIR}/cuda_probe.txt" 2>/dev/null && \
     grep -Eq '/usr/lib64|/usr/lib64/nvidia' "${REPORT_DIR}/ldconfig_cuda.txt" 2>/dev/null; then
    echo "- CUDA: Improve Fedora diagnostic when CUDA libs exist under /usr/lib64(/nvidia) but torch.cuda.is_available() is False."
  fi
} > "${REPORT_DIR}/SUMMARY.txt"

log "Done. Report ready at: ${REPORT_DIR}"
echo
echo "Open ${REPORT_DIR}/SUMMARY.txt for a one-page overview."
