#!/bin/sh
set -e

to_bool() {
  case "$1" in
    1|true|TRUE|yes|YES|y|Y) echo "true" ;;
    *) echo "false" ;;
  esac
}

mkdir -p /tmp/nginx/client_body /tmp/nginx/proxy /tmp/nginx/fastcgi /tmp/nginx/uwsgi /tmp/nginx/scgi

cat > /tmp/nginx.conf <<EOF
user node;
worker_processes auto;
pid /tmp/nginx.pid;
error_log /dev/stderr info;

events {
  worker_connections 1024;
}

http {
  include       /etc/nginx/mime.types;
  default_type  application/octet-stream;
  access_log /dev/stdout;
  sendfile on;
  keepalive_timeout 65;
  gunzip on;

  # Strip any :port from Host for upstream and redirects
  map \$http_host \$clean_host {
    "~^(?<h>[^:]+)(?::\\d+)?\$" \$h;
    default \$host;
  }

  client_body_temp_path /tmp/nginx/client_body;
  proxy_temp_path       /tmp/nginx/proxy;
  fastcgi_temp_path     /tmp/nginx/fastcgi;
  uwsgi_temp_path       /tmp/nginx/uwsgi;
  scgi_temp_path        /tmp/nginx/scgi;

  include /tmp/nginx_sites.conf;
}
EOF

cat > /tmp/nginx_sites.conf.template <<'EOF'
server {
  listen 8000;
  server_name _;
  port_in_redirect off;
  absolute_redirect off;

  location / {
    proxy_pass http://127.0.0.1:8001/;
    proxy_http_version 1.1;
    proxy_set_header Host $clean_host;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection "upgrade";
    # Force HTTPS in upstream to avoid mixed-content absolute URLs
    proxy_set_header X-Forwarded-Proto https;
    proxy_set_header X-Forwarded-Port 443;
    proxy_set_header X-Forwarded-Ssl on;
    proxy_set_header X-Forwarded-Host $clean_host;
    proxy_set_header Forwarded "proto=https;host=$host";
    proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
    # Rewrite upstream absolute redirects that include :8000
    proxy_redirect ~^http://([^:]+):8000/(.*)$ https://$1/$2;
    proxy_redirect ~^https?://([^:]+):8000/(.*)$ https://$1/$2;
  }
}
EOF
cp /tmp/nginx_sites.conf.template /tmp/nginx_sites.conf

python3 - <<'PY'
import os
import tarfile
from huggingface_hub import hf_hub_download

repo_id = os.environ.get("HF_BACKUP_REPO")
token = os.environ.get("HF_TOKEN")
if not repo_id or not token:
    print("[restore] missing HF_BACKUP_REPO or HF_TOKEN", flush=True)
    raise SystemExit(0)

try:
    path = hf_hub_download(
        repo_id=repo_id,
        repo_type="dataset",
        filename="latest.tgz",
        token=token,
    )
except Exception as e:
    print(f"[restore] download failed: {e}", flush=True)
    raise SystemExit(0)

try:
    with tarfile.open(path, "r:gz") as tar:
        tar.extractall("/home/node/app")
    print("[restore] done", flush=True)
except Exception as e:
    print(f"[restore] extract failed: {e}", flush=True)
    raise SystemExit(0)
PY

ENABLE_USER_ACCOUNTS="$(to_bool "${ST_MULTI_USER:-true}")"
ENABLE_DISCREET_LOGIN="$(to_bool "${ST_DISCREET_LOGIN:-true}")"
ENABLE_WHITELIST="$(to_bool "${ST_WHITELIST_MODE:-true}")"
ENABLE_FWD_WHITELIST="$(to_bool "${ST_FWD_WHITELIST:-false}")"
if [ "$ENABLE_WHITELIST" = "false" ]; then
  ENABLE_FWD_WHITELIST="false"
fi

