#!/usr/bin/env bash
set -euo pipefail

ROOTFS="${ROOTFS_DIR:-/opt/nvidia/Linux_for_Tegra/rootfs}"
LOADER_IMAGE_TAR="${LOADER_IMAGE_TAR:-}"
RESOURCE_API_URL="${RESOURCE_API_URL:-https://api.azaion.com}"
LOADER_DEV_STAGE="${LOADER_DEV_STAGE:-main}"
LOADER_IMAGE="${LOADER_IMAGE:-localhost:5000/loader:arm}"

if [[ ! -d "$ROOTFS" ]]; then
    echo "ERROR: Rootfs directory not found: $ROOTFS" >&2
    exit 1
fi

if [[ -z "$LOADER_IMAGE_TAR" ]]; then
    echo "ERROR: LOADER_IMAGE_TAR not set. Set it in .env to the Loader Docker image tar path." >&2
    exit 1
fi

if [[ ! -f "$LOADER_IMAGE_TAR" ]]; then
    echo "ERROR: Loader image tar not found: $LOADER_IMAGE_TAR" >&2
    exit 1
fi

cleanup_mounts() {
    for mp in proc sys dev/pts dev; do
        sudo umount "$ROOTFS/$mp" 2>/dev/null || true
    done
    if [[ -f "$ROOTFS/etc/resolv.conf.setup-bak" ]]; then
        sudo mv "$ROOTFS/etc/resolv.conf.setup-bak" "$ROOTFS/etc/resolv.conf"
    fi
}

setup_mounts() {
    for mp in proc sys dev dev/pts; do
        mountpoint -q "$ROOTFS/$mp" 2>/dev/null && sudo umount "$ROOTFS/$mp" 2>/dev/null || true
    done
    sudo mount --bind /proc "$ROOTFS/proc"
    sudo mount --bind /sys "$ROOTFS/sys"
    sudo mount --bind /dev "$ROOTFS/dev"
    sudo mount --bind /dev/pts "$ROOTFS/dev/pts"
    if [[ -f "$ROOTFS/etc/resolv.conf" ]]; then
        sudo cp "$ROOTFS/etc/resolv.conf" "$ROOTFS/etc/resolv.conf.setup-bak"
    fi
    sudo cp /etc/resolv.conf "$ROOTFS/etc/resolv.conf"
}

if [[ "$(uname -m)" != "aarch64" ]]; then
    if [[ ! -f "$ROOTFS/usr/bin/qemu-aarch64-static" ]]; then
        sudo cp /usr/bin/qemu-aarch64-static "$ROOTFS/usr/bin/"
    fi
fi

trap cleanup_mounts EXIT

echo "=== Setting up Docker in rootfs ==="
echo "  Rootfs: $ROOTFS"
echo "  Image tar: $LOADER_IMAGE_TAR"
echo ""

setup_mounts

if sudo chroot "$ROOTFS" docker --version &>/dev/null; then
    echo "[1/6] Docker already installed, skipping..."
else
    echo "[1/6] Installing Docker Engine..."
    sudo chroot "$ROOTFS" bash -c '
        apt-get update
        apt-get install -y ca-certificates curl gnupg
        install -m 0755 -d /etc/apt/keyrings
        curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
        chmod a+r /etc/apt/keyrings/docker.asc
        . /etc/os-release
        echo "deb [arch=arm64 signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu $VERSION_CODENAME stable" > /etc/apt/sources.list.d/docker.list
        apt-get update
        apt-get install -y docker-ce docker-ce-cli containerd.io docker-compose-plugin
        apt-get clean
        rm -rf /var/lib/apt/lists/*
    '
fi

if sudo chroot "$ROOTFS" dpkg -l nvidia-container-toolkit 2>/dev/null | grep -q '^ii'; then
    echo "[2/6] NVIDIA Container Toolkit already installed, skipping..."
else
    echo "[2/6] Installing NVIDIA Container Toolkit..."
    sudo chroot "$ROOTFS" bash -c '
        curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
            | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
        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" \
            > /etc/apt/sources.list.d/nvidia-container-toolkit.list
        apt-get update
        apt-get install -y nvidia-container-toolkit
        apt-get clean
        rm -rf /var/lib/apt/lists/*
    '
fi

echo "[3/6] Configuring Docker daemon (NVIDIA default runtime)..."
sudo mkdir -p "$ROOTFS/etc/docker"
sudo tee "$ROOTFS/etc/docker/daemon.json" > /dev/null <<'EOF'
{
    "default-runtime": "nvidia",
    "runtimes": {
        "nvidia": {
            "path": "nvidia-container-runtime",
            "runtimeArgs": []
        }
    }
}
EOF

echo "[4/6] Enabling Docker and containerd services..."
sudo mkdir -p "$ROOTFS/etc/systemd/system/multi-user.target.wants"
sudo ln -sf /lib/systemd/system/docker.service \
    "$ROOTFS/etc/systemd/system/multi-user.target.wants/docker.service"
sudo ln -sf /lib/systemd/system/containerd.service \
    "$ROOTFS/etc/systemd/system/multi-user.target.wants/containerd.service"

echo "[5/6] Creating Azaion application layout..."
sudo mkdir -p "$ROOTFS/opt/azaion/models"
sudo mkdir -p "$ROOTFS/opt/azaion/state"

sudo tee "$ROOTFS/opt/azaion/docker-compose.yml" > /dev/null <<EOF
services:
  loader:
    image: ${LOADER_IMAGE}
    restart: unless-stopped
    ports:
      - "8080:8080"
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock
      - /opt/azaion/docker-compose.yml:/app/docker-compose.yml:ro
      - /opt/azaion/models:/app/models
      - /opt/azaion/state:/app/state
      - /etc/azaion/device.conf:/etc/azaion/device.conf:ro
    environment:
      RESOURCE_API_URL: ${RESOURCE_API_URL}
      LOADER_COMPOSE_FILE: /app/docker-compose.yml
      LOADER_MODEL_DIR: /app/models
      LOADER_DOWNLOAD_STATE_DIR: /app/state
      LOADER_DEV_STAGE: ${LOADER_DEV_STAGE}
      LOADER_ARCH: arm64
EOF

echo "[6/6] Installing Loader image and boot service..."
sudo cp "$LOADER_IMAGE_TAR" "$ROOTFS/opt/azaion/loader-image.tar"

sudo tee "$ROOTFS/opt/azaion/boot.sh" > /dev/null <<'EOF'
#!/bin/bash
set -e
if [ -f /opt/azaion/loader-image.tar ]; then
    docker load -i /opt/azaion/loader-image.tar
    rm -f /opt/azaion/loader-image.tar
fi
docker compose -f /opt/azaion/docker-compose.yml up -d
EOF
sudo chmod 755 "$ROOTFS/opt/azaion/boot.sh"

sudo tee "$ROOTFS/etc/systemd/system/azaion-loader.service" > /dev/null <<'EOF'
[Unit]
Description=Azaion Loader
After=docker.service
Requires=docker.service

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/opt/azaion/boot.sh

[Install]
WantedBy=multi-user.target
EOF

sudo ln -sf /etc/systemd/system/azaion-loader.service \
    "$ROOTFS/etc/systemd/system/multi-user.target.wants/azaion-loader.service"

echo ""
echo "Docker setup complete."