#!/usr/bin/env python3
import os
import sys
from motor.motor_asyncio import AsyncIOMotorClient
from sanic import Sanic, exceptions
from sanic.log import logger

MONGODB_HOST = os.getenv("MONGODB_HOST")
if not MONGODB_HOST:
    raise ValueError("No MONGODB_HOST specified")

PROMETHEUS_BEARER_TOKEN = os.getenv("PROMETHEUS_BEARER_TOKEN")
if not PROMETHEUS_BEARER_TOKEN:
    raise ValueError("No PROMETHEUS_BEARER_TOKEN specified")

app = Sanic("exporter")
usernames = list(set(sys.argv[1:]))


@app.listener("before_server_start")
async def setup_db(app, loop):
    app.ctx.db = AsyncIOMotorClient(MONGODB_HOST).get_default_database()


async def wrap(i, prefix="wildduck_"):
    metrics_seen = set()
    async for name, tp, value, labels in i:
        if name not in metrics_seen:
            yield "# TYPE %s %s" % (name, tp)
            metrics_seen.add(name)
        yield "%s%s %d" % (
            prefix + name,
            ("{%s}" % ",".join(["%s=\"%s\"" % j for j in labels.items()]) if labels else ""),
            value)


async def fetch():
    last_login = {}
    user_labels = {}
    args = []
    if usernames:
        args.append({"username": {"$in": usernames}})
    async for u in app.ctx.db.users.find(*args):
        labels = {
            "username": u["username"],
            "email": u["address"],
        }
        user_labels[u["_id"]] = labels

        if u["lastLogin"]["time"]:
            last_login[u["_id"]] = u["lastLogin"]["time"]

        if u["storageUsed"]:
            yield "storage_used", "gauge", u["storageUsed"], labels
        if u["targets"]:
            yield "forwarding_addresses", "gauge", len(u["targets"]), labels
        yield "account_enabled", "gauge", not u["disabled"], labels

    # Merge application specific passwords last used timestamps
    async for a in app.ctx.db.asps.find({"user": {"$in": list(last_login.keys())}}):
        if a["used"] and a["used"] > last_login[a["user"]]:
            last_login[a["user"]] = a["used"]

    for user, dt in last_login.items():
        yield "last_login", "gauge", dt.timestamp(), user_labels[user]


@app.route("/metrics")
async def view_export(request):
    if request.token != PROMETHEUS_BEARER_TOKEN:
        raise exceptions.Forbidden("Invalid bearer token")
    response = await request.respond(content_type="text/plain")
    async for line in wrap(fetch()):
        await response.send(line + "\n")


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=3001, single_process=True)