mkdir -p /home/node/app/config
CONFIG_TMP="/tmp/st-config.yaml"
cat > "$CONFIG_TMP" <<EOF
listen: true
port: 8001
listenAddress: 0.0.0.0
enableUserAccounts: ${ENABLE_USER_ACCOUNTS}
enableDiscreetLogin: ${ENABLE_DISCREET_LOGIN}
whitelistMode: ${ENABLE_WHITELIST}
enableForwardedWhitelist: ${ENABLE_FWD_WHITELIST}
whitelist:
  - 127.0.0.1
  - ::1
  - 0.0.0.0/0
EOF
cp -f "$CONFIG_TMP" /home/node/app/config.yaml
cp -f "$CONFIG_TMP" /home/node/app/config/config.yaml
mkdir -p /home/node/app/data
cp -f "$CONFIG_TMP" /home/node/app/data/config.yaml
echo "[config] wrote /home/node/app/config.yaml, /home/node/app/config/config.yaml, /home/node/app/data/config.yaml" >&2

python3 - <<'PY'
import os
import signal
import subprocess
import sys
import threading
import time
import tarfile
from huggingface_hub import HfApi

BACKUP_EVERY = int(os.getenv("BACKUP_EVERY", "20000"))
BACKUP_MIN_INTERVAL_SEC = int(os.getenv("BACKUP_MIN_INTERVAL_SEC", "120"))
BACKUP_COOLDOWN_SEC = int(os.getenv("BACKUP_COOLDOWN_SEC", "3600"))
BACKUP_AUTOCLEAN = os.getenv("BACKUP_AUTOCLEAN", "true").strip().lower() in ("1", "true", "yes", "y")
BACKUP_CLEANUP_EVERY = int(os.getenv("BACKUP_CLEANUP_EVERY", "5"))
BACKUP_CLEANUP_MIN_INTERVAL_SEC = int(os.getenv("BACKUP_CLEANUP_MIN_INTERVAL_SEC", "21600"))
HEARTBEAT_SEC = int(os.getenv("HEARTBEAT_SEC", "30"))
if BACKUP_EVERY < 1:
    BACKUP_EVERY = 1000
if BACKUP_MIN_INTERVAL_SEC < 1:
    BACKUP_MIN_INTERVAL_SEC = 120
if BACKUP_COOLDOWN_SEC < 1:
    BACKUP_COOLDOWN_SEC = 3600
if BACKUP_CLEANUP_EVERY < 1:
    BACKUP_CLEANUP_EVERY = 5
if BACKUP_CLEANUP_MIN_INTERVAL_SEC < 1:
    BACKUP_CLEANUP_MIN_INTERVAL_SEC = 21600
if HEARTBEAT_SEC < 1:
    HEARTBEAT_SEC = 30

state = {
    "lines": 0,
    "backup_requested": False,
    "backup_running": False,
    "last_backup_ts": 0.0,
    "next_backup_ts": 0.0,
    "backups_since_cleanup": 0,
    "last_cleanup_ts": 0.0,
    "next_cleanup_ts": 0.0,
    "last_log_time": time.time(),
    "last_heartbeat_time": 0.0,
}

APP = "/home/node/app"
PATHS = [
    "config",
    "data",
    "plugins",
    "public/scripts/extensions/third-party",
]

repo_id = os.environ.get("HF_BACKUP_REPO")
token = os.environ.get("HF_TOKEN")


def emit_line(line: str) -> None:
    sys.stdout.write(line + "\n")
    sys.stdout.flush()
    state["last_log_time"] = time.time()
    state["lines"] += 1
    if state["lines"] >= BACKUP_EVERY:
        state["backup_requested"] = True
        state["lines"] = 0


