agent/token_usage.py

"""
NeuraPrompt — Token Usage Tracker (v1.0)
=========================================

Monthly token budget per user:
  Pro (Polar subscriber): 1,000,000 tokens (input + output combined)
  Free:                   25,000 tokens

Resets on the 1st of each UTC month.
MongoDB-backed (survives restarts + works across replicas). Falls back to
in-memory if Mongo is unavailable.

Document shape (in `neuraprompt.token_usage` collection):
  {
    _id: ObjectId,
    user_id: "firebase_uid",
    month_key: "2026-06",            # YYYY-MM UTC
    input_tokens: 12345,
    output_tokens: 6789,
    total_tokens: 19134,
    last_updated: ISODate,
    last_model: "poolside/laguna-m.1:free"
  }

Public API:
  get_usage(user_id) -> dict
  record_usage(user_id, input_tokens, output_tokens, model) -> dict
  check_limit(user_id, is_subscribed) -> (allowed, reason, limit, remaining)
  get_tier_limit(is_subscribed) -> int
  reset_user_usage(user_id) -> bool    # admin override

Env vars:
  MONGO_URI              — MongoDB connection string (required for persistence)
  TOKEN_PRO_LIMIT        — override pro limit (default 1000000)
  TOKEN_FREE_LIMIT       — override free limit (default 25000)
"""

from __future__ import annotations

import os
import time
import logging
from datetime import datetime, timezone
from typing import Optional, Dict, Any, Tuple

log = logging.getLogger("token_usage")

# ─────────────────────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────────────────────
MONTHLY_PRO_LIMIT  = int(os.getenv("TOKEN_PRO_LIMIT",  "10_000"))
MONTHLY_FREE_LIMIT = int(os.getenv("TOKEN_FREE_LIMIT", "1_000"))

MONGO_URI = os.getenv("MONGO_URI", "")

# ─────────────────────────────────────────────────────────────
# MONGO CONNECTION (lazy — reuses main.py's client if available)
# ─────────────────────────────────────────────────────────────
_token_col = None
_token_col_checked = False

def _get_collection():
    """Lazy resolver for the token_usage MongoDB collection.
    Reuses main.py's MongoClient if available (avoids opening a second connection).
    Falls back to in-memory if Mongo is unreachable."""
    global _token_col, _token_col_checked
    if _token_col_checked:
        return _token_col
    _token_col_checked = True

    # Try to reuse main.py's already-connected MongoClient
    try:
        import sys
        main = sys.modules.get("main")
        if main is not None and hasattr(main, "neuraprompt_db"):
            _token_col = main.neuraprompt_db["token_usage"]
            try:
                _token_col.create_index([("user_id", 1), ("month_key", 1)], unique=True)
            except Exception:
                pass  # index may already exist
            log.info("[TokenUsage] Reusing main.py's MongoDB connection.")
            return _token_col
    except Exception as e:
        log.debug(f"[TokenUsage] Could not reuse main.py's Mongo: {e}")

    # Open our own connection
    if not MONGO_URI:
        log.warning("[TokenUsage] MONGO_URI not set — using in-memory fallback (resets on restart!).")
        _token_col = None
        return None
    try:
        from pymongo import MongoClient
        from pymongo.server_api import ServerApi
        client = MongoClient(
            MONGO_URI, ssl=True,
            tlsAllowInvalidCertificates=False,
            tlsCAFile="/etc/ssl/certs/ca-certificates.crt",
            server_api=ServerApi("1"),
        )
        db = client["neuraprompt"]
        _token_col = db["token_usage"]
        _token_col.create_index([("user_id", 1), ("month_key", 1)], unique=True)
        log.info("[TokenUsage] MongoDB connected (own connection).")
        return _token_col
    except Exception as e:
        log.warning(f"[TokenUsage] Mongo connection failed — using in-memory: {e}")
        _token_col = None
        return None


# In-memory fallback (used only if Mongo is unavailable)
_memory_usage: Dict[str, dict] = {}


# ─────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────

def _month_key() -> str:
    """Current UTC month as YYYY-MM (e.g. '2026-06').
    All limits reset when this key rolls over."""
    return datetime.now(timezone.utc).strftime("%Y-%m")


