#!/usr/bin/env bash
# Analyse this machine, recommend a local coding model, and generate:
#   Dockerfile   (from Dockerfile.template, with the chosen model)
#   .ai-config   (settings used by build-ai-image.sh / run-ai-image.sh)
#   RUN_AIDER.md (how to run aider against the container on this machine)
#
# Usage: ./setup.sh [--model <ollama-tag>] [--gpu|--no-gpu] [--port <n>] [--yes]
set -euo pipefail
cd "$(dirname "$0")"
# Numeric formatting/parsing below assumes '.' as the decimal point (awk, sort -n).
# A locale like de_DE (comma decimal separator) would corrupt those comparisons.
export LC_ALL=C

FORCE_MODEL=""; FORCE_GPU=""; ASSUME_YES=0; FORCE_PORT=""
while [ $# -gt 0 ]; do
  case "$1" in
    --model) FORCE_MODEL="${2:?--model needs a value}"; shift ;;
    --gpu) FORCE_GPU=1 ;;
    --no-gpu) FORCE_GPU=0 ;;
    --port) FORCE_PORT="${2:?--port needs a value}"; shift ;;
    --yes|-y) ASSUME_YES=1 ;;
    -h|--help) sed -n '2,9p' "$0" | sed 's/^# \?//'; exit 0 ;;
    *) echo "Unknown option: $1" >&2; exit 1 ;;
  esac
  shift
done

IMAGE_NAME="local-ai"; CONTAINER_NAME="local-ai"; PORT=""
port_busy() { (exec 3<>"/dev/tcp/127.0.0.1/$1") 2>/dev/null; }
NOTES=()   # machine-specific remarks, echoed to the terminal and RUN_AIDER.md

# Distro-specific commands to make an NVIDIA GPU usable from Docker.
# $1 = need_driver (1/0)
gpu_install_steps() {
  local need_driver=$1 id="" like=""
  [ -r /etc/os-release ] && { id=$(. /etc/os-release; echo "${ID:-}"); like=$(. /etc/os-release; echo "${ID_LIKE:-}"); }
  case " $id $like " in
    *" ubuntu "*|*" debian "*)
      if [ "$need_driver" = 1 ]; then
        if [[ " $id $like " == *" ubuntu "* ]]; then
          echo "# 1. NVIDIA driver (then reboot; with Secure Boot you must enroll the MOK key when prompted)"
          echo "sudo ubuntu-drivers install"
        else
          echo "# 1. NVIDIA driver (needs the non-free component enabled; then reboot)"
          echo "sudo apt install nvidia-driver"
        fi
        echo "sudo reboot"
        echo "nvidia-smi                       # after reboot: should list your GPU"
        echo
      fi
      echo "# 2. NVIDIA container toolkit (lets Docker use the GPU)"
      echo "curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg"
      echo "curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list"
      echo "sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit"
      echo "sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker"
      ;;
    *" fedora "*|*" rhel "*|*" centos "*)
      if [ "$need_driver" = 1 ]; then
        echo "# 1. NVIDIA driver from RPM Fusion (enable RPM Fusion first: https://rpmfusion.org/Configuration), then reboot"
        echo "sudo dnf install akmod-nvidia xorg-x11-drv-nvidia-cuda && sudo reboot"
        echo
      fi
      echo "# 2. NVIDIA container toolkit"
      echo "curl -s -L https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo"
      echo "sudo dnf install -y nvidia-container-toolkit"
      echo "sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker"
      ;;
    *)
      echo "# Install the NVIDIA driver with your distribution's package manager, then the container toolkit:"
      echo "# https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html"
      echo "# and finish with: sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker"
      ;;
  esac
  echo
  echo "# 3. Verify, then re-run ./setup.sh"
  echo "docker run --rm --gpus all ubuntu nvidia-smi"
}

