#!/usr/bin/env bash
#
# install_mpi.sh - Interactive installation script for MPI.
#
# This script checks if MPI is installed (by trying to compile a simple MPI program).
# If not found, it prompts the user to install MPI using the native package manager for:
#   - FreeBSD, Ubuntu, Debian, Fedora, Arch, Mint, Manjaro, macOS, OpenSuse, and NixOS.
#
# All output is colorized for clarity and logged to mpi-install-log.txt.

set -e

# ANSI color codes.
RED="\033[0;31m"
GREEN="\033[0;32m"
YELLOW="\033[0;33m"
BLUE="\033[0;34m"
NC="\033[0m"  # No Color

# Log file.
LOGFILE="mpi-install-log.txt"
touch "$LOGFILE" # Ensure log file exists

# Log function: prints to stdout and appends to logfile.
log() {
  local message="$1"
  # Print the colored message to stdout.
  echo -e "$message"
  # Strip ANSI escape sequences and append the cleaned message to the log file.
  echo -e "$message" | sed -r 's/\x1B\[[0-9;]*[mK]//g' >> "$LOGFILE"
}

# Function to check if MPI (specifically mpicc and a compilable environment) is installed.
check_mpi_installed() {
  log "${BLUE}[Info] Checking for MPI (e.g., Open MPI, MPICH) using mpicc...${NC}"
  local tmpDir="mpiTest" # Define tmpDir locally
  
  # Clean up any previous test directory
  if [[ -d "$tmpDir" ]]; then
    rm -rf "$tmpDir"
  fi
  mkdir "$tmpDir"
  local sourceFile="$tmpDir/test_mpi.c"
  local outputFile="$tmpDir/test_mpi_app"

  # Write a minimal MPI C program
  cat > "$sourceFile" <<EOF
#include <mpi.h>
#include <stdio.h>

int main(int argc, char** argv) {
  if (MPI_Init(NULL, NULL) != MPI_SUCCESS) {
    // Using fprintf for stderr in case stdout is redirected by build tools
    fprintf(stderr, "MPI_Init failed\n");
    return 1;
  }
  int world_size;
  if (MPI_Comm_size(MPI_COMM_WORLD, &world_size) != MPI_SUCCESS) {
    fprintf(stderr, "MPI_Comm_size failed\n");
    MPI_Finalize(); // Attempt to finalize even if a step fails
    return 1;
  }
  int world_rank;
  if (MPI_Comm_rank(MPI_COMM_WORLD, &world_rank) != MPI_SUCCESS) {
    fprintf(stderr, "MPI_Comm_rank failed\n");
    MPI_Finalize();
    return 1;
  }
  // A simple printf to ensure it links and runs, though not strictly necessary for the check
  // printf("MPI Initialized successfully on rank %d of %d\\n", world_rank, world_size);
  if (MPI_Finalize() != MPI_SUCCESS) {
    fprintf(stderr, "MPI_Finalize failed\n");
    return 1;
  }
  return 0;
}
EOF

  # Try to find mpicc and compile the test program
  if command -v mpicc &> /dev/null; then
    # Log the compilation command
    log "${BLUE}[Info] Attempting to compile MPI test program with: mpicc \"$sourceFile\" -o \"$outputFile\"${NC}"
    # Attempt compilation, redirecting stdout and stderr of mpicc to the log file for details
    if mpicc "$sourceFile" -o "$outputFile" >> "$LOGFILE" 2>&1; then
      log "${GREEN}[Success] MPI compiler (mpicc) found and test program compiled successfully.${NC}"
      # Optional: Could add a mpiexec -n 1 "$outputFile" check here for further validation
      rm -rf "$tmpDir"
      return 0 # MPI seems to be installed and working
    else
      log "${RED}[Error] MPI compiler (mpicc) found, but failed to compile the test program. See $LOGFILE for details.${NC}"
    fi
  else
    log "${RED}[Error] MPI compiler (mpicc) not found in PATH.${NC}"
  fi
  
  # Cleanup and return failure
  rm -rf "$tmpDir"
  return 1
}

