startup.py

"""
startup.py — Backup/Restore + Path Patcher + Telegram Webhook Setup
====================================================================
ROOT CAUSE OF BOT FAILURE:
  HuggingFace Spaces BLOCKS outbound TCP to api.telegram.org.
  Long polling (infinity_polling) requires the bot to call Telegram's
  servers — this is blocked by HF firewall. So the bot always fails.

FIX — Webhook mode:
  Instead of bot calling Telegram (blocked), we tell Telegram to call
  OUR HF Space URL (allowed — inbound is open).
  Telegram POSTs updates to: https://<your-space>.hf.space/webhook/<token>
  Flask receives them and processes commands instantly.

SECRETS (HF Space → Settings → Variables and secrets):
  HF_TOKEN                 = hf_xxxxxxxxxxxxxxxxxxxx
  HF_BACKUP_REPO           = Jsns82882/ulp-backup
  BACKUP_INTERVAL_MINUTES  = 30
  SPACE_URL                = https://jsns82882-ulp-bot.hf.space   ← your HF space URL
"""

import os
import re
import sys
import time
import signal
import zipfile
import logging
import threading
import subprocess
from pathlib import Path
from datetime import datetime

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("startup")

# ── Config ────────────────────────────────────────────────────────────────────
HF_TOKEN            = os.environ.get("HF_TOKEN", "")
HF_BACKUP_REPO      = os.environ.get("HF_BACKUP_REPO", "Jsns82882/ulp-backup")
BACKUP_INTERVAL_MIN = int(os.environ.get("BACKUP_INTERVAL_MINUTES", "30"))
# Your HF Space public URL — set this in Secrets!
# Format: https://jsns82882-ulp-bot.hf.space  (no trailing slash)
SPACE_URL           = os.environ.get("SPACE_URL", "").rstrip("/")
DATA_DIR            = Path("/app/data")
UPLOADS_DIR         = DATA_DIR / "uploads"
APP_FILE            = Path("/app/app.py")
BACKUP_ZIP          = Path("/tmp/ulp_data_backup.zip")

# ── Path patcher ──────────────────────────────────────────────────────────────
def patch_app_paths():
    """
    1. Redirect uploads + DB to /app/data/ (persistent across restarts)
    2. Replace infinity_polling with webhook receiver route in Flask
    3. Set WEB_BASE_URL so bot links work correctly
    """
    log.info("Patching app.py ...")
    src = APP_FILE.read_text()

    patches = [
        # Persistent paths
        (
            r"app\.config\[.UPLOAD_FOLDER.\]\s*=\s*['\"]uploads['\"]",
            "app.config['UPLOAD_FOLDER'] = '/app/data/uploads'"
        ),
        (
            r'DB_PATH\s*=\s*["\']bot_database\.db["\']',
            'DB_PATH = "/app/data/bot_database.db"'
        ),
        (
            r"os\.path\.exists\(['\"]uploads['\"]\)",
            "os.path.exists('/app/data/uploads')"
        ),
        (
            r"os\.makedirs\(['\"]uploads['\"]\)",
            "os.makedirs('/app/data/uploads', exist_ok=True)"
        ),
        # PORT default
        (
            r"port\s*=\s*int\(os\.environ\.get\(['\"]PORT['\"],\s*5000\)\)",
            "port = int(os.environ.get('PORT', 7860))"
        ),
        # Disable built-in bot startup (we add webhook route instead)
        (
            r"if should_start_bot\(\):\s*\n\s*start_telegram_bot\(\)",
            "# Bot uses webhook — managed by startup.py"
        ),
    ]

    changed = False
    for pattern, replacement in patches:
        new_src, count = re.subn(pattern, replacement, src)
        if count:
            log.info(f"  Patched: {pattern[:55]}")
            src = new_src
            changed = True

    # ── Inject webhook route into app.py ──────────────────────────────────────
    # Only inject if not already there
    if "webhook_receiver" not in src:
        webhook_code = '''

# ── Telegram Webhook receiver (injected by startup.py) ──────────────────────
import json as _json

@app.route(f'/webhook/{BOT_TOKEN}', methods=['POST'])
def webhook_receiver():
    """
    Telegram calls this URL with every update.
    We pass the raw JSON update directly into telebot for processing.
    This replaces long-polling which is blocked on HuggingFace.
    """
    try:
        update_data = request.get_data(as_text=True)
        update = telebot.types.Update.de_json(update_data)
        telegram_bot.process_new_updates([update])
    except Exception as exc:
        logger.error(f"Webhook error: {exc}")
    return 'OK', 200

@app.route('/webhook_status')
def webhook_status():
    """Quick health check to confirm webhook is registered."""
    return {'status': 'webhook active', 'token_suffix': BOT_TOKEN[-6:]}
# ── End webhook injection ────────────────────────────────────────────────────
'''
        # Insert before the if __name__ == '__main__' block
        if "if __name__ == '__main__':" in src:
            src = src.replace(
                "if __name__ == '__main__':",
                webhook_code + "\nif __name__ == '__main__':"
            )
        else:
            src += webhook_code

        log.info("  Injected webhook route into app.py.")
        changed = True

    if changed:
        APP_FILE.write_text(src)
        log.info("app.py patched successfully.")
    else:
        log.info("app.py already patched.")


