#!/bin/sh
set -eu

PYTHON=/opt/venv/bin/python3
DATA_DIR="${DATA_DIR:-/app/data}"
CONFIG="$DATA_DIR/config.toml"

echo "[start] using python: $PYTHON"
echo "[start] data dir: $DATA_DIR"

# 1. 从 HF Bucket 恢复 data/
$PYTHON - <<'PYEOF'
import os
import sys
from huggingface_hub import HfFileSystem

token = os.environ.get("HF_TOKEN", "")
repo = os.environ.get("HF_BUCKET_REPO", "")
local_data = os.environ.get("DATA_DIR", "/app/data")


def to_relative_path(remote_file: str, repo_name: str) -> str:
    prefixes = (
        f"hf://buckets/{repo_name}/",
        f"buckets/{repo_name}/",
        f"/{repo_name}/",
        f"{repo_name}/",
    )
    rel_path = remote_file
    while True:
        for prefix in prefixes:
            if rel_path.startswith(prefix):
                rel_path = rel_path[len(prefix):]
                break
        else:
            break
    return rel_path.lstrip("/")

if not token or not repo:
    print("[bucket_pull] skipped: HF_TOKEN or HF_BUCKET_REPO not set")
    sys.exit(0)

fs = HfFileSystem(token=token)
bucket_root = f"hf://buckets/{repo}"
os.makedirs(local_data, exist_ok=True)

try:
    if not fs.exists(bucket_root):
        print(f"[bucket_pull] bucket {bucket_root} not found, skipping")
        sys.exit(0)

    files = fs.glob(f"{bucket_root}/**/*")
    pulled = 0
    for remote_file in files:
        if fs.isfile(remote_file):
            rel_path = to_relative_path(remote_file, repo)
            if not rel_path or rel_path.startswith("buckets/"):
                print(f"[bucket_pull] skipped invalid path: {remote_file}")
                continue
            local_file = os.path.join(local_data, rel_path)
            os.makedirs(os.path.dirname(local_file), exist_ok=True)
            fs.get(remote_file, local_file)
            print(f"[bucket_pull] restored: {rel_path}")
            pulled += 1
    print(f"[bucket_pull] done, {pulled} file(s) restored.")
except Exception as exc:
    print(f"[bucket_pull] error: {exc}")
PYEOF

# 2. 初始化存储
if [ -f /app/scripts/init_storage.sh ]; then
    /app/scripts/init_storage.sh
fi

# 3. 注入 Secrets 到 config.toml
if [ -n "${APP_KEY:-}" ] && [ -f "$CONFIG" ]; then
    sed -i "s|^app_key = .*|app_key = \"$APP_KEY\"|" "$CONFIG"
fi
if [ -n "${API_KEY:-}" ] && [ -f "$CONFIG" ]; then
    sed -i "s|^api_key = .*|api_key = \"$API_KEY\"|" "$CONFIG"
fi

export SERVER_STORAGE_TYPE="${SERVER_STORAGE_TYPE:-local}"
export SERVER_STORAGE_URL="${SERVER_STORAGE_URL:-}"

# 4. 后台：文件监听（防抖推送）+ 兜底定时同步
$PYTHON - <<'PYEOF' &
import glob
import os
import threading
import time
from huggingface_hub import HfFileSystem
from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer

token = os.environ.get("HF_TOKEN", "")
repo = os.environ.get("HF_BUCKET_REPO", "")
data_dir = os.environ.get("DATA_DIR", "/app/data")
bucket_root = f"hf://buckets/{repo}"


def should_skip_local_path(rel_path: str) -> bool:
    rel_path = rel_path.replace("\\", "/")
    return not rel_path or rel_path.startswith("buckets/")


def push():
    if not token or not repo:
        return

    fs = HfFileSystem(token=token)
    files = [
        path
        for path in glob.glob(os.path.join(data_dir, "**", "*"), recursive=True)
        if os.path.isfile(path)
    ]
    if not files:
        return

    try:
        for local_path in files:
            rel_path = os.path.relpath(local_path, data_dir).replace("\\", "/")
            if should_skip_local_path(rel_path):
                continue
            remote_path = f"{bucket_root}/{rel_path}"
            fs.put(local_path, remote_path)
        print(f"[bucket_push] pushed {len(files)} file(s)", flush=True)
    except Exception as exc:
        print(f"[bucket_push] error: {exc}", flush=True)


class Handler(FileSystemEventHandler):
    def __init__(self):
        self._timer = None
        self._lock = threading.Lock()

    def _schedule(self):
        with self._lock:
            if self._timer:
                self._timer.cancel()
            self._timer = threading.Timer(5, push)
            self._timer.start()

    def on_modified(self, event):
        if not event.is_directory:
            self._schedule()

    def on_created(self, event):
        if not event.is_directory:
            self._schedule()

    def on_deleted(self, event):
        if not event.is_directory:
            self._schedule()

    def on_moved(self, event):
        if not event.is_directory:
            self._schedule()


if token and repo:
    os.makedirs(data_dir, exist_ok=True)
    observer = Observer()
    observer.schedule(Handler(), path=data_dir, recursive=True)
    observer.start()
    print(f"[bucket_watch] watching {data_dir}...", flush=True)

    def periodic():
        while True:
            time.sleep(300)
            push()

    threading.Thread(target=periodic, daemon=True).start()

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        observer.stop()
    observer.join()
else:
    print("[bucket_watch] skipped: HF_TOKEN or HF_BUCKET_REPO not set", flush=True)
    while True:
        time.sleep(3600)
PYEOF

# 5. 启动应用
echo "[start] starting grok2api..."
exec /opt/venv/bin/granian --interface asgi \
    --host "${SERVER_HOST:-0.0.0.0}" \
    --port "${SERVER_PORT:-7860}" \
    --workers "${SERVER_WORKERS:-1}" \
    main:app