# ---------------------------------------------------------------- detection
cpu_model=$(grep -m1 'model name' /proc/cpuinfo | cut -d: -f2- | xargs || true)
cpu_threads=$(nproc)
ram_gb=$(awk '/MemTotal/ {printf "%.1f", $2/1024/1024}' /proc/meminfo)
disk_gb=$(df -Pk . | awk 'NR==2 {printf "%.0f", $4/1024/1024}')
grep -qw avx2 /proc/cpuinfo && avx2=yes || avx2=no

gpu_name=""; vram_gb=0; gpu_usable=0; nodriver_gpu=0; est_vram=0; toolkit_only=0; GPU_ADVICE=""; GPU_STEPS=""
if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
  gpu_name=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
  vram_mb=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1)
  vram_gb=$(awk -v m="$vram_mb" 'BEGIN {printf "%.1f", m/1024}')
  gpu_usable=1
elif command -v lspci >/dev/null 2>&1 && lspci | grep -iE 'vga|3d' | grep -qi nvidia; then
  gpu_name=$(lspci | grep -iE 'vga|3d' | grep -i nvidia | head -1 | sed 's/.*: //')
  nodriver_gpu=1
  # VRAM can't be read without a driver; estimate from the card name (unknown -> 0)
  case "$gpu_name" in
    *"1650"*|*"1050 Ti"*|*"3050 Mobile"*|*"3050 Laptop"*) est_vram=4 ;;
    *"1660"*|*"2060"*|*"1060"*|*"2070"*|*"2080"*|*"3060 Mobile"*|*"3060 Laptop"*|*"3070"*|*"3060 Ti"*) est_vram=6 ;;
    *"3060"*|*"4060"*|*"2080 Ti"*|*"3080"*|*"4070"*) est_vram=12 ;;
    *"3090"*|*"4090"*) est_vram=24 ;;
  esac
fi

have_docker=no; docker_ok=no; docker_nvidia=no
if command -v docker >/dev/null 2>&1; then
  have_docker=yes
  if docker info >/dev/null 2>&1; then
    docker_ok=yes
    docker info 2>/dev/null | grep -qi nvidia && docker_nvidia=yes
  fi
fi
if [ "$gpu_usable" = 1 ] && [ "$docker_nvidia" = no ]; then
  toolkit_only=1
  gpu_usable=0
fi
[ "$FORCE_GPU" = 1 ] && gpu_usable=1
[ "$FORCE_GPU" = 0 ] && gpu_usable=0
[ "$have_docker" = no ] && NOTES+=("Docker not found - install it before running ./build-ai-image.sh.")
[ "$have_docker" = yes ] && [ "$docker_ok" = no ] && NOTES+=("Docker is installed but not usable by this user (daemon down, or add yourself to the 'docker' group).")
[ "$avx2" = no ] && NOTES+=("CPU has no AVX2; CPU inference will be slow.")
command -v aider >/dev/null 2>&1 && aider_ok=yes || aider_ok=no

echo "=== This machine ==="
printf "  CPU    : %s (%s threads, AVX2: %s)\n" "$cpu_model" "$cpu_threads" "$avx2"
printf "  RAM    : %s GB\n" "$ram_gb"
if [ "$gpu_usable" = 1 ]; then printf "  GPU    : %s (%s GB VRAM) - will be used\n" "$gpu_name" "$vram_gb"
elif [ -n "$gpu_name" ]; then printf "  GPU    : %s (not usable)\n" "$gpu_name"
else printf "  GPU    : none\n"; fi
printf "  Disk   : %s GB free here\n" "$disk_gb"
printf "  Docker : %s   aider: %s\n" "$docker_ok" "$aider_ok"
echo

# ------------------------------------------------ memory budget / context
# What a model must fit into: VRAM when a GPU is used, else RAM minus headroom
# for the OS, aider and the editor. hard_budget/gpu_hard have almost no headroom -
# a model only fitting there is "tight" (usable, but close other apps first).
budget=$(awk -v r="$ram_gb" 'BEGIN {b=r-4; if (b<1) b=1; printf "%.1f", b}')
hard_budget=$(awk -v r="$ram_gb" 'BEGIN {b=r-1; if (b<1) b=1; printf "%.1f", b}')
gpu_budget=$(awk -v v="$vram_gb" 'BEGIN {printf "%.1f", v*0.9}')
gpu_hard=$(awk -v v="$vram_gb" 'BEGIN {printf "%.1f", v*0.98}')

