#!/bin/bash
set -euo pipefail

if [[ "$(id -u)" -ne 0 ]]; then
    echo "start_server.sh must run as root to prepare the read-only workspace." >&2
    exit 1
fi

JUPYTER_TOKEN="${JUPYTER_TOKEN:-huggingface}"
SCCAFM_SOURCE_URL="${SCCAFM_SOURCE_URL:-https://github.com/Catchxu/scCAFM.git}"
SCCAFM_SOURCE_REVISION="${SCCAFM_SOURCE_REVISION:-main}"
SCCAFM_MODEL_ID="${SCCAFM_MODEL_ID:-kaichenxu/scCAFM}"
SCCAFM_MODEL_REVISION="${SCCAFM_MODEL_REVISION:-main}"
SCCAFM_DATASET_ID="${SCCAFM_DATASET_ID:-kaichenxu/scCAFM-data}"
SCCAFM_DATASET_REVISION="${SCCAFM_DATASET_REVISION:-main}"
readonly SCCAFM_WORKSPACE="/workspace/scCAFM"
readonly NOTEBOOK_DIR="$SCCAFM_WORKSPACE"
SCCAFM_ASSETS_DIR="$SCCAFM_WORKSPACE/assets"
SCCAFM_DATA_DIR="$SCCAFM_WORKSPACE/tutorial_data"
JUPYTER_RUNTIME_ROOT="$(mktemp -d /tmp/sccafm-jupyter.XXXXXX)"

resolve_hf_revision() {
    local repo_type="$1"
    local repo_id="$2"
    local revision="$3"

    /home/user/miniconda/bin/python -c '
import sys
from huggingface_hub import HfApi

repo_type, repo_id, revision = sys.argv[1:]
api = HfApi()
if repo_type == "model":
    info = api.model_info(repo_id, revision=revision)
elif repo_type == "dataset":
    info = api.dataset_info(repo_id, revision=revision)
else:
    raise ValueError(f"Unsupported repository type: {repo_type}")
print(info.sha)
' "$repo_type" "$repo_id" "$revision"
}

require_commit_sha() {
    local name="$1"
    local revision="$2"

    if [[ ! "$revision" =~ ^[0-9a-f]{40}$ ]]; then
        echo "Could not resolve $name to a 40-character commit SHA: $revision" >&2
        exit 1
    fi
}

mkdir -p "$(dirname "$SCCAFM_WORKSPACE")"

if [[ ! -d "$SCCAFM_WORKSPACE/.git" ]]; then
    if [[ -e "$SCCAFM_WORKSPACE" ]] && [[ -n "$(ls -A "$SCCAFM_WORKSPACE")" ]]; then
        echo "Cannot prepare scCAFM: $SCCAFM_WORKSPACE exists and is not an empty Git checkout." >&2
        exit 1
    fi

    git clone "$SCCAFM_SOURCE_URL" "$SCCAFM_WORKSPACE"
fi

git -C "$SCCAFM_WORKSPACE" fetch --depth 1 origin "$SCCAFM_SOURCE_REVISION"
git -C "$SCCAFM_WORKSPACE" checkout --detach FETCH_HEAD
SCCAFM_SOURCE_RESOLVED_REVISION="$(git -C "$SCCAFM_WORKSPACE" rev-parse HEAD)"
require_commit_sha "scCAFM source" "$SCCAFM_SOURCE_RESOLVED_REVISION"
pip install --no-deps "$SCCAFM_WORKSPACE"

SCCAFM_MODEL_RESOLVED_REVISION="$(resolve_hf_revision model "$SCCAFM_MODEL_ID" "$SCCAFM_MODEL_REVISION")"
require_commit_sha "scCAFM model" "$SCCAFM_MODEL_RESOLVED_REVISION"
mkdir -p "$SCCAFM_ASSETS_DIR"
hf download "$SCCAFM_MODEL_ID" \
    --repo-type model \
    --revision "$SCCAFM_MODEL_RESOLVED_REVISION" \
    --local-dir "$SCCAFM_ASSETS_DIR"

SCCAFM_DATASET_RESOLVED_REVISION="$(resolve_hf_revision dataset "$SCCAFM_DATASET_ID" "$SCCAFM_DATASET_REVISION")"
require_commit_sha "scCAFM dataset" "$SCCAFM_DATASET_RESOLVED_REVISION"
mkdir -p "$SCCAFM_DATA_DIR"
hf download "$SCCAFM_DATASET_ID" \
    --repo-type dataset \
    --revision "$SCCAFM_DATASET_RESOLVED_REVISION" \
    --local-dir "$SCCAFM_DATA_DIR"

cp /home/user/app/README.md "$SCCAFM_WORKSPACE/SPACE_README.md"