# Prompt the user for a yes/no answer.
prompt_yes_no() {
  local promptMsg="$1"
  while true; do
    read -r -p "$(echo -e "${YELLOW}$promptMsg${NC} (y/n): ")" answer
    case "$answer" in
      [Yy]* ) return 0;;
      [Nn]* ) return 1;;
      * ) echo -e "${RED}Please answer yes (y) or no (n).${NC}";;
    esac
  done
}

# Detect OS and Distribution.
OS=$(uname -s)
DISTRO="unknown"
if [ -f /etc/os-release ]; then
  # shellcheck disable=SC1091
  . /etc/os-release
  DISTRO=$ID
elif [ -f /etc/lsb-release ]; then # Fallback for older systems
  # shellcheck disable=SC1091
  . /etc/lsb-release
  DISTRO=$DISTRIB_ID
fi
# Normalize some distro names
DISTRO=$(echo "$DISTRO" | tr '[:upper:]' '[:lower:]')


if [[ "$OS" == "Darwin" ]]; then
  OS="macOS"
fi

if [[ $OS == "macOS" ]]; then
    log "${BLUE}[Info] Detected OS: ${OS}${NC}"
else
    log "${BLUE}[Info] Detected OS: ${OS}  Distribution: ${DISTRO}${NC}"
fi

# Main script logic
log "${BLUE}[Info] Starting MPI installation check...${NC}"

if check_mpi_installed; then
  log "${GREEN}[Success] MPI appears to be already installed and configured on your system.${NC}"
  log "${GREEN}[Success] Your MPI-dependent application can now proceed.${NC}"
  exit 0