if   awk -v b="$budget" 'BEGIN{exit !(b<6)}';  then CONTEXT=4096
elif awk -v b="$budget" 'BEGIN{exit !(b>=40)}'; then CONTEXT=32768
elif awk -v b="$budget" 'BEGIN{exit !(b>=20)}'; then CONTEXT=16384
else CONTEXT=8192; fi

# ---------------------------------------------------------------- catalogue
mapfile -t rows < <(grep -v '^#' models.tsv | grep -v '^[[:space:]]*$' | sort -t$'\t' -k5,5n)
tags=(); status=(); speed=(); rec=-1
for i in "${!rows[@]}"; do
  IFS=$'\t' read -r tag disk mem cpu_ok rank desc <<<"${rows[$i]}"
  tags[i]=$tag
  fits_ram=$(awk -v m="$mem" -v b="$budget" 'BEGIN{print (m<=b)}')
  fits_hard_ram=$(awk -v m="$mem" -v b="$hard_budget" 'BEGIN{print (m<=b)}')
  fits_gpu=$(awk -v m="$mem" -v b="$gpu_budget" -v u="$gpu_usable" 'BEGIN{print (u && m<=b)}')
  fits_hard_gpu=$(awk -v m="$mem" -v b="$gpu_hard" -v u="$gpu_usable" 'BEGIN{print (u && m<=b)}')
  if [ "$fits_gpu" = 1 ]; then status[i]="fits"; speed[i]="fast (GPU)"
  elif [ "$fits_hard_gpu" = 1 ]; then status[i]="tight"; speed[i]="tight (GPU)"
  elif [ "$fits_ram" = 1 ] && [ "$cpu_ok" = yes ]; then status[i]="fits"; speed[i]="ok (CPU)"
  elif [ "$fits_ram" = 1 ]; then status[i]="slow"; speed[i]="slow (CPU)"
  elif [ "$fits_hard_ram" = 1 ]; then status[i]="tight"; speed[i]="tight (CPU)"
  else status[i]="too big"; speed[i]="-"; fi
  [ "${status[$i]}" = fits ] && rec=$i     # sorted by rank, so last "fits" wins
done
[ "$rec" -ge 0 ] || rec=0

# ------------------------------------------- would a GPU driver help?
if [ "$nodriver_gpu" = 1 ] || [ "$toolkit_only" = 1 ]; then
  [ "$nodriver_gpu" = 1 ] && gvram=$est_vram || gvram=$vram_gb
  best=-1
  if awk -v v="$gvram" 'BEGIN{exit !(v>0)}'; then
    for i in "${!rows[@]}"; do
      mem=$(cut -f3 <<<"${rows[$i]}")
      awk -v m="$mem" -v v="$gvram" 'BEGIN{exit !(m<=v*0.9)}' && best=$i
    done
  fi
  if [ "$nodriver_gpu" = 1 ]; then
    what="the NVIDIA driver and nvidia-container-toolkit"
    have="has no working NVIDIA driver"
  else
    what="nvidia-container-toolkit"
    have="is not usable from Docker (no NVIDIA runtime)"
  fi
  head="Your GPU ($gpu_name) $have, so the model runs on the CPU."
  if [ "$gvram" = 0 ]; then
    GPU_ADVICE="$head If it has enough VRAM, installing $what may speed things up a lot."
  elif [ "$best" -gt "$rec" ]; then
    GPU_ADVICE="$head Recommended: install $what. With ~${gvram} GB VRAM, ${tags[$best]} would run fully on the GPU (better than the current CPU pick ${tags[$rec]})."
  elif [ "$best" -ge 0 ]; then
    GPU_ADVICE="$head Optional, for speed: install $what. With ~${gvram} GB VRAM, models up to ${tags[$best]} would run fully on the GPU and larger ones such as ${tags[$rec]} partly (partial offload), but this VRAM does not unlock bigger models than ${tags[$rec]}."
  else
    GPU_ADVICE=""   # VRAM too small to matter
  fi
  if [ -n "$GPU_ADVICE" ]; then
    GPU_STEPS=$(gpu_install_steps "$nodriver_gpu")
    echo "=== GPU hint ==="
    echo "  $GPU_ADVICE"
    echo "  To do so (optional; then re-run ./setup.sh):"
    sed 's/^/    /' <<<"$GPU_STEPS"
    echo
  fi
