#!/usr/bin/env python3
"""Daily metrics snapshot for the Inference Cooperative portal.

Writes one row per day into the portal's SQLite store (`metrics_daily`, the same
`members.db` that holds the members table) so the admin panel and the member
dashboard can chart trends from one source of truth.

    metrics_snapshot.py             record / refresh TODAY's row
    metrics_snapshot.py --backfill  rebuild history (idempotent)

Rules this file follows deliberately:

- **Fail closed per metric.** A source that cannot be read yields NULL for its
  columns, never a zero. A zero reads as "nobody", which is a lie that would
  survive into a chart; NULL reads as "not measured". Each failure is printed,
  and the task's stdout lands in the app log (`cloudron logs --app <id>`).
- **Cumulative vs daily.** `members_active`, `members_total`, `newsletter_subs`
  and `matrix_space_members` are counts AS OF that day; `spend_usd` and
  `requests` are THAT day's traffic. Charts must treat them accordingly.
- **Runs in Cloudron's scheduler container**, which gets the platform env
  (declared app env + addon vars) but NOT variables the app's own start.sh
  sourced. Everything needed here is a declared app env var: LITELLM_DB_*,
  LISTMONK_*, MATRIX_BOT_TOKEN, LISTMONK_LIST_ID.
- `members_new` uses `welcome_sent_at`, the provisioning timestamp — the closest
  thing the portal stores to a join date. Manual members have no welcome email,
  so their joins are absent from the growth curve (noted, not guessed).
"""

import datetime as dt
import json
import os
import sqlite3
import sys
import urllib.parse
import urllib.request

DB_PATH = os.environ.get("DB_PATH", "/app/data/members.db")
LISTMONK_BASE = os.environ.get("LISTMONK_BASE", "https://newsletter.inference.coop")
LISTMONK_LIST_ID = int(os.environ.get("LISTMONK_LIST_ID") or 3)
MATRIX_BASE = os.environ.get("MATRIX_BASE", "https://matrix.inference.coop")
SPACE = os.environ.get("MATRIX_SPACE_ID", "!FKOD2sSqeYvEq9jzr8jDL9ogf3TdJCEI1k4IsnBZsCM")


def today() -> str:
    return dt.datetime.now(dt.timezone.utc).strftime("%Y-%m-%d")


def connect() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH, timeout=20)
    conn.execute(
        """CREATE TABLE IF NOT EXISTS metrics_daily (
               day TEXT PRIMARY KEY,
               members_active INTEGER,
               members_total INTEGER,
               members_new INTEGER,
               newsletter_subs INTEGER,
               matrix_space_members INTEGER,
               spend_usd REAL,
               requests INTEGER,
               out_tokens INTEGER,
               energy_wh REAL,
               co2e_g REAL,
               recorded_at TEXT)"""
    )
    # CREATE IF NOT EXISTS never adds columns to an existing table, so migrate.
    for col, decl in (("out_tokens", "INTEGER"), ("energy_wh", "REAL"), ("co2e_g", "REAL")):
        try:
            conn.execute(f"ALTER TABLE metrics_daily ADD COLUMN {col} {decl}")
        except sqlite3.OperationalError:
            pass  # already present
    return conn


def _get_json(url: str, headers: dict | None = None):
    req = urllib.request.Request(url, headers=headers or {})
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read() or b"{}")


# ---------------------------------------------------------------- impact factors
# Energy per 1,000 OUTPUT tokens, and the carbon intensity of the electricity.
#
# The greenpt/* rows are MEASURED. GreenPT returns an `impact` object on every
# call (energy in watt-milliseconds, emissions in µgCO2e), so these were fitted
# from live calls to their API: energy regressed on output tokens across prompt
# sizes (R² >= 0.999), carbon taken as the grid intensity their own emissions
# imply (µg/mWh and g/kWh are the same ratio). Their published unit is mWh per
# 100 output tokens, and only OUTPUT tokens count — hence completion_tokens.
#
# Everything else is an ESTIMATE and is labelled as one. Providers that do not
# disclose impact get IMPACT_PROXY: an order-of-magnitude figure near the middle
# of the measured range with an EU grid average. It exists so a day's total is
# never silently missing models, not because we know the number. Refining these
# is a data edit here, not a code change — see co-op/support#26.
MEASURED = "measured (GreenPT impact API, fitted 2026-10-28)"
SAME_FAMILY = "measured for glm-5.3-flash; assumed equal for glm-5.3"
ESTIMATED = "estimated (order-of-magnitude proxy; provider does not disclose)"

