"""
Central collector for the rclone fleet monitor.

Agents POST /api/ingest every few seconds (their heartbeat + parsed log events).
The React dashboard polls GET /api/overview and per-machine endpoints.

Storage: MongoDB
  machines   one doc per machine — last_seen, current round, live progress
  pairs      one doc per (machine, sync pair) — direction, last status/exit
  events     append-only event log (rounds, files, errors, pair results)
"""
import json
import os
import re
import subprocess
import threading
import time
from datetime import datetime, timedelta, timezone

from fastapi import FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pymongo import ASCENDING, DESCENDING, MongoClient

ONLINE_WINDOW_S = int(os.environ.get("ONLINE_WINDOW_S", "60"))  # heartbeat gap before offline
STALLED_ROUND_S = int(os.environ.get("STALLED_ROUND_S", "900"))  # online but no round_end for this long
EVENT_KEEP = {"round_start", "round_end", "pair_ok", "pair_fail",
              "file_synced", "error"}   # progress events stay live-only
LOG_TS_FMT = "%Y/%m/%d %H:%M:%S"

# tz_aware: give datetimes back with UTC tzinfo, matching what we store
client = MongoClient(os.environ.get("MONGO_URL", "mongodb://mongo:27017"), tz_aware=True)
db = client.rclonemon
db.events.create_index([("machine_id", ASCENDING), ("received_at", DESCENDING)])
db.events.create_index([("type", ASCENDING), ("received_at", DESCENDING)])
db.pairs.create_index([("machine_id", ASCENDING), ("pair", ASCENDING)], unique=True)

app = FastAPI(title="Rclone Fleet Monitor")

# Agents on other machines (a friend's system, office sites) must present this
# token. Empty AGENT_TOKEN = auth disabled (local dev only).
AGENT_TOKEN = os.environ.get("AGENT_TOKEN", "")

# Allow an external RMM frontend to call this API from another origin.
app.add_middleware(
    CORSMiddleware,
    allow_origins=os.environ.get("CORS_ORIGINS", "*").split(","),
    allow_methods=["*"],
    allow_headers=["*"],
)


def now():
    return datetime.now(timezone.utc)


# ---------------------------------------------------------------------------
# Server-side verification — the audit feature.
# The collector has its OWN rclone remote (written by start.sh). A background
# thread lists every up-pair's remote folder with `rclone lsjson`; the fleet
# view then diffs each machine's local inventory against the server listing:
#   pending = file exists locally but not (yet) on the server
#   missing = pending AND older than GRACE => should have synced => alert
# ---------------------------------------------------------------------------
VERIFY_ENABLED = bool(os.environ.get("SYNOLOGY_HOST"))
VERIFY_INTERVAL = int(os.environ.get("VERIFY_INTERVAL", "60"))
MISSING_GRACE_S = int(os.environ.get("MISSING_GRACE_S", "600"))
SERVER_LISTINGS = {}  # "remote:path" -> {"files": {relpath: size}, "at": datetime}

# Which rclone remote to list for a given site, e.g. "TEST=synoreal"
# (his real NAS) while everything else defaults to "synodrive" (local fake).
DEFAULT_REMOTE = "synodrive"
SITE_REMOTES = dict(kv.split("=", 1) for kv in
                    os.environ.get("SITE_REMOTES", "").split(",") if "=" in kv)


def _remote_for_site(site):
    return SITE_REMOTES.get(site or "", DEFAULT_REMOTE)


def _verify_loop():
    while True:
        tasks = set()
        for p in db.pairs.find({"direction": "up"}):
            label = p.get("pair", "")
            if "server:" not in label:
                continue
            m = db.machines.find_one({"machine_id": p["machine_id"]}, {"site": 1})
            remote = _remote_for_site((m or {}).get("site"))
            tasks.add((remote, label.split("server:", 1)[1].strip()))
        for remote, r in tasks:
            try:
                out = subprocess.run(
                    ["rclone", "lsjson", "--recursive", "--files-only", f"{remote}:{r}"],
                    capture_output=True, timeout=120)
                if out.returncode == 0:
                    files = {f["Path"]: f["Size"] for f in json.loads(out.stdout)}
                    SERVER_LISTINGS[f"{remote}:{r}"] = {"files": files, "at": now()}
                else:
                    print(f"[verify] {remote}:{r}: exit {out.returncode} "
                          f"{out.stderr.decode()[:120]}", flush=True)
            except Exception as e:
                print(f"[verify] {remote}:{r}: {e}", flush=True)
        time.sleep(VERIFY_INTERVAL)


if VERIFY_ENABLED:
    threading.Thread(target=_verify_loop, daemon=True).start()


def _speed_bps(s):
    """'4.2 MiB/s' -> bytes/second, or None."""
    m = re.match(r"([\d.]+)\s*([KMGT]i?)?B/s", s or "")
    if not m:
        return None
    mult = {"": 1, "K": 1e3, "Ki": 1024, "M": 1e6, "Mi": 1024**2,
            "G": 1e9, "Gi": 1024**3, "T": 1e12, "Ti": 1024**4}[m.group(2) or ""]
    return float(m.group(1)) * mult


