agent/token_usage.py

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

Monthly token budget per user (input + output combined):
  Free:                   200,000 tokens  (+ 5 daily limits — see below)
  Pro (Polar Pro):        1,000,000 tokens
  Max (Polar Max/Ultra):  1,500,000 tokens

DAILY LIMITS (Free tier only):
  The free plan has 5 daily limits per UTC day. Each daily limit is
  DAILY_UNIT_TOKENS (40,000) of usage, so 5 units = 200,000 tokens = the
  entire monthly budget in a single day. When all 5 are burned, the user
  is blocked for the rest of the UTC day with a "daily limit reached"
  message (shown INSTEAD of the monthly exhausted message).

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 200000)
  TOKEN_MAX_LIMIT        — override max limit (default 1500000)
"""

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",  "1000000"))
MONTHLY_FREE_LIMIT = int(os.getenv("TOKEN_FREE_LIMIT", "200000"))
MONTHLY_MAX_LIMIT  = int(os.getenv("TOKEN_MAX_LIMIT",  "1500000"))

# ── Daily limits (Free tier only): 5 units x 40K = 200K/day cap ──
DAILY_UNITS_PER_DAY   = int(os.getenv("DAILY_UNITS_PER_DAY",   "5"))
DAILY_UNIT_TOKENS     = int(os.getenv("DAILY_UNIT_TOKENS",     str(MONTHLY_FREE_LIMIT // DAILY_UNITS_PER_DAY)))
DAILY_FREE_TOKEN_LIMIT = DAILY_UNIT_TOKENS * DAILY_UNITS_PER_DAY   # 200,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 _resolve_tier(is_subscribed: bool, plan: str = "") -> str:
    """Map (subscribed, plan) to a tier name: free | pro | max.
    Polar plan names: 'max' / 'ultra' / 'premium' map to the Max tier."""
    p = (plan or "").strip().lower()
    if p in ("max", "ultra", "premium"):
        return "max"
    if is_subscribed or p == "pro":
        return "pro"
    return "free"


def get_tier_limit(is_subscribed: bool, plan: str = "") -> int:
    """Monthly token limit for the user's tier (Free 200K / Pro 1M / Max 1.5M)."""
    tier = _resolve_tier(is_subscribed, plan)
    if tier == "max":
        return MONTHLY_MAX_LIMIT
    if tier == "pro":
        return MONTHLY_PRO_LIMIT
    return MONTHLY_FREE_LIMIT


def get_tier_name(is_subscribed: bool, plan: str = "") -> str:
    return _resolve_tier(is_subscribed, plan)


def get_tier_label(is_subscribed: bool, plan: str = "") -> str:
    """Human label with the real limit, e.g. 'Free (200,000/month)'.
    NOTE: never hardcode 'Free (25K/month)' — that caused mismatched
    limit messages in the past."""
    tier = _resolve_tier(is_subscribed, plan)
    return f"{'Free' if tier == 'free' else tier.capitalize()} ({get_tier_limit(is_subscribed, plan):,}/month)"


# ─────────────────────────────────────────────────────────────
# DAILY USAGE (Free tier — 5 daily limits x 40K tokens)
# ─────────────────────────────────────────────────────────────
_memory_daily: Dict[str, dict] = {}

def _get_daily_collection():
    """Lazy resolver for the daily token_usage collection (separate from
    monthly so both indexes stay clean). Reuses main.py's client if present."""
    try:
        import sys
        main = sys.modules.get("main")
        if main is not None and hasattr(main, "neuraprompt_db"):
            col = main.neuraprompt_db["token_usage_daily"]
            try:
                col.create_index([("user_id", 1), ("day_key", 1)], unique=True)
            except Exception:
                pass
            return col
    except Exception as e:
        log.debug(f"[TokenUsage] daily: could not reuse main.py Mongo: {e}")

    if not MONGO_URI:
        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"),
        )
        col = client["neuraprompt"]["token_usage_daily"]
        col.create_index([("user_id", 1), ("day_key", 1)], unique=True)
        return col
    except Exception as e:
        log.warning(f"[TokenUsage] daily Mongo connection failed: {e}")
        return None