fi

# --------------------------------------------------------------------- menu
if [ -z "$FORCE_MODEL" ]; then
  echo "=== Models (budget: ${budget} GB RAM$([ "$gpu_usable" = 1 ] && echo ", ${gpu_budget} GB VRAM")) ==="
  printf "  %-3s %-24s %7s %7s  %-9s %-11s %s\n" "#" "model" "disk" "memory" "fit" "speed" ""
  for i in "${!rows[@]}"; do
    IFS=$'\t' read -r tag disk mem cpu_ok rank desc <<<"${rows[$i]}"
    mark=" "; [ "$i" = "$rec" ] && mark="*"
    printf " %s%-3s %-24s %5s GB %5s GB  %-9s %-11s %s\n" "$mark" "$((i+1))" "$tag" "$disk" "$mem" "${status[$i]}" "${speed[$i]}" "$desc"
  done
  echo
  echo "  * = recommended for this machine: ${tags[$rec]}"
  echo
  if [ "$ASSUME_YES" = 1 ]; then choice=""; else read -r -p "Select model number [$((rec+1))]: " choice; fi
  choice=${choice:-$((rec+1))}
  [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le "${#rows[@]}" ] || { echo "Invalid choice." >&2; exit 1; }
  sel=$((choice-1))
  MODEL=${tags[$sel]}
  if [ "${status[$sel]}" = "tight" ]; then
    echo "Warning: $MODEL is 'tight' on this machine (fits, but leaves little headroom - close other apps first)."
    if [ "$ASSUME_YES" != 1 ]; then read -r -p "Continue anyway? [y/N] " a; [[ "$a" =~ ^[Yy] ]] || exit 1; fi
  elif [ "${status[$sel]}" = "too big" ] || [ "${status[$sel]}" = "slow" ]; then
    echo "Warning: $MODEL is '${status[$sel]}' on this machine (may swap, crawl, or crash)."
    if [ "$ASSUME_YES" != 1 ]; then read -r -p "Continue anyway? [y/N] " a; [[ "$a" =~ ^[Yy] ]] || exit 1; fi
  fi
  sel_disk=$(cut -f2 <<<"${rows[$sel]}")
else
  MODEL=$FORCE_MODEL; sel_disk=0
  for r in "${rows[@]}"; do [ "$(cut -f1 <<<"$r")" = "$MODEL" ] && sel_disk=$(cut -f2 <<<"$r"); done
fi
if awk -v d="$sel_disk" -v f="$disk_gb" 'BEGIN{exit !(d*2>f)}'; then
  NOTES+=("Low disk space: ${disk_gb} GB free, the build needs roughly 2x the model size (${sel_disk} GB).")
fi

# ------------------------------------------------------------ host port
# Default 11434 (Ollama's standard); if something else (e.g. a host Ollama)
# already uses it, take the next free port.
if [ -n "$FORCE_PORT" ]; then PORT=$FORCE_PORT
else
  for PORT in $(seq 11434 11534); do
    # our own container from an earlier run doesn't count as busy
    [ "$(docker port "$CONTAINER_NAME" 11434/tcp 2>/dev/null | head -1 | sed 's/.*://')" = "$PORT" ] && break
    port_busy "$PORT" || break
  done
  [ "$PORT" != 11434 ] && NOTES+=("Port 11434 is already in use on this machine, so the AI container will use port $PORT.")
