#!/usr/bin/env bash
#
# install_backend.sh — build everything lektor_ai_comfy.py needs, on this or any
# other Linux machine: the ComfyUI + Higgs Audio V3 speech backend, and the
# Polish number normalizer with its own environment and models.
#
# After this finishes there is nothing left to install by hand, with one
# exception the script cannot do for you: one normalizer model is gated behind
# a licence you have to accept with your own Hugging Face account. The script
# says so and installs everything else regardless.
#
# Everything lands in three removable directories and nothing touches the system:
#   ~/miniforge3        — conda distribution (its own Python)
#   ~/comfy-tts-app     — ComfyUI, the TTS-Audio-Suite node, and the TTS model
#   ./normalizer/       — the normalizer venv and its models, beside this script
#
# No sudo, no apt, no changes to system Python. `conda init` is never run, so
# your .bashrc is left alone.
#
# Usage:
#   ./install_backend.sh                 # pinned versions (recommended)
#   ./install_backend.sh --latest        # newest ComfyUI + suite instead
#   ./install_backend.sh --skip-model    # don't pre-download the 8.7 GB TTS model
#   ./install_backend.sh --skip-normalizer   # ComfyUI only, no number normalizer
#   ./install_backend.sh --lock FILE     # restore exact package set from a lockfile
#   ./install_backend.sh --help
#
# Hugging Face credentials, if you have them, are picked up automatically from
# `hf auth login` or from HF_TOKEN in the environment.
#
# Override locations with environment variables:
#   CONDA_ROOT=/opt/miniforge3 COMFY_ROOT=/data/comfy ./install_backend.sh
#
set -euo pipefail

# ── Configuration ────────────────────────────────────────────────────────────
CONDA_ROOT="${CONDA_ROOT:-$HOME/miniforge3}"
COMFY_ROOT="${COMFY_ROOT:-$HOME/comfy-tts-app}"
ENV_NAME="${ENV_NAME:-comfy-tts}"
COMFY_PORT="${COMFY_PORT:-8189}"
PYTHON_VERSION="${PYTHON_VERSION:-3.12}"

# Exact versions of the working installation this script reproduces.
COMFYUI_REPO="https://github.com/comfyanonymous/ComfyUI.git"
COMFYUI_COMMIT="7cb784e0f48784bb6ed588912e186e5ee1e9ee68"   # v0.26.0-13-g7cb784e0
TTS_SUITE_REPO="https://github.com/diodiogod/TTS-Audio-Suite.git"
TTS_SUITE_TAG="v5.3.0"                                       # 09df1681
MINIFORGE_URL="https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh"

# PyTorch: CUDA 12.8 wheels. These carry their own CUDA runtime, so the
# system CUDA toolkit version is irrelevant — only the NVIDIA driver matters.
TORCH_INDEX="https://download.pytorch.org/whl/cu128"
TORCH_PINS=(torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0)

# The TTS model. The suite downloads this automatically on first use; doing it
# here just means the first generation isn't a surprise 8.7 GB wait.
HF_MODEL_REPO="bosonai/higgs-audio-v3-tts-4b"
MODEL_SUBDIR="models/TTS/higgs_audio_v3/higgs-audio-v3-tts-4b"

# ~10 conda env + ~3 ComfyUI + ~9 TTS model, then the normalizer: ~7 venv
# (CUDA wheels again, in its own environment) + ~8 inflection model + ~1 gated.
REQUIRED_DISK_GB=55

# ── Options ──────────────────────────────────────────────────────────────────
USE_LATEST=0
SKIP_MODEL=0
SKIP_NORMALIZER=0
LOCKFILE=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --latest)     USE_LATEST=1; shift ;;
        --skip-model) SKIP_MODEL=1; shift ;;
        --skip-normalizer) SKIP_NORMALIZER=1; shift ;;
        --lock)       LOCKFILE="${2:-}"; shift 2 ;;
        -h|--help)    awk 'NR>1 {if (/^#/) {sub(/^# ?/,""); print} else exit}' "$0"; exit 0 ;;
        *)            echo "Unknown option: $1 (try --help)" >&2; exit 1 ;;
    esac
done

# ── Output helpers ───────────────────────────────────────────────────────────
if [[ -t 1 ]]; then
    B=$'\033[1m'; G=$'\033[32m'; Y=$'\033[33m'; R=$'\033[31m'; N=$'\033[0m'
else
    B=""; G=""; Y=""; R=""; N=""
