#!/usr/bin/env bash
# Install or update the Lain parser service: one pinned model artifact, a
# private venv and a systemd user unit. Everything lands under --prefix
# (default ~/.local/share/lain-parser) and ~/.config/systemd/user; no sudo
# anywhere.
#
#   bash install.sh                    # install, or update to the channel default
#   bash install.sh --model e5-s2      # pin a specific model version
#   bash install.sh --local            # install from this repo's artifacts/ (dev)
#   bash install.sh --no-start         # install files only
#   bash install.sh --verify           # re-check installed hashes + a live parse
#
# Artifacts come from the Hugging Face repo `Enrell/lain-parser`. manifest.json
# there is the release channel: `default` names the model, and every file's
# sha256 is verified against it before the running unit is touched. A failed
# download or hash check leaves the installed service untouched.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO="$(cd "${SCRIPT_DIR}/../.." && pwd)"
HF_REPO="${LAIN_PARSER_HF_REPO:-Enrell/lain-parser}"
HF_BASE="https://huggingface.co/${HF_REPO}/resolve/main"
SERVICE_VERSION="$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["version"])' \
                   "${SCRIPT_DIR}/app/PROVENANCE.json" 2>/dev/null || echo 0.0.0)"
# --local installs from the repository's artifacts/ tree (development).
LOCAL=0
MODEL_ARG=""
SNAPSHOT_FILES=(anifilebert config.json model.safetensors tokenizer_config.json \
                vocab.json vocab.char.json label_schema.json dapt.json)
PREFIX="${LAIN_PARSER_PREFIX:-${HOME}/.local/share/lain-parser}"
UNIT_DIR="${HOME}/.config/systemd/user"
UNIT="lain-parser.service"
PORT="${LAIN_PARSER_PORT:-9377}"
HOSTS="${LAIN_PARSER_HOST:-127.0.0.1,172.17.0.1}"
# The socket lives inside Lain's data dir on purpose: a container that already
# bind-mounts that directory sees it at /data/parser.sock with no extra mount
# and no firewall rule (host TCP is firewalled from the bridge on many hosts).
SOCKET="${LAIN_PARSER_SOCKET:-${HOME}/.local/share/lain/parser.sock}"
START=1
VERIFY=0

while [[ $# -gt 0 ]]; do
  case "$1" in
    --prefix) PREFIX="$2"; shift 2 ;;
    --port) PORT="$2"; shift 2 ;;
    --host) HOSTS="$2"; shift 2 ;;
    --socket) SOCKET="$2"; shift 2 ;;
    --model) MODEL_ARG="$2"; shift 2 ;;
    --local) LOCAL=1; shift ;;
    --no-start) START=0; shift ;;
    --verify) VERIFY=1; shift ;;
    -h|--help) sed -n '2,16p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    *) echo "unknown argument: $1" >&2; exit 2 ;;
  esac
done

step() { printf '\n== %s\n' "$*"; }
die() { echo "error: $*" >&2; exit 1; }
sha_of() { sha256sum "$1" | cut -d' ' -f1; }

command -v python3 >/dev/null || die "python3 is required"
command -v curl >/dev/null || [[ "$LOCAL" == 1 ]] || die "curl is required for downloads (or use --local)"

# --- resolve the model + snapshot -----------------------------------------

MANIFEST_LOCAL="${PREFIX}/manifest.installed.json"

hf_fetch() { # hf_fetch <repo-path> <dest> — anonymous resolve-URL download
  curl -fSL --retry 3 --connect-timeout 10 "${HF_BASE}/$1" -o "$2"
}

read_manifest() { # read_manifest <file> <python-expression-over-m>
  python3 - "$1" "$2" <<'PY'
import json, sys
m = json.load(open(sys.argv[1]))
print(eval(sys.argv[2]))  # manifest is our own release file; trusted shape
PY
}