def run_backup() -> None:
    if not repo_id or not token:
        emit_line("[backup] missing HF_BACKUP_REPO or HF_TOKEN")
        return

    try:
        tar_path = "/tmp/st-backup.tgz"
        with tarfile.open(tar_path, "w:gz") as tar:
            for p in PATHS:
                full = os.path.join(APP, p)
                if os.path.exists(full):
                    tar.add(full, arcname=p)

        api = HfApi()
        api.upload_file(
            path_or_fileobj=tar_path,
            path_in_repo="latest.tgz",
            repo_id=repo_id,
            repo_type="dataset",
            token=token,
            commit_message=f"backup {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}",
        )
        state["last_backup_ts"] = time.time()
        state["backups_since_cleanup"] += 1
        emit_line("[backup] uploaded latest.tgz")

        # Keep storage in check by periodically squashing history in the backup dataset.
        if BACKUP_AUTOCLEAN:
            now_ts = time.time()
            can_cleanup = (
                state["backups_since_cleanup"] >= BACKUP_CLEANUP_EVERY
                and now_ts >= state["next_cleanup_ts"]
                and (now_ts - state["last_cleanup_ts"]) >= BACKUP_CLEANUP_MIN_INTERVAL_SEC
            )
            if can_cleanup:
                try:
                    emit_line("[cleanup] start super_squash_history")
                    api.super_squash_history(
                        repo_id=repo_id,
                        repo_type="dataset",
                        token=token,
                        commit_message=f"cleanup {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}",
                    )
                    state["last_cleanup_ts"] = time.time()
                    state["backups_since_cleanup"] = 0
                    emit_line("[cleanup] done")
                except Exception as e:
                    msg2 = str(e)
                    if "429" in msg2 or "Too Many Requests" in msg2:
                        state["next_cleanup_ts"] = time.time() + BACKUP_COOLDOWN_SEC
                        emit_line(f"[cleanup] rate limited; cooldown {BACKUP_COOLDOWN_SEC}s")
                    else:
                        emit_line(f"[cleanup] failed: {msg2}")
    except Exception as e:
        msg = str(e)
        if "429" in msg or "Too Many Requests" in msg:
            state["next_backup_ts"] = time.time() + BACKUP_COOLDOWN_SEC
            emit_line(f"[backup] rate limited; cooldown {BACKUP_COOLDOWN_SEC}s")
        else:
            emit_line(f"[backup] failed: {msg}")


def reader(proc: subprocess.Popen, prefix: str | None = None) -> None:
    if not proc.stdout:
        return
    for line in proc.stdout:
        line = line.rstrip("\n")
        if prefix:
            emit_line(f"[{prefix}] {line}")
        else:
            emit_line(line)

nginx_proc = subprocess.Popen(
    ["nginx", "-c", "/tmp/nginx.conf", "-p", "/tmp/nginx", "-g", "daemon off;"],
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    text=True,
    bufsize=1,
)
threading.Thread(target=reader, args=(nginx_proc, "nginx"), daemon=True).start()

st_proc = subprocess.Popen(
    ["/home/node/app/docker-entrypoint.sh"],
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    text=True,
    bufsize=1,
)
threading.Thread(target=reader, args=(st_proc, None), daemon=True).start()


def shutdown(signum, frame):
    try:
        if st_proc.poll() is None:
            st_proc.terminate()
        if nginx_proc.poll() is None:
            nginx_proc.terminate()
        st_proc.wait(timeout=10)
        nginx_proc.wait(timeout=10)
    except Exception:
        st_proc.kill()
        nginx_proc.kill()
    sys.exit(0)


signal.signal(signal.SIGTERM, shutdown)
signal.signal(signal.SIGINT, shutdown)

while True:
    now = time.time()

    if state["backup_requested"] and not state["backup_running"]:
        state["backup_running"] = True
        emit_line("[backup] start")
        now_ts = time.time()
        if now_ts < state["next_backup_ts"]:
            emit_line("[backup] skipped (cooldown)")
        elif now_ts - state["last_backup_ts"] < BACKUP_MIN_INTERVAL_SEC:
            emit_line("[backup] skipped (min interval)")
        else:
            run_backup()
        emit_line("[backup] done")
        state["backup_running"] = False

    if (
        now - state["last_log_time"] >= HEARTBEAT_SEC
        and now - state["last_heartbeat_time"] >= HEARTBEAT_SEC
    ):
        emit_line(f"[heartbeat] {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}")
        state["last_heartbeat_time"] = now

    if st_proc.poll() is not None:
        emit_line(f"[supervisor] sillytavern exited with {st_proc.returncode}")
        sys.exit(st_proc.returncode)
    if nginx_proc.poll() is not None:
        emit_line(f"[supervisor] nginx exited with {nginx_proc.returncode}")
        sys.exit(nginx_proc.returncode)

    time.sleep(1)
PY