#!/bin/bash
# AI Profile Router – Installation/Update auf dem Zielsystem.
# Wird als root auf dem Zielsystem ausgeführt (per SSH, vgl. deploy.sh).
# Erwartet ai_profile_router.py, image_worker.py, tts_worker.py,
# mike-ai-profile-router.service und mike-ai-kokoro.service im selben
# Verzeichnis wie dieses Skript.
set -euo pipefail

DIR="$(cd "$(dirname "$0")" && pwd)"
INSTALL_DIR=/opt/mike-ai/ai-profile-router
SERVICE=mike-ai-profile-router.service
KOKORO_SERVICE=mike-ai-kokoro.service
OLD_SERVICE=mike-ai-local-llm-router.service
OLD_DIR=/opt/mike-ai/local-llm-router
BACKUP_DIR=/opt/mike-ai/.backup-ai-profile-router-$(date +%Y%m%d-%H%M%S)
VENV="$INSTALL_DIR/venv"
MODEL_DIR=/opt/mike-ai/models/FLUX.2-klein-base-4B
IMAGE_DIR="$INSTALL_DIR/images"
KOKORO_VENV=/opt/mike-ai/kokoro/venv
KOKORO_MODEL_DIR=/opt/mike-ai/models/kokoro

echo "== AI Profile Router: Installation/Update =="

# --- 1. Bestehenden (alten) Router entfernen, mit Backup -------------------
if [ -f "/etc/systemd/system/${OLD_SERVICE}" ]; then
  echo "-- Entferne alten Router ${OLD_SERVICE} (Backup: ${BACKUP_DIR})"
  mkdir -p "$BACKUP_DIR"
  systemctl stop "$OLD_SERVICE" 2>/dev/null || true
  systemctl disable "$OLD_SERVICE" 2>/dev/null || true
  mv "/etc/systemd/system/${OLD_SERVICE}" "$BACKUP_DIR/"
  if [ -d "$OLD_DIR" ]; then
    mv "$OLD_DIR" "$BACKUP_DIR/local-llm-router"
  fi
else
  echo "-- Kein alter Router gefunden"
fi

# --- 2. Neue Dateien installieren -------------------------------------------
mkdir -p "$INSTALL_DIR" "$IMAGE_DIR"
install -m 0755 "$DIR/ai_profile_router.py" "$INSTALL_DIR/ai_profile_router.py"
install -m 0755 "$DIR/image_worker.py" "$INSTALL_DIR/image_worker.py"
install -m 0755 "$DIR/tts_worker.py" "$INSTALL_DIR/tts_worker.py"
install -m 0644 "$DIR/${SERVICE}" "/etc/systemd/system/${SERVICE}"
install -m 0644 "$DIR/${KOKORO_SERVICE}" "/etc/systemd/system/${KOKORO_SERVICE}"

# --- 3. Python-Venv mit Bild-Abhängigkeiten ---------------------------------
if [ ! -x "$VENV/bin/python" ]; then
  echo "-- Erstelle Python-Venv in $VENV"
  python3 -m venv "$VENV"
fi
echo "-- Installiere/aktualisiere Bild-Abhängigkeiten (torch, diffusers, ...)"
"$VENV/bin/pip" install --quiet --upgrade pip
"$VENV/bin/pip" install --quiet \
  torch \
  diffusers \
  transformers \
  accelerate

# --- 4. FLUX-Modell (nur wenn noch nicht vorhanden) --------------------------
if [ -f "$MODEL_DIR/model_index.json" ]; then
  echo "-- FLUX-Modell vorhanden: $MODEL_DIR"
else
  echo "-- Lade FLUX.2-klein-base-4B nach $MODEL_DIR (kann dauern)"
  "$VENV/bin/python" - <<'PY'
import os
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id="black-forest-labs/FLUX.2-klein-base-4B",
    local_dir="/opt/mike-ai/models/FLUX.2-klein-base-4B",
    local_dir_use_symlinks=False,
)
print("Modell-Download abgeschlossen")
PY
fi

# --- 5. TTS (Kokoro) ----------------------------------------------------------
# espeak-ng wird für die deutsche G2P (phonemizer) benötigt.
if ! command -v espeak-ng >/dev/null 2>&1; then
  echo "-- Installiere espeak-ng"
  apt-get install -y espeak-ng
fi

