#!/bin/bash

# TVStudy installation script for Ubuntu and WSL (Debian/Ubuntu-based)
# Enhanced with diagnostics for WSL issues

# Ensure script is run with sudo
if [ $EUID -ne 0 ]; then
  echo "Please run this script with sudo."
  exit 1
fi

set -e

# === Configuration ===
CALLING_USER=${SUDO_USER:-$USER}
USER_HOME=$(eval echo ~$CALLING_USER)
TVSTUDY_DIR="$USER_HOME/tvstudy"
LIB_DIR="$TVSTUDY_DIR/lib"
DATA_DIR="$TVSTUDY_DIR/data"
DBASE_DIR="$TVSTUDY_DIR/dbase"
URL_BASE="https://www.fcc.gov/sites/default/files/2025-06"
CDED_BASE="https://www.fcc.gov/sites/default/files"
SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)

# === WSL Detection ===
IS_WSL=false
if grep -qEi "(Microsoft|WSL)" /proc/version &> /dev/null; then
  IS_WSL=true
  echo "WSL environment detected."
else
  echo "Native Linux environment detected."
fi

# === Check MySQL Root Password Early ===
MYSQL_ROOT_PW=""
MYSQL_ROOT_PASSWORD_SET=false

echo "This script installs TVStudy 2.3.0."

if command -v mysql >/dev/null 2>&1; then
  echo "Checking MySQL root access..."
  if mysqladmin ping &> /dev/null; then
    if echo "SELECT 1;" | mysql -u root &> /dev/null; then
      echo "MySQL root access confirmed."
      echo -n "Enter MySQL root password (or leave blank to set 'tvstudy'): "
      read -r MYSQL_ROOT_PW
      [ -z "$MYSQL_ROOT_PW" ] && MYSQL_ROOT_PW="tvstudy"
    else
      echo "MySQL is installed but root access is not available without password."
      echo -n "Enter current MySQL root password: "
      read -r MYSQL_ROOT_PW
      MYSQL_ROOT_PASSWORD_SET=true
    fi
  else
    echo "MySQL server is installed but not running."
    echo -n "Enter desired MySQL root password (or leave blank to use 'tvstudy'): "
    read -r MYSQL_ROOT_PW
    [ -z "$MYSQL_ROOT_PW" ] && MYSQL_ROOT_PW="tvstudy"
  fi
else
  echo "MySQL not currently installed. Password will be set after installation."
  echo -n "Enter desired MySQL root password (or leave blank to use 'tvstudy'): "
  read -r MYSQL_ROOT_PW
  [ -z "$MYSQL_ROOT_PW" ] && MYSQL_ROOT_PW="tvstudy"
fi

# === Check Available Disk Space Function ===
check_disk_space() {
  local required_mb=$1
  local available_kb=$(df --output=avail "$PWD" | tail -n1)
  local available_mb=$((available_kb / 1024))
  if [ "$available_mb" -lt "$required_mb" ]; then
    echo "Insufficient disk space. Required: ${required_mb}MB, Available: ${available_mb}MB"
    exit 1
  fi
}

# === Offer CDED/NLCD Only If Already Installed ===
if [ -d "$TVSTUDY_DIR" ]; then
  echo "TVStudy directory exists. Do you want to install additional data (CDED/NLCD)? (y/n)"
  read -r ADD_DATA_ONLY
  if [[ "$ADD_DATA_ONLY" =~ [Yy] ]]; then
    SKIP_TVSTUDY_INSTALL=true
    echo "In that case, we'll reinstall TVStudy."
  fi
fi

# === Download Options ===
echo "Do you want to download the CDED files? (y/n)"
read -r GET_CDED

echo "Do you want to download the NLCD files? (y/n)"
read -r GET_NLCD

