#!/usr/bin/env python3
"""backhaul agent — the containerized DB-backup client. Runs as a sidecar on a database's
private Docker network: it reaches the DB internally and calls home outbound over HTTPS.
No inbound ports. Stdlib only; the DB client tools (pg_dump/mysqldump/mongodump) come from
the image.

Every minute (in-container cron) it CHECKS IN: heartbeats, reports the databases it can see,
delivers any pending dump result, and asks what to do. The app owns the schedule and replies
with a dump command when one is due; the agent then dumps and streams it home. The app also
advertises the current agent-script version, and the agent self-updates (atomic swap after a
compile check) — same mechanism as snapshoot.

Config (env, from the `docker run`):
  BACKHAUL_URL     e.g. https://backhaul.skui.io
  BACKHAUL_TOKEN   agent token from the backhaul admin
  BACKHAUL_TARGETS comma-separated DB targets, each: engine:host:port:db:user:pass
                   e.g. postgres:db:5432:nextcloud:nextcloud:nextcloud
"""
import json
import os
import subprocess
import sys
import urllib.request

VERSION = "1.3"
RESULT_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".dump-result.json")


def targets():
    """Parse BACKHAUL_TARGETS into dicts. A networked DB is engine:host:port:db:user:pass;
    a SQLite file is sqlite:<path> (a file bind-mounted into the agent, not a server)."""
    out = []
    for spec in os.environ.get("BACKHAUL_TARGETS", "").split(","):
        spec = spec.strip()
        if not spec:
            continue
        parts = spec.split(":")
        if parts[0] == "sqlite":
            path = ":".join(parts[1:]).strip()      # rejoin in case the path holds a ':'
            if not path:
                print(f"skipping sqlite target without a path: {spec}", file=sys.stderr)
                continue
            out.append({"engine": "sqlite", "path": path,
                        "name": f"sqlite/{os.path.basename(path)}"})
            continue
        if len(parts) < 6:
            print(f"skipping malformed target: {spec}", file=sys.stderr)
            continue
        engine, host, port, db, user, pw = parts[:6]
        out.append({"engine": engine, "host": host, "port": port,
                    "db": db, "user": user, "pw": pw,
                    "name": f"{engine}/{db}"})
    return out


def db_version(t):
    """Fetch the engine version string so the app can pick a matching restore image."""
    try:
        if t["engine"] == "sqlite":
            out = subprocess.run(("sqlite3", "--version"),
                                 capture_output=True, text=True, timeout=15).stdout
            return "SQLite " + ((out.strip().split(" ")[0]) or "3")
        if t["engine"] == "postgres":
            env = dict(os.environ, PGPASSWORD=t["pw"])
            out = subprocess.run(("psql", "-h", t["host"], "-p", t["port"], "-U", t["user"],
                                  "-d", t["db"], "-tAc", "select version()"),
                                 capture_output=True, text=True, timeout=30, env=env).stdout
            return out.strip().split(" on ")[0][:60]
        if t["engine"] in ("mysql", "mariadb"):
            out = subprocess.run(("mysql", "-h", t["host"], "-P", t["port"], "-u", t["user"],
                                  f"-p{t['pw']}", "-N", "-e", "select version()"),
                                 capture_output=True, text=True, timeout=30).stdout
            return "MySQL/MariaDB " + out.strip()[:40]
        if t["engine"] == "mongo":
            return "MongoDB"
    except Exception:
        pass
    return t["engine"]