# ── Webhook registration ──────────────────────────────────────────────────────
def register_webhook():
    """
    Tell Telegram to send updates to our HF Space URL.
    Called after Flask is confirmed running.
    Uses requests directly (no telebot) to avoid the blocked outbound issue
    — actually requests to api.telegram.org IS needed here too, but this
    only needs to succeed ONCE (Telegram remembers the webhook).
    We retry with backoff until it works.
    """
    if not SPACE_URL:
        log.warning("SPACE_URL not set — webhook not registered. Set it in HF Secrets!")
        log.warning("Format: https://jsns82882-ulp-bot.hf.space")
        return

    # Read BOT_TOKEN from app.py environment
    bot_token = os.environ.get("BOT_TOKEN", "8117597210:AAFBq7mAU6M5sSDP7r2GFC0uMys5JUxS4Qk")
    webhook_url = f"{SPACE_URL}/webhook/{bot_token}"

    log.info(f"Registering Telegram webhook: {webhook_url}")

    import urllib.request
    api_url = f"https://api.telegram.org/bot{bot_token}/setWebhook"
    payload = f"url={webhook_url}&drop_pending_updates=true".encode()

    for attempt in range(1, 6):
        try:
            req = urllib.request.Request(
                api_url,
                data=payload,
                headers={"Content-Type": "application/x-www-form-urlencoded"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=15) as resp:
                result = _json_loads(resp.read().decode())
                if result.get("ok"):
                    log.info(f"Webhook registered successfully: {webhook_url}")
                    return
                else:
                    log.warning(f"Webhook registration failed: {result}")
        except Exception as e:
            log.warning(f"Webhook attempt {attempt}/5 failed: {e}")
        time.sleep(attempt * 5)

    log.error("Could not register webhook after 5 attempts.")


def _json_loads(s):
    import json
    try:
        return json.loads(s)
    except Exception:
        return {}


# ── HuggingFace backup helpers ────────────────────────────────────────────────
def get_hf_api():
    if not HF_TOKEN or not HF_BACKUP_REPO:
        log.warning("HF_TOKEN or HF_BACKUP_REPO not set — backup/restore disabled.")
        return None
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=HF_TOKEN)
        try:
            api.repo_info(repo_id=HF_BACKUP_REPO, repo_type="dataset")
        except Exception:
            log.info(f"Creating backup dataset: {HF_BACKUP_REPO}")
            api.create_repo(repo_id=HF_BACKUP_REPO, repo_type="dataset", private=False)
        return api
    except Exception as e:
        log.error(f"HF API init failed: {e}")
        return None


def zip_data() -> bool:
    try:
        data_files = [f for f in DATA_DIR.rglob("*") if f.is_file()]
        if not data_files:
            log.info("No data to back up yet.")
            return False
        with zipfile.ZipFile(BACKUP_ZIP, "w", zipfile.ZIP_DEFLATED) as zf:
            for f in data_files:
                zf.write(f, f.relative_to(DATA_DIR))
        size_kb = BACKUP_ZIP.stat().st_size // 1024
        log.info(f"Zipped {len(data_files)} files ({size_kb} KB)")
        return True
    except Exception as e:
        log.error(f"Zip failed: {e}")
        return False


def upload_backup(api):
    try:
        timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        log.info(f"Uploading backup to {HF_BACKUP_REPO} ...")
        api.upload_file(
            path_or_fileobj=str(BACKUP_ZIP),
            path_in_repo="latest/ulp_data_backup.zip",
            repo_id=HF_BACKUP_REPO,
            repo_type="dataset",
            commit_message=f"Auto-backup {timestamp}",
        )
        api.upload_file(
            path_or_fileobj=str(BACKUP_ZIP),
            path_in_repo=f"history/{timestamp}_backup.zip",
            repo_id=HF_BACKUP_REPO,
            repo_type="dataset",
            commit_message=f"History {timestamp}",
        )
        log.info("Backup uploaded.")
        _cleanup_old_backups(api)
    except Exception as e:
        log.error(f"Upload failed: {e}")