fi
step() { echo; echo "${B}==> $*${N}"; }
ok()   { echo "  ${G}[+]${N} $*"; }
warn() { echo "  ${Y}[!]${N} $*"; }
die()  { echo "  ${R}[X]${N} $*" >&2; exit 1; }
have() { command -v "$1" >/dev/null 2>&1; }

# Activate the conda env. conda's shell functions dislike `set -u`, so it is
# lifted for the duration of the call only.
activate_env() {
    # shellcheck disable=SC1091
    source "$CONDA_ROOT/etc/profile.d/conda.sh"
    set +u
    conda activate "$ENV_NAME"
    set -u
}

ENV_PY="$CONDA_ROOT/envs/$ENV_NAME/bin/python"
pip_install() { "$ENV_PY" -m pip install --no-input "$@"; }

# ── 0. Preflight ─────────────────────────────────────────────────────────────
step "Preflight checks"

[[ "$(uname -s)" == "Linux" ]] || die "This script targets Linux (found $(uname -s))."

have git || die "git is required. Install it first (e.g. sudo apt install git)."
if have curl; then DL=(curl -fL -o); elif have wget; then DL=(wget -O); else
    die "Need curl or wget to download Miniforge."
fi
ok "git and a downloader are present"

if have nvidia-smi && nvidia-smi >/dev/null 2>&1; then
    ok "GPU: $(nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader | head -1)"
else
    warn "No usable nvidia-smi. Install will proceed but TTS will run on CPU (very slow)."
    warn "The cu128 wheels need an NVIDIA driver >= 525; the CUDA toolkit is NOT required."
fi

mkdir -p "$(dirname "$COMFY_ROOT")"
avail_gb=$(df -BG --output=avail "$(dirname "$COMFY_ROOT")" | tail -1 | tr -dc '0-9')
if (( avail_gb < REQUIRED_DISK_GB )); then
    warn "Only ${avail_gb} GB free where $COMFY_ROOT will live; ~${REQUIRED_DISK_GB} GB recommended."
    warn "Use --skip-model, or set COMFY_ROOT to a roomier disk."
else
    ok "Disk space: ${avail_gb} GB available"
fi

if have ffmpeg; then
    ok "ffmpeg found at $(command -v ffmpeg) (used by lektor_ai_comfy.py, not by ComfyUI)"
else
    warn "ffmpeg not found. ComfyUI does not need it, but lektor_ai_comfy.py does."
    warn "Install it with: conda install -n $ENV_NAME -c conda-forge ffmpeg   (no sudo)"
fi

if (( USE_LATEST )); then
    warn "--latest: installing newest ComfyUI and TTS-Audio-Suite rather than the pinned"
    warn "versions this project was validated against. Expect to debug breakage yourself."
else
    ok "Pinned: ComfyUI ${COMFYUI_COMMIT:0:8} (v0.26.0), TTS-Audio-Suite $TTS_SUITE_TAG"
fi

# ── 1. Miniforge ─────────────────────────────────────────────────────────────
step "Miniforge (self-contained conda in $CONDA_ROOT)"

if [[ -x "$CONDA_ROOT/bin/conda" ]]; then
    ok "Already installed: $("$CONDA_ROOT/bin/conda" --version)"
else
    tmp_installer="$(mktemp -t miniforge-XXXXXX.sh)"
    trap 'rm -f "$tmp_installer"' EXIT
    echo "  Downloading Miniforge…"
    "${DL[@]}" "$tmp_installer" "$MINIFORGE_URL"
    echo "  Installing to $CONDA_ROOT (batch mode, no shell integration)…"
    bash "$tmp_installer" -b -p "$CONDA_ROOT"
    rm -f "$tmp_installer"; trap - EXIT
    ok "Installed $("$CONDA_ROOT/bin/conda" --version)"
fi

# ── 2. Conda environment ─────────────────────────────────────────────────────
step "Conda environment '$ENV_NAME' (Python $PYTHON_VERSION)"

# shellcheck disable=SC1091
source "$CONDA_ROOT/etc/profile.d/conda.sh"
if conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
    ok "Environment already exists — reusing it"
else
    conda create -n "$ENV_NAME" "python=$PYTHON_VERSION" -y
    ok "Created"
fi

activate_env
ok "Active: $(python -V) at $(command -v python)"

[[ "$(command -v python)" == "$CONDA_ROOT/envs/$ENV_NAME/"* ]] \
    || die "Wrong Python active — refusing to install into the wrong environment."

pip_install --upgrade pip setuptools wheel >/dev/null
ok "pip toolchain up to date"