def db_stats(t):
    """Read-only live metrics for the dashboard's 'nerd stats' — size, tables, rows, index
    footprint, connections, cache-hit ratio, the biggest tables, server version + uptime.
    Best-effort: any failure returns None and never disturbs the heartbeat. Postgres only."""
    if t["engine"] == "sqlite":
        return _sqlite_stats(t)
    if t["engine"] != "postgres":
        return None
    env = dict(os.environ, PGPASSWORD=t["pw"])

    def q(sql):
        r = subprocess.run(("psql", "-h", t["host"], "-p", t["port"], "-U", t["user"],
                            "-d", t["db"], "-tAc", sql),
                           capture_output=True, text=True, timeout=30, env=env)
        return r.stdout.strip()

    try:
        top = q("SELECT COALESCE(json_agg(x),'[]') FROM (SELECT relname AS name, "
                "pg_total_relation_size(relid) AS bytes, n_live_tup AS rows "
                "FROM pg_stat_user_tables ORDER BY pg_total_relation_size(relid) DESC "
                "LIMIT 10) x")
        hit = q("SELECT round(100.0*sum(blks_hit)/NULLIF(sum(blks_hit)+sum(blks_read),0),2) "
                "FROM pg_stat_database WHERE datname=current_database()")
        up = q("SELECT EXTRACT(EPOCH FROM (now()-pg_postmaster_start_time()))::bigint")
        return {
            "size_bytes": int(q("SELECT pg_database_size(current_database())") or 0),
            "tables": int(q("SELECT count(*) FROM pg_stat_user_tables") or 0),
            "rows_est": int(float(q("SELECT COALESCE(SUM(n_live_tup),0) "
                                    "FROM pg_stat_user_tables") or 0)),
            "index_bytes": int(q("SELECT COALESCE(SUM(pg_indexes_size(relid)),0) "
                                 "FROM pg_stat_user_tables") or 0),
            "connections": int(q("SELECT count(*) FROM pg_stat_activity "
                                 "WHERE datname=current_database()") or 0),
            "cache_hit_pct": float(hit) if hit else None,
            "server_version": q("SHOW server_version"),
            "uptime_s": int(up) if up else None,
            "top_tables": json.loads(top or "[]"),
        }
    except Exception:
        return None


def _sqlite_stats(t):
    """Live metrics for a SQLite file: file size, table count, row totals, largest tables."""
    def q(sql):
        r = subprocess.run(("sqlite3", t["path"], sql),
                           capture_output=True, text=True, timeout=30)
        return r.stdout.strip()
    try:
        names = [n for n in q("SELECT name FROM sqlite_master WHERE type='table' "
                              "AND name NOT LIKE 'sqlite_%' ORDER BY name").splitlines() if n]
        top = []
        rows = 0
        for n in names:
            c = int(q(f'SELECT count(*) FROM "{n}"') or 0)
            rows += c
            top.append({"name": n, "bytes": 0, "rows": c})
        top.sort(key=lambda x: x["rows"], reverse=True)
        size = os.path.getsize(t["path"]) if os.path.exists(t["path"]) else 0
        return {"size_bytes": size, "tables": len(names), "rows_est": rows,
                "index_bytes": 0, "connections": 0, "cache_hit_pct": None,
                "server_version": "SQLite " + q("SELECT sqlite_version()"),
                "uptime_s": None, "top_tables": top[:10]}
    except Exception:
        return None


def dump_cmd(t):
    """The dump command + environment for an engine. Streams compressed to stdout."""
    if t["engine"] == "sqlite":
        # `.dump` runs in a read transaction → a consistent SQL snapshot even under writes.
        # The file (and its -wal/-shm) must be readable in the agent container (bind-mount).
        return (("sqlite3", t["path"], ".dump"), dict(os.environ))
    if t["engine"] == "postgres":
        env = dict(os.environ, PGPASSWORD=t["pw"])
        # custom format is already compressed and is what pg_restore wants
        return (("pg_dump", "-h", t["host"], "-p", t["port"], "-U", t["user"],
                 "-d", t["db"], "-Fc"), env)
    if t["engine"] in ("mysql", "mariadb"):
        return (("sh", "-c",
                 f"mysqldump -h {t['host']} -P {t['port']} -u {t['user']} "
                 f"-p{t['pw']} --single-transaction --routines {t['db']} | gzip"),
                dict(os.environ))
    if t["engine"] == "mongo":
        return (("mongodump", "--host", t["host"], "--port", t["port"],
                 "--db", t["db"], "--archive", "--gzip"), dict(os.environ))
    raise ValueError(f"unknown engine {t['engine']}")


UPLOAD_CHUNK = 64 * 1024 * 1024  # keep each request well under Cloudflare's 100 MB body cap


def _post_chunk(url, token, name, meta, upload_id, index, final, data):
    """POST one slice of a dump. Each request has a normal Content-Length under the CDN limit,
    so this works over the public HTTPS URL from any server, not just co-located ones."""
    req = urllib.request.Request(
        url.rstrip("/") + "/api/dump", data=bytes(data),
        headers={"X-Agent-Token": token, "X-DB-Name": name, "X-DB-Meta": meta,
                 "X-Upload-Id": upload_id, "X-Chunk-Index": str(index),
                 "X-Chunk-Final": "1" if final else "0",
                 "Content-Type": "application/octet-stream",
                 "User-Agent": f"backhaul-agent/{VERSION}"})
    with urllib.request.urlopen(req, timeout=3600) as resp:
        return resp.status == 200


