#!/usr/bin/env bash
# Setup 1TB data disk on GMKtec K11
# nvme0n1 = 1TB data  |  nvme1n1 = 128GB system
set -euo pipefail

DATA_DISK=/dev/nvme0n1
DATA_PART=/dev/nvme0n1p1
MOUNT_POINT=/data

echo "=== Sprawdzam dyski ==="
lsblk -o NAME,SIZE,FSTYPE,MOUNTPOINT,MODEL

# Safety: refuse if DATA_DISK has mountpoints (system disk)
if lsblk -n -o MOUNTPOINT "$DATA_DISK" | grep -q '/'; then
  echo "BŁĄD: $DATA_DISK ma partycje zamontowane — przerwanie!"
  exit 1
fi

# Create partition if missing
if ! lsblk -n "$DATA_PART" &>/dev/null; then
  echo "=== Partycjonowanie $DATA_DISK ==="
  parted -s "$DATA_DISK" mklabel gpt
  parted -s "$DATA_DISK" mkpart primary ext4 0% 100%
  partprobe "$DATA_DISK"
  sleep 2
fi

# Format if no filesystem
if ! blkid -o value -s TYPE "$DATA_PART" 2>/dev/null | grep -q ext4; then
  echo "=== Formatowanie $DATA_PART (ext4) ==="
  mkfs.ext4 -F -L data1tb "$DATA_PART"
fi

UUID=$(blkid -o value -s UUID "$DATA_PART")
echo "UUID=$UUID"

# Mount point
mkdir -p "$MOUNT_POINT"
if ! mountpoint -q "$MOUNT_POINT"; then
  echo "=== Montowanie $DATA_PART -> $MOUNT_POINT ==="
  mount "$DATA_PART" "$MOUNT_POINT"
fi

# fstab
FSTAB_LINE="UUID=${UUID} ${MOUNT_POINT} ext4 defaults,noatime 0 2"
if ! grep -q "$MOUNT_POINT" /etc/fstab; then
  echo "=== Dodaję wpis do /etc/fstab ==="
  cp /etc/fstab /etc/fstab.bak.$(date +%Y%m%d%H%M%S)
  echo "$FSTAB_LINE" >> /etc/fstab
else
  echo "=== fstab już zawiera $MOUNT_POINT ==="
fi

mount -a
df -h "$MOUNT_POINT"

# Directory structure
echo "=== Tworzę strukturę katalogów ==="
mkdir -p "$MOUNT_POINT/docker"
mkdir -p "$MOUNT_POINT/apps/vllm/huggingface"
mkdir -p "$MOUNT_POINT/apps/gguf/"{qwen3.6-27b,gemma-4-12b}
mkdir -p "$MOUNT_POINT/apps/localai/"{models,backends,configuration,images,data}
mkdir -p "$MOUNT_POINT/apps/comfyui/storage"
mkdir -p "$MOUNT_POINT/apps/comfyui/cache/"{hf-hub,torch-hub}
mkdir -p "$MOUNT_POINT/apps/comfyui/"{models,input,output,custom_nodes,workflows}
chown -R tomasz-syn-grzegorza:tomasz-syn-grzegorza "$MOUNT_POINT/apps"

# Docker data-root migration
echo "=== Migracja Docker data-root ==="
systemctl stop docker docker.socket 2>/dev/null || true

if [ -d /var/lib/docker ] && [ "$(ls -A /var/lib/docker 2>/dev/null)" ]; then
  echo "Kopiuję /var/lib/docker -> $MOUNT_POINT/docker ..."
  rsync -aHAX /var/lib/docker/ "$MOUNT_POINT/docker/"
fi

mkdir -p "$MOUNT_POINT/docker"

python3 <<'PYEOF'
import json
from pathlib import Path
p = Path("/etc/docker/daemon.json")
cfg = json.loads(p.read_text()) if p.exists() else {}
cfg["data-root"] = "/data/docker"
p.write_text(json.dumps(cfg, indent=2) + "\n")
print(p.read_text())
PYEOF

systemctl start docker
sleep 2

echo "=== Weryfikacja ==="
docker info | grep "Docker Root Dir"
df -h / "$MOUNT_POINT"
lsblk -o NAME,SIZE,FSTYPE,MOUNTPOINT,MODEL

echo "=== Test Docker ==="
docker run --rm hello-world
docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu24.04 nvidia-smi

echo ""
echo "=== GOTOWE ==="
echo "Dysk danych: $MOUNT_POINT ($UUID)"
echo "Docker Root: /data/docker"
echo "Modele vLLM: /data/apps/vllm/huggingface"
echo "LocalAI:     /data/apps/localai/models"