app/main.py

import os
import json
import time
import secrets
import threading
from datetime import datetime, timedelta

import psycopg
import gradio as gr

import backup
from fastapi import FastAPI, Depends, HTTPException, status
from fastapi.responses import JSONResponse
from fastapi.security import HTTPBasic, HTTPBasicCredentials

PG_USER = os.environ.get("POSTGRES_USER", "demo")
PG_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "")
PG_DB = os.environ.get("POSTGRES_DB", "demo")
PG_PORT = os.environ.get("PGPORT", "5432")

# If set, the Gradio UI and the JSON API require this password (username: admin).
APP_PASSWORD = os.environ.get("APP_PASSWORD", "")
APP_USER = os.environ.get("APP_USER", "admin")

TUNNEL_FILE = "/tmp/tunnel.json"   # written by start.sh (bore or ngrok)

BOOT = datetime.now()
HISTORY: list[tuple[str, bool, bool]] = []   # (HH:MM:SS, postgres_up, tunnel_up)
HISTORY_MAX = 500

# Keep-alive counter (incremented by the in-container loop AND any external cron).
KEEPALIVE = {"count": 0, "last": None, "src": "—"}


# ---------------------------------------------------------------- helpers ----
def get_tunnel():
    """Return (host, port) of the public TCP tunnel from /tmp/tunnel.json, or None."""
    try:
        with open(TUNNEL_FILE) as f:
            d = json.load(f)
        if d.get("host") and d.get("port"):
            return d["host"], str(d["port"])
    except Exception:
        return None
    return None


def pg_ok() -> bool:
    try:
        with psycopg.connect(
            host="/tmp", port=PG_PORT, user=PG_USER,
            dbname=PG_DB, password=PG_PASSWORD, connect_timeout=3,
        ):
            return True
    except Exception:
        return False


def build_url(tunnel):
    if not tunnel:
        return None
    host, port = tunnel
    return f"postgresql://{PG_USER}:{PG_PASSWORD}@{host}:{port}/{PG_DB}"


def humanize(delta: timedelta) -> str:
    s = int(delta.total_seconds())
    d, s = divmod(s, 86400)
    h, s = divmod(s, 3600)
    m, s = divmod(s, 60)
    if d:
        return f"{d}d {h}h {m}m"
    if h:
        return f"{h}h {m}m {s}s"
    return f"{m}m {s}s"


def record_check():
    ok = pg_ok()
    tn = get_tunnel() is not None
    HISTORY.append((datetime.now().strftime("%H:%M:%S"), ok, tn))
    del HISTORY[:-HISTORY_MAX]
    return ok, tn


def _watcher_loop():
    while True:
        try:
            record_check()
        except Exception:
            pass
        time.sleep(30)


# Background uptime watcher — records a sample every 30s while the container lives.
threading.Thread(target=_watcher_loop, daemon=True).start()

# Periodic Postgres backups (to BACKUP_DIR and/or a private HF Dataset).
backup.start_scheduler()


# ------------------------------------------------------------- JSON / API ----
api = FastAPI(title="Remote Postgres", docs_url=None, redoc_url=None)
security = HTTPBasic(auto_error=False)


def require_auth(credentials: HTTPBasicCredentials | None = Depends(security)):
    if not APP_PASSWORD:
        return
    if credentials is None or not (
        secrets.compare_digest(credentials.username, APP_USER)
        and secrets.compare_digest(credentials.password, APP_PASSWORD)
    ):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Authentication required",
            headers={"WWW-Authenticate": "Basic"},
        )


@api.get("/health")
async def health():
    """Open endpoint — used by the self keep-alive and external uptime monitors."""
    tunnel = get_tunnel()
    ok = pg_ok()
    return JSONResponse({
        "status": "ok" if ok else "degraded",
        "postgres": ok,
        "tunnel": bool(tunnel),
        "uptime": humanize(datetime.now() - BOOT),
    })


@api.get("/keepalive")
async def keepalive(src: str = "self"):
    """Hit to reset HF's idle timer (by the in-container loop AND any external cron).
    Open + logged so you can confirm pings are landing."""
    KEEPALIVE["count"] += 1
    KEEPALIVE["last"] = datetime.now()
    KEEPALIVE["src"] = src
    print(f"[keepalive] hit #{KEEPALIVE['count']} from '{src}' at "
          f"{KEEPALIVE['last'].strftime('%Y-%m-%d %H:%M:%S')}", flush=True)
    return {"ok": True, "count": KEEPALIVE["count"], "src": src}