if [[ "$VERIFY" == 1 ]]; then
  step "verifying ${PREFIX}"
  [[ -x "${PREFIX}/venv/bin/python" ]] || die "venv missing under ${PREFIX}"
  [[ -f "$MANIFEST_LOCAL" ]] || die "installed manifest missing — was this installed by this installer?"
  installed_name="$(read_manifest "$MANIFEST_LOCAL" 'm["installed_model"]')"
  want="$(read_manifest "$MANIFEST_LOCAL" "m['models']['${installed_name}']['files']['model.int8.pt']")"
  sum="$(sha_of "${PREFIX}/artifact/${installed_name}/model.int8.pt")"
  [[ "$sum" == "$want" ]] || die "artifact hash mismatch: ${sum} != ${want}"
  echo "artifact ${installed_name} hash ok"
  "${PREFIX}/venv/bin/python" "${PREFIX}/app/serve.py" --check \
    "Example.Show.S01E02.1080p-GROUP.mkv" \
    --artifact "${PREFIX}/artifact/${installed_name}" \
    --snapshot "${PREFIX}/anifilebert" >/dev/null
  echo "model loads and parses"
  exit 0
fi

if [[ "$LOCAL" == 1 ]]; then
  ARTIFACT_NAME="${MODEL_ARG:-e5-s2}"
  ARTIFACT_SRC="${REPO}/artifacts/${ARTIFACT_NAME}"
  SNAPSHOT_SRC="${ANIFILEBERT_SNAPSHOT:-${REPO}/artifacts/anifilebert-dapt512-long}"
  [[ -f "${ARTIFACT_SRC}/model.int8.pt" ]] || die "local artifact not found at ${ARTIFACT_SRC}"
  [[ -f "${ARTIFACT_SRC}/config.json" ]] || die "local artifact config not found"
  [[ -d "$SNAPSHOT_SRC" ]] || die "local snapshot not found at ${SNAPSHOT_SRC}"
  MANIFEST_TMP=""
else
  step "fetching release manifest (${HF_REPO})"
  mkdir -p "${PREFIX}/.download"
  MANIFEST_TMP="${PREFIX}/.download/manifest.json"
  hf_fetch "manifest.json" "$MANIFEST_TMP" || die "manifest download failed"
  ARTIFACT_NAME="${MODEL_ARG:-$(read_manifest "$MANIFEST_TMP" 'm["default"]')}"
  read_manifest "$MANIFEST_TMP" "'${ARTIFACT_NAME}' in m['models']" >/dev/null 2>&1 || \
    die "model ${ARTIFACT_NAME} not in manifest"
  # Refuse cleanly when a model needs a newer service than the vendored app ships.
  need="$(read_manifest "$MANIFEST_TMP" "m['models']['${ARTIFACT_NAME}'].get('min_service','0')")"
  python3 - "$need" "$SERVICE_VERSION" <<'PY' || die "model requires service >= $need; update the installer first"
import sys
key = lambda v: [int(x) for x in v.split('.')]
sys.exit(0 if key(sys.argv[2]) >= key(sys.argv[1]) else 1)
PY
  SNAP_NAME="$(read_manifest "$MANIFEST_TMP" "m['models']['${ARTIFACT_NAME}']['snapshot']")"
  echo "channel default: $(read_manifest "$MANIFEST_TMP" 'm["default"]') -> installing ${ARTIFACT_NAME} (snapshot ${SNAP_NAME})"
fi

