FROM diegosouzapw/omniroute:latest

# 1. Switch to root to install system packages
USER root

# 2. Install Python for the backup script (Debian-based)
RUN apt-get update && \
    apt-get install -y --no-install-recommends python3 python3-pip tar openssl && \
    pip3 install --break-system-packages --no-cache-dir huggingface_hub && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /app

# 3. Create the Python sync script (FIXED: tarball mirrors real FS layout
#    + marker-based legacy healer + file-level verify)
RUN cat > /app/sync.py <<'PY'
import os
import shutil
import tarfile
import tempfile
import time
import sys
from pathlib import Path
from huggingface_hub import HfApi, hf_hub_download

REPO_ID = os.environ.get("DATASET_REPO")
TOKEN = os.environ.get("HF_TOKEN")
BACKUP_FILE = "/tmp/omniroute_backup.tar.gz"
DATA_DIR = Path("/app/data")
CONFIG_DIR = Path("/root/.omniroute")


def safe_extract(tar, path):
    """Extract `tar` into `path`, rejecting members that would escape it."""
    base = Path(path).resolve()
    for member in tar.getmembers():
        target = (base / member.name).resolve()
        if not str(target).startswith(str(base)):
            raise RuntimeError(f"Unsafe tar member: {member.name}")
    # Python 3.12+ requires the `filter` argument (and we WANT data filter,
    # which strips leading slashes / .. and refuses absolute paths).
    if sys.version_info >= (3, 12):
        tar.extractall(path, filter="data")
    else:
        tar.extractall(path)


def heal_legacy_layout():
    """
    Older backups were packed as `./data/...` and `./.omniroute/...`, which the
    old code extracted to `/data` and `/.omniroute` instead of `/app/data` and
    `/root/.omniroute`.
    IMPORTANT: run.sh `mkdir -p`s both target dirs BEFORE we run, so we cannot
    just check `dst.exists()` — the target is always there (empty). We must
    check the actual *file* the app needs:
      - /app/data/storage.sqlite    (target)  vs  /data/storage.sqlite    (misplaced)
      - /root/.omniroute/server.env (target)  vs  /.omniroute/server.env  (misplaced)
    If the misplaced copy exists and the target marker is missing, we swap.
    If the target is already populated, we REFUSE to clobber it.
    """
    plans = [
        (Path("/data"),       "storage.sqlite", Path("/app/data")),
        (Path("/.omniroute"), "server.env",     Path("/root/.omniroute")),
    ]
    for src, marker_name, dst in plans:
        src_marker = src / marker_name
        dst_marker = dst / marker_name

        if not src_marker.exists():
            continue  # nothing misplaced for this pair

        if dst_marker.exists():
            print(
                f">>> [Sync] Healer: target {dst_marker} already present, "
                f"leaving {src} alone (would not clobber real data)",
                flush=True,
            )
            continue

        print(f">>> [Sync] Healer: misplaced legacy content at {src}", flush=True)

        # Target dir exists but is missing its marker -> it must be empty.
        # Try to remove it so shutil.move produces a clean swap; if it's not
        # empty (unexpected), bail out instead of nesting src inside dst.
        if dst.exists():
            try:
                dst.rmdir()
                print(f">>> [Sync] Healer: removed empty target {dst}", flush=True)
            except OSError:
                print(
                    f">>> [Sync] Healer: ABORT — {dst} is not empty, "
                    f"refusing to clobber",
                    flush=True,
                )
                continue

        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(str(src), str(dst))
        print(f">>> [Sync] Healer: relocated {src} -> {dst}", flush=True)


def upload():
    if not REPO_ID or not TOKEN:
        print(">>> [Sync] DATASET_REPO or HF_TOKEN not configured. Skipping backup.")
        return

    with tempfile.TemporaryDirectory() as tmpdir:
        tmp_path = Path(tmpdir)

        # Mirror the REAL filesystem layout. After upload, the tar contains:
        #   ./app/data/...
        #   ./root/.omniroute/...
        # so extracting to "/" puts files back where the app expects them.
        if DATA_DIR.exists():
            target = tmp_path / "app" / "data"
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copytree(DATA_DIR, target)
            print(f">>> [Sync] Packed {DATA_DIR} -> app/data", flush=True)

        if CONFIG_DIR.exists():
            target = tmp_path / "root" / ".omniroute"
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copytree(CONFIG_DIR, target)
            print(f">>> [Sync] Packed {CONFIG_DIR} -> root/.omniroute", flush=True)

        try:
            with tarfile.open(BACKUP_FILE, "w:gz") as tar:
                tar.add(str(tmp_path), arcname=".")
            HfApi(token=TOKEN).upload_file(
                path_or_fileobj=BACKUP_FILE,
                path_in_repo="omniroute_backup.tar.gz",
                repo_id=REPO_ID,
                repo_type="dataset",
            )
            print(f">>> [Sync] Backup successful: {time.strftime('%Y-%m-%d %H:%M:%S')}", flush=True)
        except Exception as e:
            print(f">>> [Sync] Backup failed: {e}", flush=True)