@api.get("/api/connection")
async def connection(_: None = Depends(require_auth)):
    tunnel = get_tunnel()
    host, port = tunnel if tunnel else (None, None)
    return {
        "postgres_up": pg_ok(),
        "connection_url": build_url(tunnel),
        "host": host, "port": port,
        "user": PG_USER, "password": PG_PASSWORD, "database": PG_DB,
    }


@api.post("/admin/backup")
async def admin_backup(_: None = Depends(require_auth)):
    try:
        path = backup.backup()
        return {"ok": True, "file": os.path.basename(path)}
    except Exception as e:
        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)


@api.post("/admin/restore")
async def admin_restore(_: None = Depends(require_auth)):
    """Reload the latest backup (whole cluster). Cleanest on a fresh/empty DB;
    on a populated DB it best-effort merges and skips rows that already exist."""
    try:
        done = backup.restore()
        return {"ok": bool(done)}
    except Exception as e:
        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)


# ------------------------------------------------- Gradio uptime watcher ----
def snapshot():
    tunnel = get_tunnel()
    up = pg_ok()
    url = build_url(tunnel) or "— tunnel starting… refresh in a few seconds —"
    host, port = tunnel if tunnel else ("—", "—")

    samples = len(HISTORY)
    pct = (100 * sum(1 for _, o, _ in HISTORY if o) / samples) if samples else 100.0
    status_md = (
        f"### Status\n"
        f"- **Postgres:** {'🟢 up' if up else '🔴 down'}\n"
        f"- **Public tunnel:** {'🟢 up' if tunnel else '🟡 off'}\n"
        f"- **Container uptime:** {humanize(datetime.now() - BOOT)}\n"
        f"- **PG availability (this container):** {pct:.1f}%  ·  {samples} checks\n"
        f"- **Keep-alive hits:** {KEEPALIVE['count']}  ·  last from "
        f"`{KEEPALIVE['src']}` at "
        f"{KEEPALIVE['last'].strftime('%H:%M:%S') if KEEPALIVE['last'] else 'never'}\n"
        f"- **Backups:** {backup.STATE['count']} taken  ·  last: "
        f"{backup.STATE['last'].strftime('%H:%M:%S') if backup.STATE['last'] else 'none yet'}"
        f"  ·  HF off-Space: {'on' if backup.STATE['hf'] else 'off'}\n"
        f"- **Last checked:** {datetime.now().strftime('%H:%M:%S')}"
    )
    table = [
        [t, "✅" if o else "❌", "✅" if n else "❌"]
        for (t, o, n) in reversed(HISTORY[-25:])
    ]
    return status_md, url, host, port, PG_USER, PG_PASSWORD, PG_DB, table


with gr.Blocks(title="Remote Postgres — Uptime Watcher", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🐘 Remote Postgres — Uptime Watcher")
    gr.Markdown(
        "Throwaway Postgres for demo projects, exposed over a TCP tunnel (bore). "
        "Data is backed up, but the URL **changes on every restart**."
    )
    status_box = gr.Markdown()
    url_box = gr.Textbox(
        label="Connection URL — paste into your projects",
        interactive=False, show_copy_button=True,
    )
    with gr.Row():
        host_box = gr.Textbox(label="Host", interactive=False)
        port_box = gr.Textbox(label="Port", interactive=False)
    with gr.Row():
        user_box = gr.Textbox(label="User", interactive=False)
        pw_box = gr.Textbox(label="Password", interactive=False)
        db_box = gr.Textbox(label="Database", interactive=False)
    history_box = gr.Dataframe(
        headers=["time", "postgres", "tunnel"],
        label="Recent checks (newest first)",
        interactive=False,
    )
    with gr.Row():
        refresh_btn = gr.Button("Refresh now", variant="primary")
        backup_btn = gr.Button("Back up now")
        restore_btn = gr.Button("Restore latest")

    def backup_now():
        try:
            backup.backup()
        except Exception as e:
            print(f"[ui] manual backup failed: {e}", flush=True)
        return snapshot()

    def restore_now():
        try:
            backup.restore()
        except Exception as e:
            print(f"[ui] manual restore failed: {e}", flush=True)
        return snapshot()

    outputs = [status_box, url_box, host_box, port_box,
               user_box, pw_box, db_box, history_box]
    timer = gr.Timer(10)
    timer.tick(snapshot, outputs=outputs)
    refresh_btn.click(snapshot, outputs=outputs)
    backup_btn.click(backup_now, outputs=outputs)
    restore_btn.click(restore_now, outputs=outputs)
    demo.load(snapshot, outputs=outputs)


auth = (APP_USER, APP_PASSWORD) if APP_PASSWORD else None
app = gr.mount_gradio_app(api, demo, path="/", auth=auth)