def run_dump(t, url, token):
    """Stream a dump home in <100 MB slices so it passes the CDN's upload limit from any
    server; the app reassembles by upload id and finalizes on the last slice. The dump is never
    fully buffered — only one slice is held at a time."""
    cmd, env = dump_cmd(t)
    meta = json.dumps({"engine": t["engine"], "version": db_version(t)})
    upload_id = os.urandom(8).hex()
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, env=env)
    buf = bytearray()
    index = 0
    ok = True
    try:
        while True:
            b = proc.stdout.read(1 << 20)
            if not b:
                break
            buf += b
            while len(buf) >= UPLOAD_CHUNK:
                if not _post_chunk(url, token, t["name"], meta, upload_id, index, False,
                                   buf[:UPLOAD_CHUNK]):
                    ok = False
                    break
                del buf[:UPLOAD_CHUNK]
                index += 1
            if not ok:
                break
        proc.wait()
        if ok:                          # final slice finalizes the backup (even if empty)
            ok = _post_chunk(url, token, t["name"], meta, upload_id, index, True, bytes(buf))
    except Exception as e:
        print("dump upload failed:", e, file=sys.stderr)
        ok = False
    finally:
        if proc.poll() is None:
            proc.kill()
            proc.wait()
    return {"name": t["name"], "ok": ok and proc.returncode == 0}


def maybe_self_update(answer, url, token):
    """Swap our own script when the app advertises a newer version — atomic, compile-checked.
    Disable with BACKHAUL_AUTOUPDATE=0."""
    latest = answer.get("agent_latest")
    if (not latest or latest == VERSION
            or os.environ.get("BACKHAUL_AUTOUPDATE", "1") == "0"):
        return
    me = os.path.abspath(__file__)
    try:
        req = urllib.request.Request(
            url.rstrip("/") + "/api/agent-script",
            headers={"X-Agent-Token": token,
                     "User-Agent": f"backhaul-agent/{VERSION}"})  # else Cloudflare 403s urllib
        with urllib.request.urlopen(req, timeout=30) as r:
            source = r.read()
        text = source.decode()
        if "backhaul-agent" not in text or f'VERSION = "{latest}"' not in text:
            return
        compile(text, me, "exec")
        tmp = me + ".new"
        with open(tmp, "wb") as f:
            f.write(source)
        os.chmod(tmp, os.stat(me).st_mode)
        os.replace(tmp, me)
        print(f"self-updated {VERSION} -> {latest}")
    except Exception as e:
        print("self-update failed:", e, file=sys.stderr)


def checkin(url, token, tgts, pending):
    """POST the heartbeat: which DBs we see + any finished dump result. The reply tells us
    what to do next and the latest agent version."""
    stats = {}
    for t in tgts:
        s = db_stats(t)
        if s:
            stats[t["name"]] = s
    body = json.dumps({"agent_version": VERSION,
                       "databases": [{"name": t["name"], "engine": t["engine"]} for t in tgts],
                       "results": pending, "stats": stats}).encode()
    req = urllib.request.Request(url.rstrip("/") + "/api/checkin", data=body,
                                 headers={"X-Agent-Token": token,
                                          "Content-Type": "application/json",
                                          "User-Agent": f"backhaul-agent/{VERSION}"})
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read() or b"{}")


def main():
    url = os.environ.get("BACKHAUL_URL", "").rstrip("/")
    token = os.environ.get("BACKHAUL_TOKEN", "")
    if not (url and token):
        sys.exit("BACKHAUL_URL and BACKHAUL_TOKEN must be set")
    tgts = targets()

    pending = []
    if os.path.exists(RESULT_FILE):
        try:
            pending = json.load(open(RESULT_FILE))
        except ValueError:
            pending = []

    answer = checkin(url, token, tgts, pending)
    if pending:
        os.remove(RESULT_FILE)

    # run any dump the app asked for (matched by name), record result for next check-in
    results = []
    for want in answer.get("dump") or []:
        t = next((t for t in tgts if t["name"] == want), None)
        if t:
            print(f"dumping {t['name']}...", flush=True)
            results.append(run_dump(t, url, token))
    if results:
        json.dump(results, open(RESULT_FILE, "w"))

    maybe_self_update(answer, url, token)


if __name__ == "__main__":
    main()