def _day_key() -> str:
    """Current UTC day as YYYY-MM-DD. Daily limits reset when this rolls over."""
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def get_daily_usage(user_id: str) -> dict:
    """Token usage for the CURRENT UTC day. Shape: {total_tokens, day_key}"""
    day = _day_key()
    col = _get_daily_collection()
    if col is not None:
        try:
            doc = col.find_one({"user_id": user_id, "day_key": day})
            if doc:
                return {"total_tokens": int(doc.get("total_tokens", 0)), "day_key": day}
        except Exception as e:
            log.warning(f"[TokenUsage] get_daily_usage Mongo error: {e}")
    key = f"{user_id}|{day}"
    doc = _memory_daily.get(key, {})
    return {"total_tokens": int(doc.get("total_tokens", 0)), "day_key": day}


def _record_daily(user_id: str, total: int) -> None:
    """Increment the user's daily total (in Mongo + memory fallback)."""
    day = _day_key()
    col = _get_daily_collection()
    written = False
    if col is not None:
        try:
            col.update_one(
                {"user_id": user_id, "day_key": day},
                {"$inc": {"total_tokens": total},
                 "$set": {"last_updated": datetime.now(timezone.utc)}},
                upsert=True,
            )
            written = True
        except Exception as e:
            log.warning(f"[TokenUsage] _record_daily Mongo error: {e}")
    if not written:
        key = f"{user_id}|{day}"
        doc = _memory_daily.setdefault(key, {"total_tokens": 0})
        doc["total_tokens"] += total