# Longest prefix must come first: matching is startswith().
IMPACT_FACTORS = {
    "greenpt/deepseek-v4.1-flash": {"wh_per_1k_out": 0.3641, "gco2e_per_kwh": 193.0, "source": MEASURED},
    "greenpt/glm-5.3-flash":       {"wh_per_1k_out": 1.0782, "gco2e_per_kwh": 193.0, "source": MEASURED},
    "greenpt/glm-5.3":             {"wh_per_1k_out": 1.0782, "gco2e_per_kwh": 193.0, "source": SAME_FAMILY},
    "greenpt/green-r":             {"wh_per_1k_out": 0.5535, "gco2e_per_kwh": 23.6,  "source": MEASURED},
}
IMPACT_PROXY = {"wh_per_1k_out": 0.5000, "gco2e_per_kwh": 250.0, "source": ESTIMATED}

# ---------------------------------------------------------------- sources

def listmonk_joins() -> list[str] | None:
    """The date each current member joined the members list (YYYY-MM-DD).

    This is the honest source for a membership curve: the list is maintained by
    the portal's nightly reconcile to mirror ACTIVE members, and each
    subscription carries `subscription_created_at`. `welcome_sent_at` in the
    members table is NOT equivalent — most members were provisioned before the
    portal sent welcome emails, so a curve built on it lags by dozens of people.
    """
    token, user = os.environ.get("LISTMONK_API_TOKEN"), os.environ.get("LISTMONK_API_USER") or "portal-sync"
    if not token:
        print("metrics: LISTMONK_API_TOKEN not set — membership history unavailable")
        return None
    joins, page = [], 1
    try:
        while True:
            data = _get_json(
                f"{LISTMONK_BASE}/api/subscribers?list_id={LISTMONK_LIST_ID}&per_page=100&page={page}",
                {"Authorization": f"token {user}:{token}"},
            ) or {}
            rows = ((data.get("data") or {}).get("results")) or []
            for s in rows:
                for membership in s.get("lists") or []:
                    if membership.get("id") == LISTMONK_LIST_ID:
                        stamp = membership.get("subscription_created_at") or s.get("created_at")
                        if stamp:
                            joins.append(str(stamp)[:10])
            if not rows or len(rows) < 100 or page >= 20:
                break
            page += 1
    except Exception as exc:  # noqa: BLE001
        print(f"metrics: listmonk joins read failed ({exc})")
        return None
    return joins or None


def from_sqlite() -> dict:
    """Members, plus the provisioning dates that give the growth curve."""
    conn = connect()
    try:
        active = conn.execute("SELECT COUNT(*) FROM members WHERE active=1").fetchone()[0]
        total = conn.execute("SELECT COUNT(*) FROM members").fetchone()[0]
        joins = [
            row[0]
            for row in conn.execute(
                "SELECT welcome_sent_at FROM members WHERE welcome_sent_at IS NOT NULL AND welcome_sent_at != ''"
            )
        ]
    finally:
        conn.close()
    return {"active": active, "total": total, "joins": joins}