def _cleanup_old_backups(api, keep=5):
    try:
        from huggingface_hub import list_repo_files
        files = sorted([
            f for f in list_repo_files(HF_BACKUP_REPO, repo_type="dataset")
            if f.startswith("history/") and f.endswith(".zip")
        ])
        for f in files[:-keep]:
            api.delete_file(path_in_repo=f, repo_id=HF_BACKUP_REPO, repo_type="dataset")
    except Exception as e:
        log.warning(f"Cleanup warning: {e}")


def do_backup(api):
    log.info("Backup starting ...")
    if zip_data():
        upload_backup(api)
    log.info("Backup done.")


def restore_from_backup(api):
    try:
        log.info("Looking for existing backup ...")
        from huggingface_hub import hf_hub_download
        local = hf_hub_download(
            repo_id=HF_BACKUP_REPO,
            filename="latest/ulp_data_backup.zip",
            repo_type="dataset",
            token=HF_TOKEN,
        )
        log.info(f"Restoring from {local} ...")
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(local, "r") as zf:
            zf.extractall(DATA_DIR)
        log.info("Data restored successfully.")
    except Exception as e:
        log.info(f"No backup found, starting fresh: {e}")
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        UPLOADS_DIR.mkdir(parents=True, exist_ok=True)


def backup_scheduler(api, interval_minutes):
    log.info(f"Backup scheduler: every {interval_minutes} min.")
    while True:
        time.sleep(interval_minutes * 60)
        try:
            do_backup(api)
        except Exception as e:
            log.error(f"Scheduled backup error: {e}")


# ── Flask app process ─────────────────────────────────────────────────────────
app_process = None

def start_app():
    global app_process
    env = os.environ.copy()
    env["PORT"] = env.get("PORT", "7860")
    env["DISABLE_TELEBOT"] = "1"
    if SPACE_URL:
        env["WEB_BASE_URL"] = SPACE_URL
    log.info(f"Starting Flask on port {env['PORT']} ...")
    app_process = subprocess.Popen(
        [sys.executable, "app.py"],
        cwd="/app",
        env=env,
    )
    log.info(f"Flask started (PID {app_process.pid})")


def wait_for_flask(timeout=30):
    """Wait until Flask is accepting connections before registering webhook."""
    import socket
    log.info("Waiting for Flask to be ready ...")
    for _ in range(timeout):
        try:
            s = socket.create_connection(("127.0.0.1", 7860), timeout=1)
            s.close()
            log.info("Flask is ready.")
            return True
        except Exception:
            time.sleep(1)
    log.warning("Flask did not become ready in time.")
    return False


def graceful_shutdown(signum, frame):
    log.info("Shutdown — running final backup ...")
    api = get_hf_api()
    if api:
        do_backup(api)
    if app_process:
        app_process.terminate()
        try:
            app_process.wait(timeout=10)
        except Exception:
            app_process.kill()
    sys.exit(0)


# ── Main ──────────────────────────────────────────────────────────────────────
def main():
    log.info("=" * 55)
    log.info("  ULP Startup — Patch + Restore + Flask + Webhook Bot")
    log.info("=" * 55)

    # 1. Patch app.py
    patch_app_paths()

    # 2. Ensure data dirs
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    UPLOADS_DIR.mkdir(parents=True, exist_ok=True)

    # 3. HF API + restore
    api = get_hf_api()
    if api:
        restore_from_backup(api)

    # 4. Shutdown hooks
    signal.signal(signal.SIGTERM, graceful_shutdown)
    signal.signal(signal.SIGINT, graceful_shutdown)

    # 5. Start Flask
    start_app()

    # 6. Wait for Flask, then register Telegram webhook
    def register_webhook_when_ready():
        if wait_for_flask():
            time.sleep(2)  # small buffer
            register_webhook()

    wt = threading.Thread(target=register_webhook_when_ready, daemon=True)
    wt.start()

    # 7. Backup scheduler
    if api:
        bt = threading.Thread(
            target=backup_scheduler,
            args=(api, BACKUP_INTERVAL_MIN),
            daemon=True,
        )
        bt.start()
        log.info("Backup scheduler started.")

    # 8. Watch Flask — restart if it crashes
    while True:
        ret = app_process.wait()
        log.warning(f"Flask exited (code {ret}) — restarting in 5s ...")
        time.sleep(5)
        start_app()
        # Re-register webhook after restart
        wt2 = threading.Thread(target=register_webhook_when_ready, daemon=True)
        wt2.start()


if __name__ == "__main__":
    main()