# ── Exact restore path ───────────────────────────────────────────────────────
if [[ -n "$LOCKFILE" ]]; then
    step "Restoring exact package set from $LOCKFILE"
    [[ -f "$LOCKFILE" ]] || die "Lockfile not found: $LOCKFILE"
    pip_install --extra-index-url "$TORCH_INDEX" -r "$LOCKFILE"
    ok "Lockfile applied — skipping individual package steps"
fi

# ── 3. PyTorch ───────────────────────────────────────────────────────────────
step "PyTorch (CUDA 12.8 wheels)"

if [[ -z "$LOCKFILE" ]]; then
    if (( USE_LATEST )); then
        pip_install --index-url "$TORCH_INDEX" torch torchvision torchaudio
    else
        pip_install --index-url "$TORCH_INDEX" "${TORCH_PINS[@]}"
    fi
fi
python - <<'PY'
import torch
print(f"  torch {torch.__version__}  cuda_available={torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"  device: {torch.cuda.get_device_name(0)}")
PY
ok "PyTorch ready"

# ── 4. ComfyUI ───────────────────────────────────────────────────────────────
step "ComfyUI"

mkdir -p "$COMFY_ROOT"
COMFY_DIR="$COMFY_ROOT/ComfyUI"

if [[ -d "$COMFY_DIR/.git" ]]; then
    ok "Repository already present at $COMFY_DIR"
else
    git clone "$COMFYUI_REPO" "$COMFY_DIR"
fi

if (( ! USE_LATEST )); then
    git -C "$COMFY_DIR" fetch --quiet origin
    git -C "$COMFY_DIR" checkout --quiet "$COMFYUI_COMMIT"
    ok "Checked out $COMFYUI_COMMIT (detached HEAD — expected)"
else
    git -C "$COMFY_DIR" pull --quiet --ff-only origin master || warn "Could not fast-forward; leaving as-is"
fi

[[ -z "$LOCKFILE" ]] && pip_install -r "$COMFY_DIR/requirements.txt"
ok "ComfyUI dependencies installed"

# ── 5. TTS-Audio-Suite ───────────────────────────────────────────────────────
step "TTS-Audio-Suite custom node"

SUITE_DIR="$COMFY_DIR/custom_nodes/TTS-Audio-Suite"
mkdir -p "$COMFY_DIR/custom_nodes"

if [[ -d "$SUITE_DIR/.git" ]]; then
    ok "Repository already present"
else
    git clone "$TTS_SUITE_REPO" "$SUITE_DIR"
fi

if (( ! USE_LATEST )); then
    git -C "$SUITE_DIR" fetch --quiet --tags origin
    git -C "$SUITE_DIR" checkout --quiet "tags/$TTS_SUITE_TAG"
    ok "Checked out $TTS_SUITE_TAG"
else
    git -C "$SUITE_DIR" pull --quiet --ff-only || warn "Could not fast-forward; leaving as-is"
fi

if [[ -z "$LOCKFILE" ]]; then
    # install.py resolves this suite's many conflicting dependencies itself
    # (selective --no-deps installs, numpy pinning). Do not replace it with a
    # plain `pip install -r requirements.txt` — that file deliberately omits
    # the packages that need special handling.
    echo "  Running the suite's own installer (this pulls a lot and takes a while)…"
    ( cd "$SUITE_DIR" && python install.py )
fi
ok "TTS-Audio-Suite installed"

# ── 6. Model ─────────────────────────────────────────────────────────────────
step "Higgs Audio V3 model ($HF_MODEL_REPO)"

MODEL_DIR="$COMFY_DIR/$MODEL_SUBDIR"
if (( SKIP_MODEL )); then
    warn "Skipped (--skip-model). It downloads automatically on the first generation."
elif [[ -f "$MODEL_DIR/model.safetensors" ]]; then
    ok "Already present ($(du -h "$MODEL_DIR/model.safetensors" | cut -f1))"
else
    echo "  Downloading ~8.7 GB into $MODEL_DIR …"
    mkdir -p "$MODEL_DIR"
    if ! HF_MODEL_REPO="$HF_MODEL_REPO" MODEL_DIR="$MODEL_DIR" python - <<'PY'
import os, sys
from huggingface_hub import snapshot_download
try:
    snapshot_download(repo_id=os.environ["HF_MODEL_REPO"],
                      local_dir=os.environ["MODEL_DIR"])
except Exception as exc:
    print(f"download failed: {exc}", file=sys.stderr)
    sys.exit(1)
PY
    then
        warn "Model download failed."
        warn "This model is licensed for research / non-commercial use and may require"
        warn "accepting its terms on huggingface.co first, then: hf auth login"
        warn "The suite will retry automatically on the first generation."
    else
        ok "Model downloaded"
    fi
fi

# ── 7. Number normalizer ─────────────────────────────────────────────────────
# Its own venv, deliberately: the normalizer needs transformers>=5 and the TTS
# environment above is pinned to versions that must not move. They talk over a
# pipe instead, so the conflict never has to be resolved.
NORMALIZER_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/normalizer"

if [[ $SKIP_NORMALIZER -eq 1 ]]; then
    step "Number normalizer (skipped)"
    warn "Digits will reach the TTS as digits, which it often misreads."
elif [[ ! -x "$NORMALIZER_DIR/setup.sh" ]]; then
    step "Number normalizer"
    warn "$NORMALIZER_DIR/setup.sh not found — skipping."
else
    step "Number normalizer (own venv + models)"
    if "$NORMALIZER_DIR/setup.sh" --model; then
        ok "Normalizer environment ready"
    else
        warn "Normalizer setup reported a problem — see above."
        warn "The pipeline still runs; numbers fall back to num2words."
    fi
fi

# ── 8. Launcher ──────────────────────────────────────────────────────────────
step "Launcher script"

LAUNCHER="$COMFY_ROOT/start_comfyui.sh"
cat > "$LAUNCHER" <<EOF
#!/bin/bash
# Generated by install_backend.sh — starts ComfyUI on port $COMFY_PORT.
source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate $ENV_NAME
cd "$COMFY_DIR" && python main.py --port $COMFY_PORT
EOF
chmod +x "$LAUNCHER"
ok "Wrote $LAUNCHER"

# ── 9. Verification ──────────────────────────────────────────────────────────
step "Verifying installation"

python - <<'PY'
import importlib, sys
failed = []
for mod in ("torch", "torchaudio", "transformers", "safetensors", "soundfile", "librosa"):
    try:
        m = importlib.import_module(mod)
        print(f"  [+] {mod:<13} {getattr(m, '__version__', '?')}")
    except Exception as exc:
        print(f"  [X] {mod:<13} {exc}")
        failed.append(mod)
sys.exit(1 if failed else 0)
PY
ok "Core imports succeed"

NORMALIZER_STATE="not installed"
if [[ -x "$NORMALIZER_DIR/venv/bin/python" ]]; then
    if "$NORMALIZER_DIR/venv/bin/python" "$NORMALIZER_DIR/normalizer.py" --status \
            2>/dev/null | grep -q '"mode": "model"'; then
        NORMALIZER_STATE="ready (neural)"
        ok "Normalizer ready"
    else
        NORMALIZER_STATE="fallback only (gated model missing)"
        warn "Normalizer installed but running on the num2words fallback"
    fi
fi

echo
echo "${B}${G}Installation complete.${N}"
echo
echo "  Normalizer:      $NORMALIZER_STATE"
echo
echo "  Start ComfyUI:   $LAUNCHER"
echo "  Web UI:          http://127.0.0.1:$COMFY_PORT"
echo "  Then run:        python3 lektor_ai_comfy.py <video.mkv> --embedded-subs"
echo
echo "  lektor_ai_comfy.py expects these paths — edit its constants if you changed them:"
echo "    COMFY_URL         = \"http://127.0.0.1:$COMFY_PORT\""
echo "    COMFY_OUTPUT_DIR  = \"$COMFY_DIR/output\""
echo "    COMFY_INPUT_DIR   = \"$COMFY_DIR/input\""
echo
echo "  Record this environment for later:"
echo "    source $CONDA_ROOT/etc/profile.d/conda.sh && conda activate $ENV_NAME"
echo "    pip freeze > comfyui-requirements.lock.txt"
echo
echo "  Remove everything:"
echo "    conda env remove -n $ENV_NAME && rm -rf $COMFY_ROOT"
echo "    rm -rf $NORMALIZER_DIR/venv $NORMALIZER_DIR/model $NORMALIZER_DIR/models"
echo "    rm -rf $CONDA_ROOT   # only if nothing else uses it"

if [[ "$NORMALIZER_STATE" == fallback* ]]; then
    echo
    echo "  ${Y}One model still needs you:${N} it is gated behind a licence."
    echo "    1. Open https://huggingface.co/Folx/qwen3-0.6b-pl-text-normalization"
    echo "       while signed in, and accept the terms"
    echo "    2. $NORMALIZER_DIR/venv/bin/hf auth login"
    echo "    3. $NORMALIZER_DIR/setup.sh --model"
    echo "  Until then times, dates and identifiers fall back to num2words."
fi