def from_listmonk() -> int | None:
    """Confirmed subscribers on the members list, counted rather than trusted.

    Listmonk's cached `subscriber_count` on the list object is stale and
    double-counted (it read 142 for a list that truly holds 71), so we page the
    subscribers endpoint and count confirmed memberships ourselves.
    """
    token, user = os.environ.get("LISTMONK_API_TOKEN"), os.environ.get("LISTMONK_API_USER") or "portal-sync"
    if not token:
        print("metrics: LISTMONK_API_TOKEN not set — newsletter_subs left NULL")
        return None
    try:
        total, page, seen = 0, 1, 0
        while True:
            data = _get_json(
                f"{LISTMONK_BASE}/api/subscribers?list_id={LISTMONK_LIST_ID}&per_page=100&page={page}",
                {"Authorization": f"token {user}:{token}"},
            ) or {}
            payload = data.get("data") or {}
            rows = payload.get("results") or []
            # The subscriber's own `status` is enabled/blocked; the LIST
            # membership status is what means "on this list" — counting the
            # former returned a silent, wrong 0.
            for s in rows:
                for membership in s.get("lists") or []:
                    if membership.get("id") == LISTMONK_LIST_ID and (
                        membership.get("subscription_status") or ""
                    ) == "confirmed":
                        total += 1
            seen += len(rows)
            if not rows or len(rows) < 100 or page >= 20:
                break
            page += 1
        if seen == 0:
            # Nothing came back at all: the endpoint or payload shape changed.
            # Record "not measured" (NULL) rather than a zero that reads as
            # "nobody is subscribed".
            print("metrics: listmonk returned no rows — newsletter_subs left NULL")
            return None
        return total
    except Exception as exc:  # noqa: BLE001 - report, never guess
        print(f"metrics: listmonk read failed ({exc}) — newsletter_subs left NULL")
        return None


def from_matrix() -> int | None:
    token = os.environ.get("MATRIX_BOT_TOKEN")
    if not token:
        print("metrics: MATRIX_BOT_TOKEN not set — matrix_space_members left NULL")
        return None
    try:
        data = _get_json(
            f"{MATRIX_BASE}/_matrix/client/v3/rooms/{urllib.parse.quote(SPACE)}/joined_members",
            {"Authorization": f"Bearer {token}"},
        ) or {}
        return len(data.get("joined") or {})
    except Exception as exc:  # noqa: BLE001
        print(f"metrics: matrix read failed ({exc}) — matrix_space_members left NULL")
        return None


def spend_by_day(since: dt.date | None = None) -> dict[str, tuple[float, int]]:
    """{day: (spend_usd, requests)} straight from LiteLLM's spend logs.

    Prefers `model_group` the way aggregate_model_usage does, but here we only
    need volume, so every row in the window counts.
    """
    try:
        import psycopg2
    except Exception as exc:  # noqa: BLE001
        print(f"metrics: psycopg2 unavailable ({exc})")
        return {}
    try:
        conn = psycopg2.connect(
            host=os.environ.get("LITELLM_DB_HOST"),
            port=os.environ.get("LITELLM_DB_PORT"),
            user=os.environ.get("LITELLM_DB_USER"),
            password=os.environ.get("LITELLM_DB_PASSWORD"),
            dbname=os.environ.get("LITELLM_DB_NAME"),
            connect_timeout=10,
            sslmode="prefer",
        )
    except Exception as exc:  # noqa: BLE001
        print(f"metrics: LiteLLM DB connect failed ({exc}) — spend left NULL")
        return {}
    try:
        where, params = "", []
        if since:
            where, params = 'WHERE "startTime" >= %s', [since.isoformat()]
        with conn.cursor() as cur:
            cur.execute(
                f"""SELECT to_char("startTime", 'YYYY-MM-DD') AS day,
                           COALESCE(SUM(spend), 0), COUNT(*)
                    FROM "LiteLLM_SpendLogs" {where}
                    GROUP BY 1 ORDER BY 1""",
                params,
            )
            out = {row[0]: (round(float(row[1]), 6), int(row[2])) for row in cur.fetchall()}
    finally:
        conn.close()
    return out