fi

# ------------------------------------------------------------------- output
USE_GPU=$gpu_usable
sed -e "s|@@MODEL@@|$MODEL|g" -e "s|@@CONTEXT@@|$CONTEXT|g" Dockerfile.template > Dockerfile
cat > .ai-config <<EOF
# generated by ./setup.sh - machine specific, not committed
MODEL="$MODEL"
CONTEXT=$CONTEXT
USE_GPU=$USE_GPU
IMAGE_NAME="$IMAGE_NAME"
CONTAINER_NAME="$CONTAINER_NAME"
PORT=$PORT
EOF

if [ "$USE_GPU" = 1 ]; then where="on your GPU ($gpu_name)"; else where="on your CPU ($cpu_threads threads, no GPU)"; fi
case "$MODEL" in
  *:1.5b|*:3b) EDIT_HINT=$'\nSmall models struggle with diff-style edits; add `--edit-format whole`.\n' ;;
  *) EDIT_HINT="" ;;
esac
if [ "$USE_GPU" = 1 ]; then SPEED="Expect fast responses."
else SPEED="CPU-only: expect a few tokens/second, and a slow first response while the model loads into RAM. Keep prompts and added files small."; fi

{
cat <<EOF
# Running aider with the local AI

Generated by \`./setup.sh\` for this machine:

- CPU: $cpu_model ($cpu_threads threads)
- RAM: $ram_gb GB, GPU: ${gpu_name:-none}$([ "$USE_GPU" = 1 ] && echo " ($vram_gb GB VRAM)" || echo " (not used)")
- Model: \`$MODEL\` running $where, context window $CONTEXT tokens

$SPEED

## 1. Build and start the container

\`\`\`bash
./build-ai-image.sh     # once; downloads the model into the image
./run-ai-image.sh       # starts the server on http://127.0.0.1:$PORT
./run-ai-image.sh stop  # when you are done
\`\`\`

## 2. Install aider (once)

\`\`\`bash
python3 -m pip install --user pipx && pipx install aider-chat
# or: python3 -m pip install aider-install && aider-install
\`\`\`

## 3. Run aider in your project (a git repo)

\`\`\`bash
cd /path/to/your/project
export OLLAMA_API_BASE=http://127.0.0.1:$PORT
aider --model ollama_chat/$MODEL
\`\`\`
$EDIT_HINT
Tip: put those in \`~/.aider.conf.yml\` to skip the flags:

\`\`\`yaml
model: ollama_chat/$MODEL
\`\`\`

(and keep \`export OLLAMA_API_BASE=...\` in your shell profile).

## Troubleshooting

- \`./run-ai-image.sh logs\` shows the server log; \`./run-ai-image.sh status\` shows its state.
- Nothing is sent off the machine once the image is built.
- To switch model or after changing hardware, re-run \`./setup.sh\` then \`./build-ai-image.sh\`.
EOF
if [ -n "$GPU_ADVICE" ]; then
  echo; echo "## Optional: use your GPU"; echo
  echo "$GPU_ADVICE"; echo
  echo '```bash'; echo "$GPU_STEPS"; echo '```'
fi
if [ "${#NOTES[@]}" -gt 0 ]; then
  echo; echo "## Notes about this machine"; echo
  for n in "${NOTES[@]}"; do echo "- $n"; done
fi
} > RUN_AIDER.md

echo
echo "Selected: $MODEL (context $CONTEXT, GPU: $([ "$USE_GPU" = 1 ] && echo yes || echo no))"
for n in "${NOTES[@]}"; do echo "Note: $n"; done
echo "Generated: Dockerfile  .ai-config  RUN_AIDER.md"
echo
echo "Next steps:"
echo "  ./build-ai-image.sh"
echo "  ./run-ai-image.sh"
echo "  then follow RUN_AIDER.md to start aider"