step "installing into ${PREFIX}"
mkdir -p "${PREFIX}/app/fp" "${PREFIX}/artifact" "${PREFIX}/anifilebert"
cp -f "${SCRIPT_DIR}/app/serve.py" "${PREFIX}/app/serve.py"
cp -f "${SCRIPT_DIR}"/app/fp/*.py "${PREFIX}/app/fp/"
cp -f "${SCRIPT_DIR}/requirements.txt" "${PREFIX}/requirements.txt"

if [[ "$LOCAL" == 1 ]]; then
  cp -f "${ARTIFACT_SRC}/model.int8.pt" "${ARTIFACT_SRC}/config.json" \
        "${PREFIX}/artifact/${ARTIFACT_NAME}/" 2>/dev/null || {
    mkdir -p "${PREFIX}/artifact/${ARTIFACT_NAME}"
    cp -f "${ARTIFACT_SRC}/model.int8.pt" "${ARTIFACT_SRC}/config.json" \
          "${PREFIX}/artifact/${ARTIFACT_NAME}/"; }
  for entry in "${SNAPSHOT_FILES[@]}"; do
    cp -RL "${SNAPSHOT_SRC}/${entry}" "${PREFIX}/anifilebert/" 2>/dev/null || true
  done
else
  # Download each needed file to a staging dir, verify sha256 against the
  # manifest, then move into place. Nothing below $PREFIX/artifact or
  # $PREFIX/anifilebert changes until every hash checks out.
  stage="${PREFIX}/.download/${ARTIFACT_NAME}"
  mkdir -p "${stage}/model" "${stage}/snapshot"
  while read -r f sha; do
    [[ "$f" == "model.int8.pt" || "$f" == "config.json" ]] || continue
    dest="${stage}/model/${f}"
    installed="${PREFIX}/artifact/${ARTIFACT_NAME}/${f}"
    if [[ -f "$dest" && "$(sha_of "$dest")" == "$sha" ]] \
       || { [[ -f "$installed" && "$(sha_of "$installed")" == "$sha" ]] \
            && cp -f "$installed" "$dest"; }; then
      echo "already have ${f}"; continue
    fi
    echo "fetching models/${ARTIFACT_NAME}/${f}"
    hf_fetch "models/${ARTIFACT_NAME}/${f}" "$dest" || die "download failed: $f"
    [[ "$(sha_of "$dest")" == "$sha" ]] || die "hash mismatch after download: $f"
  done < <(read_manifest "$MANIFEST_TMP" \
           "'\n'.join(f'{k} {v}' for k,v in m['models']['${ARTIFACT_NAME}']['files'].items())")
  # snapshot files the service needs: the anifilebert/ package plus the
  # top-level files listed in SNAPSHOT_FILES (dapt.json tolerated if absent).
  needed="$(printf '%s\n' "${SNAPSHOT_FILES[@]}" | grep -v '^anifilebert$')"
  while read -r f sha; do
    [[ -n "$f" ]] || continue
    dest="${stage}/snapshot/${f}"
    mkdir -p "$(dirname "$dest")"
    installed="${PREFIX}/anifilebert/${f}"
    if [[ -f "$dest" && "$(sha_of "$dest")" == "$sha" ]] \
       || { [[ -f "$installed" && "$(sha_of "$installed")" == "$sha" ]] \
            && cp -f "$installed" "$dest"; }; then
      echo "already have snapshot/${f}"; continue
    fi
    echo "fetching snapshots/${SNAP_NAME}/${f}"
    if hf_fetch "snapshots/${SNAP_NAME}/${f}" "$dest" && [[ "$(sha_of "$dest")" == "$sha" ]]; then
      continue
    fi
    if [[ "$f" == "dapt.json" ]]; then
      echo "warning: ${f} unavailable — skipping (provenance only)"
      rm -f "$dest"; continue
    fi
    die "snapshot file ${f} failed to fetch or verify"
  done < <(python3 - "$MANIFEST_TMP" "$SNAP_NAME" $needed <<'PY'
import json, sys
m = json.load(open(sys.argv[1]))
snap = m["snapshots"][sys.argv[2]]
want = set(sys.argv[3:])
for rel, sha in sorted(snap["sha256"].items()):
    if rel in want or rel.startswith("anifilebert/"):
        print(rel, sha)
PY
)
  mkdir -p "${PREFIX}/artifact/${ARTIFACT_NAME}"
  cp -f "${stage}/model/model.int8.pt" "${stage}/model/config.json" \
        "${PREFIX}/artifact/${ARTIFACT_NAME}/"
  cp -Rf "${stage}/snapshot/." "${PREFIX}/anifilebert/"
fi

# Record the manifest used for this install (offline verify + update checks).
if [[ -n "${MANIFEST_TMP}" ]]; then
  python3 - "$MANIFEST_TMP" "$MANIFEST_LOCAL" "$ARTIFACT_NAME" <<'PY'
import json, sys
m = json.load(open(sys.argv[1])); m["installed_model"] = sys.argv[3]
json.dump(m, open(sys.argv[2], "w"), indent=2)
PY
fi

step "recording hashes"
{
  echo "# installed $(date -u +%Y-%m-%dT%H:%M:%SZ) by $(id -un)@$(hostname)"
  (cd "$PREFIX" && find app artifact anifilebert requirements.txt -type f \
     -exec sha256sum {} + | sort -k2)
} > "${PREFIX}/SHA256SUMS"
echo "wrote ${PREFIX}/SHA256SUMS"

step "creating the venv (CPU wheels only)"
if [[ ! -x "${PREFIX}/venv/bin/python" ]]; then
  if command -v uv >/dev/null 2>&1; then
    uv venv --python "${LAIN_PARSER_PYTHON:-3.12}" "${PREFIX}/venv"
    uv pip install --python "${PREFIX}/venv/bin/python" \
      --index-url https://download.pytorch.org/whl/cpu "torch==2.14.0"
    uv pip install --python "${PREFIX}/venv/bin/python" -r "${PREFIX}/requirements.txt"
  else
    python3 -m venv "${PREFIX}/venv"
    "${PREFIX}/venv/bin/pip" install --upgrade pip >/dev/null
    "${PREFIX}/venv/bin/pip" install --index-url https://download.pytorch.org/whl/cpu \
      "torch==2.14.0"
    "${PREFIX}/venv/bin/pip" install -r "${PREFIX}/requirements.txt"
  fi
else
  echo "venv already present, keeping it"
fi
"${PREFIX}/venv/bin/python" -c "import torch, transformers; print('torch', torch.__version__, '| transformers', transformers.__version__)"

step "smoke test before touching the service"
"${PREFIX}/venv/bin/python" "${PREFIX}/app/serve.py" --check \
  "Example.Show.S01E02.1080p-GROUP.mkv" \
  --artifact "${PREFIX}/artifact/${ARTIFACT_NAME}" \
  --snapshot "${PREFIX}/anifilebert" >/dev/null \
  || die "new artifact failed the smoke parse — the running service was left alone"

step "systemd user unit"
mkdir -p "${UNIT_DIR}"
unit_tmp="$(mktemp)"
cat > "$unit_tmp" <<UNIT
# Generated by filename-parser packaging/lain-parser/install.sh -- yours to edit.
[Unit]
Description=Lain local filename parser (${ARTIFACT_NAME})
After=network.target docker.service

[Service]
WorkingDirectory=${PREFIX}
ExecStart=${PREFIX}/venv/bin/python ${PREFIX}/app/serve.py --artifact ${PREFIX}/artifact/${ARTIFACT_NAME} --snapshot ${PREFIX}/anifilebert --host ${HOSTS} --port ${PORT} --unix-socket ${SOCKET}
Environment=PYTHONUNBUFFERED=1
Restart=on-failure
RestartSec=2

[Install]
WantedBy=default.target
UNIT
if [[ ! -f "${UNIT_DIR}/${UNIT}" ]]; then
  install -m 0644 "$unit_tmp" "${UNIT_DIR}/${UNIT}"
  echo "wrote ${UNIT_DIR}/${UNIT}"
elif diff -q "$unit_tmp" "${UNIT_DIR}/${UNIT}" >/dev/null; then
  echo "unit already up to date"
else
  # A model bump changes ExecStart; keep the old unit as .bak for inspection.
  cp -f "${UNIT_DIR}/${UNIT}" "${UNIT_DIR}/${UNIT}.bak"
  install -m 0644 "$unit_tmp" "${UNIT_DIR}/${UNIT}"
  echo "unit updated (old copy: ${UNIT}.bak)"
fi
rm -f "$unit_tmp"

if [[ "$START" == 1 ]] && command -v systemctl >/dev/null 2>&1; then
  systemctl --user daemon-reload
  systemctl --user enable --now "$UNIT" || true
  systemctl --user restart "$UNIT" || true
  if command -v curl >/dev/null 2>&1; then
    ok=""
    for _ in $(seq 1 20); do
      sleep 1
      for addr in ${HOSTS//,/ }; do
        if curl -fsS --max-time 2 "http://${addr}:${PORT}/health" >/dev/null 2>&1; then ok="$addr"; break 2; fi
      done
    done
    if [[ -n "$ok" ]]; then
      curl -fsS --max-time 3 "http://${ok}:${PORT}/health"; echo
    else
      echo "warning: health check did not answer in 20s (see: journalctl --user -u ${UNIT})"
    fi
  fi
fi

step "done"
cat <<SUMMARY
prefix      ${PREFIX}
artifact    ${ARTIFACT_NAME}
source      $([[ "$LOCAL" == 1 ]] && echo "local artifacts/" || echo "${HF_REPO} (manifest)")
unit        ${UNIT_DIR}/${UNIT}
native URL  http://127.0.0.1:${PORT}
docker URL  http://172.17.0.1:${PORT}   (or host.docker.internal with host-gateway)
docker UDS  unix:///data/parser.sock    (socket at ${SOCKET})
update      bash install.sh            (re-run: installs the channel default)
rollback    bash install.sh --model <previous-name>
verify      bash install.sh --verify
SUMMARY