def tokens_by_day_model(since: dt.date | None = None) -> dict[str, dict[str, int]]:
    """{day: {model: output_tokens}} from LiteLLM's spend logs.

    OUTPUT tokens only — that is the unit GreenPT's published energy metric is
    normalised to ("tokens = computational distance"; input tokens are excluded
    so comparisons hold across prompts and models). Read from our own logs, so
    member-visible impact does not depend on any client persisting anything.
    """
    try:
        import psycopg2
    except Exception as exc:  # noqa: BLE001
        print(f"metrics: psycopg2 unavailable ({exc})")
        return {}
    try:
        conn = psycopg2.connect(
            host=os.environ.get("LITELLM_DB_HOST"),
            port=os.environ.get("LITELLM_DB_PORT"),
            user=os.environ.get("LITELLM_DB_USER"),
            password=os.environ.get("LITELLM_DB_PASSWORD"),
            dbname=os.environ.get("LITELLM_DB_NAME"),
            connect_timeout=10,
            sslmode="prefer",
        )
    except Exception as exc:  # noqa: BLE001
        print(f"metrics: LiteLLM DB connect failed ({exc}) — tokens left NULL")
        return {}
    try:
        where, params = "", []
        if since:
            where, params = 'WHERE "startTime" >= %s', [since.isoformat()]
        with conn.cursor() as cur:
            # model_group carries our provider/model alias (e.g.
            # greenpt/deepseek-v4.1-flash) — the name the factor table is keyed
            # on. Raw `model` holds the upstream name (openai/deepseek-v4.1-flash)
            # and would match nothing. Same convention as aggregate_model_usage.
            cur.execute(
                f"""SELECT to_char("startTime", 'YYYY-MM-DD') AS day,
                           COALESCE(NULLIF(model_group, ''), NULLIF(model, ''), 'unknown') AS alias,
                           COALESCE(SUM(completion_tokens), 0)
                    FROM "LiteLLM_SpendLogs" {where}
                    GROUP BY 1, 2""",
                params,
            )
            out: dict[str, dict[str, int]] = {}
            for day, alias, toks in cur.fetchall():
                out.setdefault(day, {})[alias or "unknown"] = int(toks)
    finally:
        conn.close()
    return out


def impact_for(models: dict[str, int]) -> tuple[float | None, float | None]:
    """(Wh, gCO2e) for a day's {model: output_tokens}.

    Factors live in IMPACT_FACTORS. Unknown models fall back to the documented
    proxy row rather than being skipped, so a total is never quietly missing a
    model — but the proxy is labelled as an order-of-magnitude estimate and the
    docs say so. Energy: Wh per 1,000 output tokens. Carbon: gCO2e per kWh.
    """
    if not models:
        return None, None
    if not sum(t for t in models.values() if t):
        # No measured tokens for this day: report "not measured" (NULL), never a
        # zero. A zero reads as "the co-op used no energy" — a lie that would
        # survive into a chart. The first backfill shipped exactly that, because
        # it summed the wrong column and every model totalled 0.
        return None, None
    wh = 0.0
    co2 = 0.0
    for model, toks in models.items():
        if not toks:
            continue
        f = None
        for prefix, factor in IMPACT_FACTORS.items():
            if model == prefix or model.startswith(prefix):
                f = factor
                break
        f = f or IMPACT_PROXY
        wh += toks / 1000.0 * f["wh_per_1k_out"]
        # Wh -> kWh is /1e3, then x gCO2e/kWh. This divided by 1e6 once, which
        # understated carbon exactly 1000x: the energy column stayed plausible
        # while carbon read ~0.2 g/day for a day using ~1 kWh.
        co2 += toks / 1000.0 * f["wh_per_1k_out"] / 1e3 * f["gco2e_per_kwh"]
    # Tripwire against unit errors: the blend of the factors above must imply a
    # grid intensity inside a sane band (the factors themselves span 24-250).
    if wh > 0:
        implied = co2 / wh * 1000.0  # gCO2e per kWh
        if not (5.0 <= implied <= 500.0):
            print(
                f"metrics: WARNING implied grid intensity {implied:.1f} gCO2e/kWh is outside "
                f"the plausible band — suspect a unit error before trusting these numbers"
            )
    return round(wh, 6), round(co2, 6)


# ---------------------------------------------------------------- writes