def download():
    if not REPO_ID or not TOKEN:
        print(">>> [Sync] DATASET_REPO or HF_TOKEN not configured. Skipping restore.")
        return
    try:
        print(">>> [Sync] Downloading data from Dataset...", flush=True)
        path = hf_hub_download(
            repo_id=REPO_ID,
            filename="omniroute_backup.tar.gz",
            repo_type="dataset",
            token=TOKEN,
        )
        print(f">>> [Sync] Downloaded archive to {path}", flush=True)

        with tarfile.open(path, "r:gz") as tar:
            members = tar.getmembers()
            print(f">>> [Sync] Extracting {len(members)} members to / ...", flush=True)
            safe_extract(tar, "/")

        # Self-heal: handle backups created before the path fix.
        heal_legacy_layout()

        # Verify the app can actually see the restored files.
        # Check the REAL files, not just the dirs (mkdir -p makes them lie).
        critical = [
            (DATA_DIR,   "storage.sqlite"),
            (CONFIG_DIR, "server.env"),
        ]
        for d, fname in critical:
            marker = d / fname
            status = "OK" if marker.exists() else "MISSING"
            print(f">>> [Sync] Verify {marker}: {status}", flush=True)

        print(">>> [Sync] Data restore successful.", flush=True)
    except Exception as e:
        print(f">>> [Sync] Skipping restore (might be first run): {e}", flush=True)


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "upload":
        upload()
    elif len(sys.argv) > 1 and sys.argv[1] == "download":
        download()
PY

# 4. Create the startup shell script (FIXED: don't clobber restored server.env)
RUN cat > /app/run.sh <<'SH'
#!/bin/sh
set -e

mkdir -p /app/data
mkdir -p /root/.omniroute

echo '=== 1. Restoring Data & Config ==='
python3 /app/sync.py download

echo '=== 2. Starting Scheduled Backup (Every 10 minutes) ==='
(while true; do sleep 600; python3 /app/sync.py upload; done) &

echo '=== 3. Generating & Persisting Secrets ==='
# Only mint fresh secrets if restore didn't bring a server.env with us.
# Otherwise we'd rotate JWT_SECRET etc. on every container restart and
# invalidate every active session / API key / initial-password.
if [ ! -f /root/.omniroute/server.env ]; then
  echo '>>> No restored server.env found, generating fresh secrets'
  export JWT_SECRET="${JWT_SECRET:-$(openssl rand -base64 48 | tr -d '\n')}"
  export API_KEY_SECRET="${API_KEY_SECRET:-$(openssl rand -hex 32 | tr -d '\n')}"
  export STORAGE_ENCRYPTION_KEY="${STORAGE_ENCRYPTION_KEY:-$(openssl rand -hex 32 | tr -d '\n')}"
  export OMNIROUTE_WS_BRIDGE_SECRET="${OMNIROUTE_WS_BRIDGE_SECRET:-$(openssl rand -base64 32 | tr -d '\n')}"
  export INITIAL_PASSWORD="${INITIAL_PASSWORD:-$(openssl rand -hex 16 | tr -d '\n')}"

  cat > /root/.omniroute/server.env <<EOF
JWT_SECRET=${JWT_SECRET}
API_KEY_SECRET=${API_KEY_SECRET}
STORAGE_ENCRYPTION_KEY=${STORAGE_ENCRYPTION_KEY}
OMNIROUTE_WS_BRIDGE_SECRET=${OMNIROUTE_WS_BRIDGE_SECRET}
INITIAL_PASSWORD=${INITIAL_PASSWORD}
DATA_DIR=/app/data
PORT=7860
NODE_ENV=production
NEXT_PUBLIC_BASE_URL=https://${SPACE_HOST:-localhost:7860}
BASE_URL=http://localhost:7860
AUTH_COOKIE_SECURE=true
REQUIRE_API_KEY=true
CORS_ORIGIN=*
OMNIROUTE_MEMORY_MB=14745
PRICING_SYNC_ENABLED=true
OMNIROUTE_AUTO_FREE_FALLBACK_TO_FULL_POOL=true
OMNIROUTE_SHOW_LOG=0
APP_LOG_LEVEL=error
EOF
else
  echo '>>> Restored server.env found, keeping restored secrets'
fi

# Source whichever server.env we ended up with (restored or freshly minted).
set -a
. "/root/.omniroute/server.env"
set +a

export NODE_OPTIONS="--max-old-space-size=${OMNIROUTE_MEMORY_MB}"

echo '=== 4. Starting OmniRoute ==='
exec node server.js
SH

RUN chmod +x /app/run.sh

EXPOSE 7860
CMD ["/bin/sh", "/app/run.sh"]