def _annotate_pending(m, pairs):
    """Attach pending/missing/backlog info to up pairs.
    Returns (pending_files, pending_bytes, eta_s) — all None if NO pair could
    be verified (so the UI shows '–' instead of a misleading 0)."""
    inv = m.get("inventory") or {}
    remote_name = _remote_for_site(m.get("site"))
    total_pending = 0
    total_bytes = 0
    verified_any = False
    for p in pairs:
        label = p.get("pair", "")
        if p.get("direction") != "up" or "server:" not in label:
            continue
        local = label.split("→")[0].strip()
        rpath = label.split("server:", 1)[1].strip()
        listing = SERVER_LISTINGS.get(f"{remote_name}:{rpath}")
        files = inv.get(local)
        if listing is None or files is None:
            continue
        verified_any = True
        pending = [f for f in files if f["name"] not in listing["files"]]
        p["pending_count"] = len(pending)
        p["pending_bytes"] = sum(f.get("size", 0) for f in pending)
        p["verified_at"] = listing["at"]
        p["missing"] = [f["name"] for f in pending
                        if time.time() - f["mtime"] > MISSING_GRACE_S][:20]
        total_pending += len(pending)
        total_bytes += p["pending_bytes"]
    if not verified_any:
        return None, None, None
    # backlog ETA: remaining bytes / the transfer speed rclone last reported
    bps = _speed_bps((m.get("progress") or {}).get("speed"))
    eta_s = int(total_bytes / bps) if bps and bps > 0 and total_bytes > 0 else None
    return total_pending, total_bytes, eta_s


def require_agent_token(authorization: str | None):
    if AGENT_TOKEN and authorization != f"Bearer {AGENT_TOKEN}":
        raise HTTPException(401, "missing or invalid agent token")


@app.get("/api/health")
def health():
    """Liveness probe for RMM uptime checks: is the collector + DB reachable?"""
    db.command("ping")
    return {"ok": True, "time": now()}


@app.post("/api/ingest")
def ingest(payload: dict, authorization: str | None = Header(None)):
    require_agent_token(authorization)
    mid = payload.get("machine_id")
    if not mid:
        raise HTTPException(400, "machine_id required")
    ts = now()
    machine_set = {"machine_id": mid, "site": payload.get("site"), "last_seen": ts}
    if "inventory" in payload:
        machine_set["inventory"] = payload["inventory"]

    to_insert = []
    for ev in payload.get("events", []):
        etype = ev.get("type")
        if etype == "round_start":
            machine_set["current_round"] = ev.get("round")
            machine_set["round_started_at"] = ts
            machine_set["progress"] = None
        elif etype == "round_end":
            machine_set["last_round"] = ev.get("round")
            machine_set["last_round_at"] = ts
            machine_set["progress"] = None
        elif etype == "progress":
            machine_set["progress"] = {k: ev.get(k) for k in
                                       ("pair", "direction", "percent", "speed", "eta", "done", "total")}
        elif etype in ("pair_start", "pair_ok", "pair_fail"):
            status = {"pair_start": "running", "pair_ok": "ok", "pair_fail": "fail"}[etype]
            upd = {"machine_id": mid, "pair": ev.get("pair"),
                   "direction": ev.get("direction"), "last_status": status,
                   "updated_at": ts}
            if etype == "pair_fail":
                upd["last_exit"] = ev.get("exit_code")
                upd["last_fail_at"] = ts
            elif etype == "pair_ok":
                upd["last_exit"] = 0
                upd["last_ok_at"] = ts
            db.pairs.update_one({"machine_id": mid, "pair": ev.get("pair")},
                                {"$set": upd}, upsert=True)
        if etype in EVENT_KEEP:
            to_insert.append({**ev, "machine_id": mid, "received_at": ts})

    db.machines.update_one({"machine_id": mid}, {"$set": machine_set}, upsert=True)
    if to_insert:
        db.events.insert_many(to_insert)
    # Rounds can be far shorter than the poll interval (a 5 GB upload finishing
    # in 15s), so live progress is easy to miss entirely. Summarise every
    # finished round instead: files, bytes, duration, average throughput.
    for ev in payload.get("events", []):
        if ev.get("type") == "round_end" and ev.get("round") is not None:
            _summarise_round(mid, ev, payload.get("inventory") or {}, ts)
    return {"ok": True, "stored": len(to_insert)}