NLCD_FILES=()
if [[ "$GET_NLCD" =~ [Yy] ]]; then
  echo "Download all NLCD files? (y/n)"
  read -r GET_ALL_NLCD

  if [[ "$GET_ALL_NLCD" =~ [Yy] ]]; then
    NLCD_FILES+=("nlcd.tar.xz" "nlcd2011.tar.xz" "nlcd2016.tar.xz" "nlcd2021.tar.xz")
  else
    echo "Select NLCD files to download:"
    echo "Download NLCD 2021? (y/n)"; read -r GET_NLCD21
    echo "Download NLCD 2016? (y/n)"; read -r GET_NLCD16
    echo "Download NLCD 2011? (y/n)"; read -r GET_NLCD11
    echo "Download NLCD 2006? (y/n)"; read -r GET_NLCD0

    [[ "$GET_NLCD21" =~ [Yy] ]] && NLCD_FILES+=("nlcd2021.tar.xz")
    [[ "$GET_NLCD16" =~ [Yy] ]] && NLCD_FILES+=("nlcd2016.tar.xz")
    [[ "$GET_NLCD11" =~ [Yy] ]] && NLCD_FILES+=("nlcd2011.tar.xz")
    [[ "$GET_NLCD0" =~ [Yy] ]] && NLCD_FILES+=("nlcd.tar.xz")
  fi
fi

# === Dynamic Disk Space Calculation ===
base_required=38912
cded_required=40960
nlcd_required=0

for file in "${NLCD_FILES[@]}"; do
  nlcd_required=$((nlcd_required + 3072))
done

[[ "$GET_CDED" =~ [Yy] ]] && cded_selected=$cded_required || cded_selected=0
[[ -z "$SKIP_TVSTUDY_INSTALL" ]] && base_selected=$base_required || base_selected=0

total_required=$(( base_selected + cded_selected + nlcd_required ))
check_disk_space $total_required