def daily_units_info(user_id: str) -> dict:
    """How many of the 5 daily limits the user has burned today (Free tier)."""
    used = get_daily_usage(user_id)["total_tokens"]
    units_used = 0
    if used > 0:
        units_used = min(DAILY_UNITS_PER_DAY, -(-used // DAILY_UNIT_TOKENS))  # ceil
    return {
        "daily_used":       used,
        "daily_limit":      DAILY_FREE_TOKEN_LIMIT,
        "units_used":       units_used,
        "units_remaining":  DAILY_UNITS_PER_DAY - units_used,
        "units_per_day":    DAILY_UNITS_PER_DAY,
        "unit_tokens":      DAILY_UNIT_TOKENS,
        "day_key":          _day_key(),
    }


# ─────────────────────────────────────────────────────────────
# 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 = "", reserved_input: int = 0) -> dict:
    """Add tokens to the user's monthly total. Returns the updated usage.

    Args:
        reserved_input: Amount to release from reservation (from check_and_reserve).
                       If 0, attempts to auto-release based on input_tokens."""

    # Release reservation first
    release = reserved_input or input_tokens
    release_reservation(user_id, release)
    """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)

    # Release the reservation (actual usage may differ from estimate)
    if reserved_input > 0:
        release_reservation(user_id, reserved_input)

    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

    # Daily tracking (used by the Free tier's 5-daily-limits system)
    _record_daily(user_id, total)

    return get_usage(user_id)


def check_limit(user_id: str, is_subscribed: bool, plan: str = "") -> Tuple[bool, str, int, int]:
    """Check if the user has any remaining token budget.
    Returns (allowed, reason, limit, remaining).
      reason — "" | "monthly_limit_reached" | "daily_limit_reached"
    Free tier ALSO checks the 5-daily-limits system (5 x 40K = 200K/day)."""
    tier = _resolve_tier(is_subscribed, plan)
    usage = get_usage(user_id)
    limit = get_tier_limit(is_subscribed, plan)
    remaining = limit - usage["total_tokens"]
    if remaining <= 0:
        return False, "monthly_limit_reached", limit, 0
    if tier == "free":
        daily = get_daily_usage(user_id)["total_tokens"]
        if daily >= DAILY_FREE_TOKEN_LIMIT:
            return False, "daily_limit_reached", limit, 0
    return True, "", limit, remaining


# In-memory reservation tracking (production: use Redis with TTL)
_reserved_tokens: Dict[str, int] = {}

def check_and_reserve(user_id: str, is_subscribed: bool,
                       estimated_input_tokens: int, plan: str = "") -> Tuple[bool, str, dict]:
    """Pre-flight check + ATOMIC reservation before an LLM call.

    This function BOTH checks the limit AND reserves the estimated tokens
    so parallel requests cannot all slip through. Reserved tokens are
    released and actual usage is recorded by record_usage() after the call.

    Free tier ALSO enforces the 5-daily-limits system: 5 units of
    DAILY_UNIT_TOKENS (40K) per UTC day = 200K/day cap. When burned out,
    reason is "daily_limit_reached" (callers show that INSTEAD of the
    monthly exhausted message).

    Returns (allowed, reason, info_dict).
      allowed — True if the user has enough budget
      reason  — "" if allowed, else:
                "monthly_limit_reached" | "daily_limit_reached" | "insufficient_for_input"
      info    — {limit, used, remaining, requested, tier, reserved,
                 tier_label, month_key, + daily_* fields for free tier}
    """
    usage = get_usage(user_id)
    tier = _resolve_tier(is_subscribed, plan)
    limit = get_tier_limit(is_subscribed, plan)
    used = usage["total_tokens"]
    month_key = usage.get("month_key", _month_key())

    # Include reserved tokens in the calculation (prevents parallel over-spend)
    reserved = _reserved_tokens.get(user_id, 0)
    effective_used = used + reserved
    remaining = limit - effective_used

    base_info = {
        "limit": limit, "used": used, "remaining": max(0, remaining),
        "requested": estimated_input_tokens, "tier": tier, "reserved": reserved,
        "tier_label": get_tier_label(is_subscribed, plan),
        "month_key": month_key,
    }

    # ── Daily limit check (Free tier only) ──
    if tier == "free":
        daily = daily_units_info(user_id)
        daily_info = {
            "daily_used":      daily["daily_used"],
            "daily_limit":     daily["daily_limit"],
            "units_used":      daily["units_used"],
            "units_remaining": daily["units_remaining"],
            "units_per_day":   daily["units_per_day"],
            "unit_tokens":     daily["unit_tokens"],
        }
        base_info.update(daily_info)
        if daily["daily_used"] >= DAILY_FREE_TOKEN_LIMIT:
            return False, "daily_limit_reached", base_info

    if remaining <= 0:
        return False, "monthly_limit_reached", base_info
    if estimated_input_tokens > remaining:
        return False, "insufficient_for_input", base_info

    # RESERVE the tokens atomically
    _reserved_tokens[user_id] = reserved + estimated_input_tokens
    base_info["reserved"] = reserved + estimated_input_tokens
    base_info["remaining"] = remaining - estimated_input_tokens

    return True, "", base_info


def release_reservation(user_id: str, reserved_amount: int) -> None:
    """Release reserved tokens (call if the LLM call fails or is cancelled)."""
    current = _reserved_tokens.get(user_id, 0)
    new_val = max(0, current - reserved_amount)
    if new_val == 0:
        _reserved_tokens.pop(user_id, None)
    else:
        _reserved_tokens[user_id] = new_val


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, plan: str = "") -> dict:
    """Full summary for the /kype/usage and /agent usage endpoints."""
    usage = get_usage(user_id)
    tier = _resolve_tier(is_subscribed, plan)
    limit = get_tier_limit(is_subscribed, plan)
    out = {
        "user_id":         user_id,
        "tier":            tier,
        "tier_label":      get_tier_label(is_subscribed, plan),
        "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)",
        "free_limit":      MONTHLY_FREE_LIMIT,
        "pro_limit":       MONTHLY_PRO_LIMIT,
        "max_limit":       MONTHLY_MAX_LIMIT,
    }
    if tier == "free":
        out["daily"] = daily_units_info(user_id)
    return out


# ─────────────────────────────────────────────────────────────
# 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),
                           plan: str = Query("")):
        """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, plan=plan)

    @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,
            "max_limit":        MONTHLY_MAX_LIMIT,
            "daily_unit_tokens": DAILY_UNIT_TOKENS,
            "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