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

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 = "", 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

    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


# 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) -> 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.

    Returns (allowed, reason, info_dict).
      allowed — True if the user has enough budget
      reason  — "" if allowed, else: "monthly_limit_reached" | "insufficient_for_input"
      info    — {limit, used, remaining, requested, tier, reserved}
    """
    usage = get_usage(user_id)
    limit = get_tier_limit(is_subscribed)
    used = usage["total_tokens"]

    # 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
    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, "reserved": reserved,
        }
    if estimated_input_tokens > remaining:
        return False, "insufficient_for_input", {
            "limit": limit, "used": used, "remaining": remaining,
            "requested": estimated_input_tokens, "tier": tier, "reserved": reserved,
        }

    # RESERVE the tokens atomically
    _reserved_tokens[user_id] = reserved + estimated_input_tokens

    return True, "", {
        "limit": limit, "used": used, "remaining": remaining - estimated_input_tokens,
        "requested": estimated_input_tokens, "tier": tier, "reserved": reserved + estimated_input_tokens,
    }


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) -> 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