def _summarise_round(mid, end_ev, inventory, ts):
    rnd = end_ev["round"]
    if db.events.find_one({"machine_id": mid, "type": "round_summary", "round": rnd}):
        return
    sizes = {f["name"]: f.get("size", 0) for files in inventory.values() for f in files}
    synced = list(db.events.find({"machine_id": mid, "type": "file_synced", "round": rnd},
                                 {"file": 1, "direction": 1}))
    # only up-pairs have local sizes to look up; downloads land outside WATCH_DIRS
    total_bytes = sum(sizes.get(e.get("file"), 0) for e in synced)

    duration = None
    start_ev = db.events.find_one({"machine_id": mid, "type": "round_start", "round": rnd})
    try:
        t0 = datetime.strptime(start_ev["log_ts"].strip(), LOG_TS_FMT)
        t1 = datetime.strptime(end_ev["log_ts"].strip(), LOG_TS_FMT)
        duration = (t1 - t0).total_seconds()
    except Exception:
        pass

    avg_bps = total_bytes / duration if duration and duration > 0 and total_bytes else None
    summary = {"type": "round_summary", "machine_id": mid, "round": rnd,
               "files": len(synced), "bytes": total_bytes,
               "duration_s": duration, "avg_bps": avg_bps, "received_at": ts}
    db.events.insert_one(dict(summary))
    if len(synced):  # remember the last round that actually moved data
        summary.pop("_id", None)
        db.machines.update_one({"machine_id": mid}, {"$set": {"last_transfer": summary}})


def _machine_view(m, ts):
    last_seen = m.get("last_seen")
    online = bool(last_seen) and (ts - last_seen).total_seconds() < ONLINE_WINDOW_S
    hour_ago = ts - timedelta(hours=1)
    mid = m["machine_id"]
    pairs = list(db.pairs.find({"machine_id": mid}, {"_id": 0}))
    if VERIFY_ENABLED:
        pending_total, pending_bytes, backlog_eta_s = _annotate_pending(m, pairs)
    else:
        pending_total = pending_bytes = backlog_eta_s = None
    return {
        "pending_total": pending_total,
        "pending_bytes": pending_bytes,
        "backlog_eta_s": backlog_eta_s,
        "machine_id": mid,
        "site": m.get("site"),
        "online": online,
        "last_seen": last_seen,
        "current_round": m.get("current_round"),
        "last_round": m.get("last_round"),
        "last_round_at": m.get("last_round_at"),
        "progress": m.get("progress"),
        "last_transfer": m.get("last_transfer"),
        "pairs": pairs,
        "files_1h": db.events.count_documents(
            {"machine_id": mid, "type": "file_synced", "received_at": {"$gt": hour_ago}}),
        "errors_1h": db.events.count_documents(
            {"machine_id": mid, "type": "error", "received_at": {"$gt": hour_ago}}),
    }


def _fleet():
    ts = now()
    machines = [_machine_view(m, ts) for m in db.machines.find()]
    machines.sort(key=lambda m: m["machine_id"])

    alerts = []
    for m in machines:
        if not m["online"]:
            alerts.append({"severity": "critical", "machine_id": m["machine_id"],
                           "message": "machine offline — no heartbeat"})
        for p in m["pairs"]:
            if p.get("last_status") == "fail":
                alerts.append({"severity": "serious", "machine_id": m["machine_id"],
                               "message": f"sync pair failing (exit {p.get('last_exit')}): {p['pair']}"})
        if m["online"] and m["errors_1h"] > 0:
            alerts.append({"severity": "warning", "machine_id": m["machine_id"],
                           "message": f"{m['errors_1h']} rclone error(s) in the last hour"})
        for p in m["pairs"]:
            if p.get("missing"):
                alerts.append({"severity": "serious", "machine_id": m["machine_id"],
                               "message": f"{len(p['missing'])} file(s) NOT on server after "
                                          f"{MISSING_GRACE_S // 60}min grace: "
                                          f"{', '.join(p['missing'][:3])}…"})
        # "alive but not producing": agent heartbeats fine, but the sync
        # container hasn't finished a round in far too long
        if m["online"] and m.get("last_round_at"):
            stalled = (ts - m["last_round_at"]).total_seconds()
            if stalled > STALLED_ROUND_S:
                age = f"{stalled / 3600:.1f}h" if stalled >= 3600 else f"{int(stalled // 60)}min"
                alerts.append({"severity": "serious", "machine_id": m["machine_id"],
                               "message": f"agent alive but NO sync round completed in {age} "
                                          f"— sync container stopped or stuck?"})
    return ts, machines, alerts


@app.get("/api/overview")
def overview():
    ts, machines, alerts = _fleet()
    return {"generated_at": ts, "machines": machines, "alerts": alerts}


@app.get("/api/alerts")
def alerts_only():
    """Lightweight endpoint for external RMM polling — alerts, no machine detail."""
    ts, _, alerts = _fleet()
    return {"generated_at": ts, "count": len(alerts), "alerts": alerts}


@app.get("/api/machines/{machine_id}/events")
def machine_events(machine_id: str, type: str | None = None, limit: int = 50):
    q = {"machine_id": machine_id}
    if type:
        q["type"] = {"$in": type.split(",")}
    evs = list(db.events.find(q, {"_id": 0}).sort("received_at", DESCENDING).limit(min(limit, 500)))
    return {"machine_id": machine_id, "events": evs}


# React build (mounted at /app/static) — must be mounted AFTER the API routes
if os.path.isdir("static"):
    app.mount("/", StaticFiles(directory="static", html=True), name="ui")