# === Prerequisite Installation ===
if [ -z "$SKIP_TVSTUDY_INSTALL" ]; then

  echo "Updating package lists and installing prerequisites..." 

  # Check for WSL and handle PPA/dist-upgrade
  if [ "$IS_WSL" = true ]; then
    echo "WSL detected. Checking for existing Kisak PPA..."
    
    # Check if PPA already exists
    if grep -r "kisak" /etc/apt/sources.list.d/ &>/dev/null; then
      echo "Kisak PPA already installed."
    else
      echo "Adding Kisak PPA..."
      # Ensure software-properties-common is installed for add-apt-repository
      if ! command -v add-apt-repository &> /dev/null; then
          echo "'add-apt-repository' command not found. Installing 'software-properties-common'..."
          apt-get update
          apt-get install -y software-properties-common
      fi
      
      # Add PPA with error handling
      if add-apt-repository ppa:kisak/kisak-meta -y; then
        echo "PPA added successfully."
      else
        echo "Warning: PPA addition had issues, but continuing..."
      fi
    fi
    
    echo "Performing system update..."
    apt update
    apt dist-upgrade -y
  else
    # Standard update for non-WSL systems
    apt update
  fi
  
  echo "Installing TVStudy-specific prerequisites..."
  apt install -y \
    openjdk-21-jre \
    mysql-server \
    libmysqlclient21 \
    ghostscript \
    wget \
    unzip \
    python3 \
    tar \
    pv

  if [ "$MYSQL_ROOT_PASSWORD_SET" = false ]; then
    echo "Setting MySQL root password..."
    echo "ALTER USER 'root'@'localhost' IDENTIFIED WITH mysql_native_password BY '$MYSQL_ROOT_PW'; FLUSH PRIVILEGES;" | mysql -u root || \
      echo "Failed to reset root password. You may need to run mysql_secure_installation manually."
  fi

  echo "Downloading or using local base TVStudy files..."
  FILES=(
    "2013Jan_tvstudy_files.tar.xz"
    "2024May_tvstudy_data_updates.tar.xz"
    "2024May_tvstudy_linux64_files.tar.xz"
    "2024May_checkTerrain.zip"
  )

  mkdir -p "$TVSTUDY_DIR"

  for file in "${FILES[@]}"; do
    if [ -f "$SCRIPT_DIR/$file" ]; then
      echo "Using local copy of $file"
      [ "$SCRIPT_DIR" != "$PWD" ] && cp "$SCRIPT_DIR/$file" .
    else
      echo "Downloading $file..."
      wget -U 'Mozilla/5.0' -c "$URL_BASE/$file" || { echo "Failed to download $file"; exit 1; }
    fi
    chown "$CALLING_USER:$CALLING_USER" "$file"
  done

  echo "Extracting base TVStudy files..."
  
  # Enhanced extraction with diagnostics
  echo "Checking contents of 2013Jan_tvstudy_files.tar.xz..."
  tar -tJf 2013Jan_tvstudy_files.tar.xz | head -10
  
  echo "Extracting to $USER_HOME..."
  pv 2013Jan_tvstudy_files.tar.xz | tar -xJf - -C "$USER_HOME"
  
  # Check what was actually extracted
  echo "Contents of $USER_HOME after extraction:"
  ls -la "$USER_HOME" | grep -E "(tvstudy|TVStudy)"
  
  # Handle potential directory mismatch
  if [ ! -d "$TVSTUDY_DIR" ] && [ -d "$USER_HOME/tvstudy_pack" ]; then
    echo "Found tvstudy_pack directory instead of tvstudy. Renaming..."
    mv "$USER_HOME/tvstudy_pack" "$TVSTUDY_DIR"
  elif [ ! -d "$TVSTUDY_DIR" ]; then
    echo "ERROR: Neither tvstudy nor tvstudy_pack directory found after extraction."
    echo "Available directories in $USER_HOME:"
    ls -la "$USER_HOME"
    exit 1
  fi
  
  # Update directory variables in case of rename
  LIB_DIR="$TVSTUDY_DIR/lib"
  DATA_DIR="$TVSTUDY_DIR/data"
  DBASE_DIR="$TVSTUDY_DIR/dbase"
  
  echo "Extracting remaining files..."
  pv 2024May_tvstudy_data_updates.tar.xz | tar -xJf - -C "$TVSTUDY_DIR" --strip-components=1
  pv 2024May_tvstudy_linux64_files.tar.xz | tar -xJf - -C "$TVSTUDY_DIR"
  unzip -o 2024May_checkTerrain.zip -d "$DBASE_DIR"

  echo "Creating symlink for ghostscript..."
  ln -sf "$(which gs)" "$LIB_DIR/gs"

  echo "Setting ownership and permissions..."
  chmod +x "$LIB_DIR"/* || true
  chown -R "$CALLING_USER:$CALLING_USER" "$TVSTUDY_DIR"
fi

# === Additional Data Downloads ===
if [[ "$GET_CDED" =~ [Yy] ]]; then
  if [ -f "$SCRIPT_DIR/cded.tar.xz" ]; then
    echo "Using local copy of cded.tar.xz"
    [ "$SCRIPT_DIR" != "$PWD" ] && cp "$SCRIPT_DIR/cded.tar.xz" .
  else
    wget -c "$CDED_BASE/cded.tar.xz"
  fi
  pv cded.tar.xz | tar -xJf - -C "$DBASE_DIR"
fi

for nlcd_file in "${NLCD_FILES[@]}"; do
  if [ -f "$SCRIPT_DIR/$nlcd_file" ]; then
    echo "Using local copy of $nlcd_file"
    [ "$SCRIPT_DIR" != "$PWD" ] && cp "$SCRIPT_DIR/$nlcd_file" .
  else
    wget -c "$URL_BASE/$nlcd_file" || { echo "Failed to download $nlcd_file"; exit 1; }
  fi
  pv "$nlcd_file" | tar -xJf - -C "$DBASE_DIR"
done

echo "TVStudy setup completed."
echo "If this was an update, your data has been installed."
echo "If this was a fresh install, TVStudy is now ready to use."
echo "Once comfortable that this was successful, you may delete the downloaded .xz files and .zip file."