{
    printf 'component\trepository\tselector\tresolved_revision\n'
    printf 'source\t%s\t%s\t%s\n' "$SCCAFM_SOURCE_URL" "$SCCAFM_SOURCE_REVISION" "$SCCAFM_SOURCE_RESOLVED_REVISION"
    printf 'model\t%s\t%s\t%s\n' "$SCCAFM_MODEL_ID" "$SCCAFM_MODEL_REVISION" "$SCCAFM_MODEL_RESOLVED_REVISION"
    printf 'dataset\t%s\t%s\t%s\n' "$SCCAFM_DATASET_ID" "$SCCAFM_DATASET_REVISION" "$SCCAFM_DATASET_RESOLVED_REVISION"
} > "$SCCAFM_WORKSPACE/SPACE_VERSIONS.tsv"

echo "Resolved source revision:  $SCCAFM_SOURCE_RESOLVED_REVISION"
echo "Resolved model revision:   $SCCAFM_MODEL_RESOLVED_REVISION"
echo "Resolved dataset revision: $SCCAFM_DATASET_RESOLVED_REVISION"

jupyter labextension disable "@jupyterlab/apputils-extension:announcements"

chown user:user "$JUPYTER_RUNTIME_ROOT"
chmod 0700 "$JUPYTER_RUNTIME_ROOT"
install -d -o user -g user -m 0700 \
    "$JUPYTER_RUNTIME_ROOT/cache" \
    "$JUPYTER_RUNTIME_ROOT/config" \
    "$JUPYTER_RUNTIME_ROOT/data" \
    "$JUPYTER_RUNTIME_ROOT/ipython" \
    "$JUPYTER_RUNTIME_ROOT/runtime"

chown -R root:root "$SCCAFM_WORKSPACE" /home/user /data/huggingface
chmod -R a-w "$SCCAFM_WORKSPACE" /home/user /data/huggingface
chmod 0755 /tmp /var/tmp

if ! runuser -u user -- test -w "$JUPYTER_RUNTIME_ROOT/runtime"; then
    echo "Refusing to start: the Jupyter runtime directory is not writable by the Jupyter user." >&2
    exit 1
fi

if runuser -u user -- test -w "$SCCAFM_WORKSPACE"; then
    echo "Refusing to start: the scCAFM workspace is still writable by the Jupyter user." >&2
    exit 1
fi

echo "Read-only workspace verified at $SCCAFM_WORKSPACE"

exec runuser -u user -- env -i \
    HOME=/home/user \
    PATH=/home/user/miniconda/bin:/usr/local/bin:/usr/bin:/bin \
    PYTHONPATH=/home/user/app \
    PYTHONDONTWRITEBYTECODE=1 \
    LANG=C.UTF-8 \
    LC_ALL=C.UTF-8 \
    LD_LIBRARY_PATH="${LD_LIBRARY_PATH:-}" \
    NVIDIA_VISIBLE_DEVICES="${NVIDIA_VISIBLE_DEVICES:-}" \
    NVIDIA_DRIVER_CAPABILITIES="${NVIDIA_DRIVER_CAPABILITIES:-}" \
    HF_HOME="$JUPYTER_RUNTIME_ROOT/cache/huggingface" \
    HF_HUB_DISABLE_TELEMETRY=1 \
    XDG_CACHE_HOME="$JUPYTER_RUNTIME_ROOT/cache" \
    JUPYTER_CONFIG_DIR="$JUPYTER_RUNTIME_ROOT/config" \
    JUPYTER_DATA_DIR="$JUPYTER_RUNTIME_ROOT/data" \
    JUPYTER_RUNTIME_DIR="$JUPYTER_RUNTIME_ROOT/runtime" \
    IPYTHONDIR="$JUPYTER_RUNTIME_ROOT/ipython" \
    TMPDIR="$JUPYTER_RUNTIME_ROOT" \
    /home/user/miniconda/bin/jupyter-lab \
    --ip 0.0.0.0 \
    --port 7860 \
    --no-browser \
    --ServerApp.token="$JUPYTER_TOKEN" \
    --ServerApp.contents_manager_class=readonly_contents.ReadOnlyContentsManager \
    --ServerApp.terminals_enabled=False \
    --ServerApp.tornado_settings="{'headers': {'Content-Security-Policy': 'frame-ancestors *'}}" \
    --ServerApp.cookie_options="{'SameSite': 'None', 'Secure': True}" \
    --ServerApp.disable_check_xsrf=True \
    --LabApp.news_url=None \
    --LabApp.check_for_updates_class="jupyterlab.NeverCheckForUpdate" \
    --notebook-dir="$NOTEBOOK_DIR"