def upsert(day: str, row: dict) -> None:
    conn = connect()
    try:
        conn.execute(
            """INSERT INTO metrics_daily
                   (day, members_active, members_total, members_new, newsletter_subs,
                    matrix_space_members, spend_usd, requests, out_tokens, energy_wh,
                    co2e_g, recorded_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
               ON CONFLICT(day) DO UPDATE SET
                   members_active=excluded.members_active,
                   members_total=excluded.members_total,
                   members_new=excluded.members_new,
                   newsletter_subs=excluded.newsletter_subs,
                   matrix_space_members=excluded.matrix_space_members,
                   spend_usd=excluded.spend_usd,
                   requests=excluded.requests,
                   out_tokens=excluded.out_tokens,
                   energy_wh=excluded.energy_wh,
                   co2e_g=excluded.co2e_g,
                   recorded_at=excluded.recorded_at""",
            (
                day,
                row.get("members_active"),
                row.get("members_total"),
                row.get("members_new"),
                row.get("newsletter_subs"),
                row.get("matrix_space_members"),
                row.get("spend_usd"),
                row.get("requests"),
                row.get("out_tokens"),
                row.get("energy_wh"),
                row.get("co2e_g"),
                dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"),
            ),
        )
        conn.commit()
    finally:
        conn.close()


def snapshot(day: str | None = None) -> dict:
    day = day or today()
    m = from_sqlite()
    subs = from_listmonk()
    joins = listmonk_joins()
    space = from_matrix()
    since = dt.date.fromisoformat(day) - dt.timedelta(days=1)
    spend = spend_by_day(since=since)
    today_spend = spend.get(day, (None, None))
    day_models = (tokens_by_day_model(since=since).get(day)) or {}
    tot = sum(day_models.values())
    out_tokens = tot or None
    energy_wh, co2e_g = impact_for(day_models) if tot else (None, None)
    row = {
        "members_active": m["active"],
        "members_total": m["total"],
        "members_new": (sum(1 for j in joins if j == day) if joins is not None else None),
        "newsletter_subs": subs,
        "matrix_space_members": space,
        "spend_usd": today_spend[0],
        "requests": today_spend[1],
        "out_tokens": out_tokens,
        "energy_wh": energy_wh,
        "co2e_g": co2e_g,
    }
    upsert(day, row)
    print(f"metrics: recorded {day} -> {json.dumps(row)}")
    return row


def backfill() -> None:
    """Rebuild history: membership from list-join dates, usage from spend logs.

    Membership uses Listmonk's per-subscription `subscription_created_at` — the
    join date of every CURRENT member. Members who joined and later left are not
    in it (their departure is likewise absent), so past days are the cumulative
    joins of today's members: an honest floor for "how many of us were here by
    that day", not a reconstruction of the roster at the time.
    """
    m = from_sqlite()
    joins = sorted(listmonk_joins() or [])
    spend = spend_by_day()
    tokens = tokens_by_day_model()
    if not joins:
        print("metrics: no membership history (listmonk unavailable) — past days left NULL for members")
    if not joins and not spend:
        print("metrics: backfill found no history at all")
        return

    candidates = list(joins) + list(spend)
    first = min(candidates)
    last = today()
    print(
        f"metrics: backfilling {first}..{last} "
        f"({len(joins)} join dates, {len(spend)} spend days, {len(tokens)} token days)"
    )
    day = dt.date.fromisoformat(first)
    end = dt.date.fromisoformat(last)
    written = 0
    while day <= end:
        key = day.isoformat()
        # Cumulative members AS OF this day: current members whose list-join date
        # is on or before it. Days before anyone joined stay 0.
        members_total = sum(1 for j in joins if j <= key) if joins else None
        row = {
            "members_total": members_total,
            "members_new": sum(1 for j in joins if j == key),
            "spend_usd": spend.get(key, (None, None))[0],
            "requests": spend.get(key, (None, None))[1],
        }
        day_models = tokens.get(key) or {}
        if sum(day_models.values()):
            row["out_tokens"] = sum(day_models.values())
            row["energy_wh"], row["co2e_g"] = impact_for(day_models)
        if key == today():
            # today's row carries the exact live metrics
            row.update(
                members_active=m["active"],
                members_total=m["total"],
                newsletter_subs=from_listmonk(),
                matrix_space_members=from_matrix(),
            )
        upsert(key, row)
        written += 1
        day += dt.timedelta(days=1)
    print(f"metrics: backfill wrote {written} day(s)")


if __name__ == "__main__":
    if "--backfill" in sys.argv:
        backfill()
    else:
        snapshot()