else
  log "${YELLOW}[Warning] MPI was not detected or is not properly configured on your system.${NC}"
  if prompt_yes_no "[Query] Would you like this script to attempt to install MPI (Open MPI) now?"; then
    log "${BLUE}[Info] Attempting to install MPI...${NC}"
    case "$OS" in
      "macOS")
        if ! command -v brew &> /dev/null; then
          log "${YELLOW}[Warning] Homebrew is not installed. It is highly recommended for installing MPI on macOS.${NC}"
          if prompt_yes_no "[Query] Would you like to install Homebrew now?"; then
            log "${BLUE}[Info] Installing Homebrew...${NC}"
            # Piping to tee to ensure output is captured in the log file
            /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)" 2>&1 | tee -a "$LOGFILE"
          else
            log "${RED}[Error] Homebrew is required for an easy MPI installation on macOS using this script. Aborting.${NC}"
            exit 1
          fi
        fi
        log "${BLUE}[Info] Running: brew install open-mpi${NC}"
        brew install open-mpi 2>&1 | tee -a "$LOGFILE"
        ;;
      "Linux")
        case "$DISTRO" in
          "ubuntu"|"debian"|"linuxmint"|"pop"|"raspbian") # Added pop and raspbian
            log "${BLUE}[Info] Running: sudo apt-get update && sudo apt-get install -y libopenmpi-dev${NC}"
            sudo apt-get update 2>&1 | tee -a "$LOGFILE"
            sudo apt-get install -y libopenmpi-dev 2>&1 | tee -a "$LOGFILE"
            ;;
          "fedora"|"centos"|"rhel") # Added centos and rhel (dnf is common now)
             if command -v dnf &> /dev/null; then
                log "${BLUE}[Info] Running: sudo dnf install -y openmpi-devel${NC}"
                sudo dnf install -y openmpi-devel 2>&1 | tee -a "$LOGFILE"
             elif command -v yum &> /dev/null; then # Fallback for older CentOS/RHEL
                log "${BLUE}[Info] Running: sudo yum install -y openmpi-devel${NC}"
                sudo yum install -y openmpi-devel 2>&1 | tee -a "$LOGFILE"
             else
                log "${RED}[Error] Neither dnf nor yum found. Cannot install MPI on this Fedora/RHEL-like system automatically.${NC}"
                exit 1
             fi
            ;;
          "arch"|"manjaro"|"endeavouros") # Added endeavouros
            log "${BLUE}[Info] Running: sudo pacman -S --noconfirm openmpi${NC}"
            sudo pacman -S --noconfirm openmpi 2>&1 | tee -a "$LOGFILE"
            ;;
          "opensuse"|"opensuse-leap"|"opensuse-tumbleweed") # More specific opensuse IDs
            log "${BLUE}[Info] Running: sudo zypper install -y openmpi-devel${NC}" # Or openmpi4-devel, openmpi-devel is often a meta package
            sudo zypper install -y openmpi-devel 2>&1 | tee -a "$LOGFILE"
            ;;
          "nixos")
            log "${BLUE}[Info] On NixOS, you typically manage packages declaratively or with nix-shell.${NC}"
            log "${BLUE}[Info] To install Open MPI imperatively: nix-env -iA nixpkgs.openmpi${NC}"
            log "${YELLOW}[Warning] This script will not run nix-env automatically. Please install manually if desired.${NC}"
            log "${YELLOW}[Info] For declarative environment, add 'openmpi' to your configuration.nix or shell.nix.${NC}"
            # nix-env -iA nixpkgs.openmpi 2>&1 | tee -a "$LOGFILE" # Decided against auto-running this for NixOS
            ;;
          *)
            log "${RED}[Error] Your Linux distribution ($DISTRO) is not recognized or not explicitly supported for automatic MPI installation by this script.${NC}"
            log "${RED}[Info] Please try installing Open MPI or another MPI implementation manually.${NC}"
            exit 1
            ;;
        esac
        ;;
      "FreeBSD")
        log "${BLUE}[Info] Running: sudo pkg install openmpi${NC}"
        sudo pkg install -y openmpi 2>&1 | tee -a "$LOGFILE"
        ;;
      *)
        log "${RED}[Error] Automatic MPI installation is not supported on OS: ${OS} by this script.${NC}"
        log "${RED}[Info] Please try installing Open MPI or another MPI implementation manually.${NC}"
        exit 1
        ;;
    esac

    # Verify MPI installation after attempting to install
    log "${BLUE}[Info] Verifying MPI installation attempt...${NC}"
    if check_mpi_installed; then
      log "${GREEN}[Success] MPI installation appears to have succeeded.${NC}"
    else
      log "${RED}[Error] MPI installation appears to have failed or is still not correctly configured. Please check $LOGFILE and install MPI manually.${NC}"
      exit 1
    fi

  else
    log "${YELLOW}[Info] User chose not to install MPI automatically.${NC}"
    log "${RED}[Error] MPI is required for the application to proceed.${NC}"
    log "${YELLOW}Please install an MPI implementation (e.g., Open MPI, MPICH) manually using the appropriate command for your system:${NC}"
    case "$OS" in
      "macOS")
        log "${YELLOW}  Example: brew install open-mpi (Install Homebrew from https://brew.sh if not present)${NC}"
        ;;
      "Linux")
        case "$DISTRO" in
          "ubuntu"|"debian"|"linuxmint"|"pop"|"raspbian")
            log "${YELLOW}  Example: sudo apt-get install libopenmpi-dev${NC}"
            ;;
          "fedora"|"centos"|"rhel")
            log "${YELLOW}  Example: sudo dnf install openmpi-devel (or sudo yum install openmpi-devel)${NC}"
            ;;
          "arch"|"manjaro"|"endeavouros")
            log "${YELLOW}  Example: sudo pacman -S openmpi${NC}"
            ;;
          "opensuse"|"opensuse-leap"|"opensuse-tumbleweed")
            log "${YELLOW}  Example: sudo zypper install openmpi-devel${NC}"
            ;;
          "nixos")
            log "${YELLOW}  Example: nix-env -iA nixpkgs.openmpi or add to configuration.nix${NC}"
            ;;
          *)
            log "${YELLOW}  Please consult your distribution's documentation for installing an MPI development package (e.g., Open MPI, MPICH).${NC}"
            ;;
        esac
        ;;
      "FreeBSD")
        log "${YELLOW}  Example: sudo pkg install openmpi${NC}"
        ;;
      *)
        log "${YELLOW}  Please consult your operating system's documentation for installing an MPI development package.${NC}"
        ;;
    esac
    exit 1
  fi
fi

log "${GREEN}[SUCCESS] MPI is installed and ready for use!${NC}"
log "${GREEN}[SUCCESS] 4DSSE installation can now proceed.${NC}"
exit 0