def _estimate_tokens(text: str) -> int:
    """Rough token estimate when the API doesn't return usage data.
    4 chars ≈ 1 token. Over-estimates non-Latin text (safer — denies earlier)."""
    return max(1, len(text or "") // 4)


def get_tier_limit(is_subscribed: bool) -> int:
    """Monthly token limit for the user's tier."""
    return MONTHLY_PRO_LIMIT if is_subscribed else MONTHLY_FREE_LIMIT


def get_tier_name(is_subscribed: bool) -> str:
    return "pro" if is_subscribed else "free"


# ─────────────────────────────────────────────────────────────
# PUBLIC API
# ─────────────────────────────────────────────────────────────

def get_usage(user_id: str) -> dict:
    """Returns the user's token usage for the CURRENT month.
    Shape: {input_tokens, output_tokens, total_tokens, month_key}"""
    month = _month_key()
    col = _get_collection()
    if col is not None:
        try:
            doc = col.find_one({"user_id": user_id, "month_key": month})
            if doc:
                return {
                    "input_tokens":  int(doc.get("input_tokens", 0)),
                    "output_tokens": int(doc.get("output_tokens", 0)),
                    "total_tokens":  int(doc.get("total_tokens", 0)),
                    "month_key":     month,
                }
        except Exception as e:
            log.warning(f"[TokenUsage] get_usage Mongo error: {e}")
    # In-memory fallback
    key = f"{user_id}|{month}"
    doc = _memory_usage.get(key, {})
    return {
        "input_tokens":  int(doc.get("input_tokens", 0)),
        "output_tokens": int(doc.get("output_tokens", 0)),
        "total_tokens":  int(doc.get("total_tokens", 0)),
        "month_key":     month,
    }


def record_usage(user_id: str, input_tokens: int, output_tokens: int,
                 model: str = "") -> dict:
    """Add tokens to the user's monthly total. Returns the updated usage.
    `input_tokens` + `output_tokens` are the actual counts from the API
    response (usage.prompt_tokens / usage.completion_tokens)."""
    if input_tokens < 0:
        input_tokens = 0
    if output_tokens < 0:
        output_tokens = 0
    total = input_tokens + output_tokens
    month = _month_key()
    now = datetime.now(timezone.utc)

    col = _get_collection()
    if col is not None:
        try:
            col.update_one(
                {"user_id": user_id, "month_key": month},
                {
                    "$inc": {
                        "input_tokens":  input_tokens,
                        "output_tokens": output_tokens,
                        "total_tokens":  total,
                    },
                    "$set": {
                        "last_updated": now,
                        "last_model":   model,
                    },
                },
                upsert=True,
            )
        except Exception as e:
            log.warning(f"[TokenUsage] record_usage Mongo error (using memory): {e}")
            # Fall through to in-memory
            key = f"{user_id}|{month}"
            doc = _memory_usage.setdefault(key, {
                "input_tokens": 0, "output_tokens": 0, "total_tokens": 0,
            })
            doc["input_tokens"]  += input_tokens
            doc["output_tokens"] += output_tokens
            doc["total_tokens"]  += total
            doc["last_updated"]   = now
            doc["last_model"]     = model
    else:
        key = f"{user_id}|{month}"
        doc = _memory_usage.setdefault(key, {
            "input_tokens": 0, "output_tokens": 0, "total_tokens": 0,
        })
        doc["input_tokens"]  += input_tokens
        doc["output_tokens"] += output_tokens
        doc["total_tokens"]  += total
        doc["last_updated"]   = now
        doc["last_model"]     = model

    return get_usage(user_id)


def check_limit(user_id: str, is_subscribed: bool) -> Tuple[bool, str, int, int]:
    """Check if the user has any remaining token budget for this month.
    Returns (allowed, reason, limit, remaining).
      allowed   — True if remaining > 0
      reason    — "" if allowed, else a short code: "monthly_limit_reached"
      limit     — the monthly limit for the user's tier
      remaining — tokens remaining (>= 0)"""
    usage = get_usage(user_id)
    limit = get_tier_limit(is_subscribed)
    remaining = limit - usage["total_tokens"]
    if remaining <= 0:
        return False, "monthly_limit_reached", limit, 0
    return True, "", limit, remaining


def check_and_reserve(user_id: str, is_subscribed: bool,
                       estimated_input_tokens: int) -> Tuple[bool, str, dict]:
    """Pre-flight check before an LLM call.
    Returns (allowed, reason, info_dict).
      allowed — True if the user has enough budget for the estimated input
      reason  — "" if allowed, else: "monthly_limit_reached" | "insufficient_for_input"
      info    — {limit, used, remaining, requested, tier}
    """
    usage = get_usage(user_id)
    limit = get_tier_limit(is_subscribed)
    used = usage["total_tokens"]
    remaining = limit - used
    tier = get_tier_name(is_subscribed)

    if remaining <= 0:
        return False, "monthly_limit_reached", {
            "limit": limit, "used": used, "remaining": 0,
            "requested": estimated_input_tokens, "tier": tier,
        }
    if estimated_input_tokens > remaining:
        return False, "insufficient_for_input", {
            "limit": limit, "used": used, "remaining": remaining,
            "requested": estimated_input_tokens, "tier": tier,
        }
    return True, "", {
        "limit": limit, "used": used, "remaining": remaining,
        "requested": estimated_input_tokens, "tier": tier,
    }


def reset_user_usage(user_id: str) -> bool:
    """Admin override: zero out a user's usage for the current month.
    Returns True if successful."""
    month = _month_key()
    col = _get_collection()
    if col is not None:
        try:
            col.update_one(
                {"user_id": user_id, "month_key": month},
                {"$set": {
                    "input_tokens": 0, "output_tokens": 0, "total_tokens": 0,
                    "last_updated": datetime.now(timezone.utc),
                    "reset_at": datetime.now(timezone.utc),
                }},
                upsert=True,
            )
            return True
        except Exception as e:
            log.error(f"[TokenUsage] reset_user_usage failed: {e}")
            return False
    else:
        key = f"{user_id}|{month}"
        _memory_usage[key] = {
            "input_tokens": 0, "output_tokens": 0, "total_tokens": 0,
        }
        return True


def get_usage_summary(user_id: str, is_subscribed: bool) -> dict:
    """Full summary for the /kype/usage endpoint."""
    usage = get_usage(user_id)
    limit = get_tier_limit(is_subscribed)
    return {
        "user_id":         user_id,
        "tier":            get_tier_name(is_subscribed),
        "subscribed":      is_subscribed,
        "month_key":       usage["month_key"],
        "input_tokens":    usage["input_tokens"],
        "output_tokens":   usage["output_tokens"],
        "total_used":      usage["total_tokens"],
        "monthly_limit":   limit,
        "remaining":       max(0, limit - usage["total_tokens"]),
        "resets_on":       f"{usage['month_key']}-01 → next month (UTC 1st)",
        "pro_limit":       MONTHLY_PRO_LIMIT,
        "free_limit":      MONTHLY_FREE_LIMIT,
    }


# ─────────────────────────────────────────────────────────────
# FASTAPI ROUTER (optional — mount on your app to expose usage stats)
# ─────────────────────────────────────────────────────────────

try:
    from fastapi import APIRouter, HTTPException, Query
    token_router = APIRouter(prefix="/tokens")

    @token_router.get("/usage")
    async def get_my_usage(user_id: str = Query(...), subscribed: bool = Query(False)):
        """Get the caller's token usage for the current month.
        Frontend should pass `subscribed` based on the Polar check result,
        OR the backend can resolve it from Polar directly (preferred)."""
        return get_usage_summary(user_id, is_subscribed=subscribed)

    @token_router.get("/health")
    async def token_health():
        col = _get_collection()
        return {
            "status":           "ok",
            "storage":          "mongodb" if col is not None else "in-memory",
            "pro_limit":        MONTHLY_PRO_LIMIT,
            "free_limit":       MONTHLY_FREE_LIMIT,
            "current_month":    _month_key(),
        }

    @token_router.post("/admin/reset")
    async def admin_reset(user_id: str = Query(...)):
        """Admin-only: reset a user's monthly usage. Protect this route with
        your admin auth middleware."""
        ok = reset_user_usage(user_id)
        return {"reset": ok, "user_id": user_id, "month": _month_key()}

except ImportError:
    # FastAPI not available — module used as a library only
    pass