# Kokoro-Venv (CPU-only torch + semidark/kokoro + semidark/misaki + soundfile + lameenc)
if [ ! -x "$KOKORO_VENV/bin/python" ]; then
  echo "-- Erstelle Kokoro-Venv in $KOKORO_VENV"
  mkdir -p "$KOKORO_VENV"
  python3 -m venv "$KOKORO_VENV"
  "$KOKORO_VENV/bin/pip" install --quiet --upgrade pip
  echo "-- Installiere CPU-only torch"
  "$KOKORO_VENV/bin/pip" install --quiet torch \
    --index-url https://download.pytorch.org/whl/cpu
  echo "-- Installiere Kokoro-Abhängigkeiten"
  "$KOKORO_VENV/bin/pip" install --quiet \
    huggingface-hub loguru transformers regex num2words \
    soundfile lameenc
  # Offizieller Kikiri-Deutsche-Pfad:
  # - semidark/misaki-Fork (0.9.4) mit [de]-Extra (phonemizer-fork + espeakng-loader)
  # - semidark/kokoro-Fork (0.9.4) mit --no-deps (vermeidet misaki[en] → spacy-curated-transformers)
  # - spacy + num2words für misaki.en-Import (Englisch, nicht für deutschen Pfad)
  "$KOKORO_VENV/bin/pip" install --quiet "spacy<4"
  "$KOKORO_VENV/bin/pip" install --quiet \
    "misaki[de] @ git+https://github.com/semidark/misaki.git"
  "$KOKORO_VENV/bin/pip" install --quiet \
    "kokoro @ git+https://github.com/semidark/kokoro.git" --no-deps
fi

# Kokoro-Modelle (kikiri-german-martin + kikiri-german-victoria, Apache 2.0)
if [ ! -f "$KOKORO_MODEL_DIR/kikiri-german-martin/kikiri_german_martin_ep10.pth" ] \
   || [ ! -f "$KOKORO_MODEL_DIR/kikiri-german-victoria/kikiri_german_victoria_ep10.pth" ]; then
  echo "-- Lade Kokoro-Modelle nach $KOKORO_MODEL_DIR (kann dauern)"
  "$KOKORO_VENV/bin/python" - <<'PY'
import os
from huggingface_hub import hf_hub_download

base = "/opt/mike-ai/models/kokoro"
files = [
    ("kikiri-tts/kikiri-german-martin", "kikiri_german_martin_ep10.pth", "kikiri-german-martin"),
    ("kikiri-tts/kikiri-german-martin", "voices/martin.pt", "kikiri-german-martin"),
    ("kikiri-tts/kikiri-german-victoria", "kikiri_german_victoria_ep10.pth", "kikiri-german-victoria"),
    ("kikiri-tts/kikiri-german-victoria", "voices/victoria.pt", "kikiri-german-victoria"),
    ("hexgrad/Kokoro-82M", "config.json", "kikiri-german-martin"),
    ("hexgrad/Kokoro-82M", "config.json", "kikiri-german-victoria"),
]
for repo, fname, subdir in files:
    dest = os.path.join(base, subdir, fname)
    if os.path.exists(dest) and os.path.getsize(dest) > 1000:
        print(f"vorhanden: {subdir}/{fname}")
        continue
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    hf_hub_download(repo_id=repo, filename=fname, local_dir=os.path.join(base, subdir))
    print(f"geladen: {subdir}/{fname}")
PY
fi

# --- 6. Services aktivieren und starten ---------------------------------------
systemctl daemon-reload
systemctl enable "$SERVICE" "$KOKORO_SERVICE"
systemctl restart "$KOKORO_SERVICE"
systemctl restart "$SERVICE"

# --- 7. Verifikation ----------------------------------------------------------
sleep 1
if ! systemctl is-active --quiet "$SERVICE"; then
  echo "-- FEHLER: Router-Service läuft nicht" >&2
  journalctl -u "$SERVICE" -n 20 --no-pager >&2
  exit 1
fi
echo "-- Router-Service läuft"
if ! systemctl is-active --quiet "$KOKORO_SERVICE"; then
  echo "-- FEHLER: Kokoro-Service läuft nicht" >&2
  journalctl -u "$KOKORO_SERVICE" -n 20 --no-pager >&2
  exit 1
fi
echo "-- Kokoro-Service läuft"
curl -sf "http://127.0.0.1:8081/status" | python3 -m json.tool
echo "== Fertig =="