helper/ratelimit.py

import asyncio
import hashlib
import os
import re
import time
from typing import Any, Dict, List, Optional

from fastapi import HTTPException, Request

from helper.subscriptions import (
    TIER_CONFIG,
    USAGE_PERIODS,
    fetch_subscription,
    normalize_plan_key,
    usage_locks,
    usage_store,
)


# -------------------------------------------------------------------
# Configuration / Constants
# -------------------------------------------------------------------

IDENTITY_CACHE_TTL_SECONDS = 60

CLIENT_BIND_TTL_SECONDS = int(
    os.getenv("CLIENT_BIND_TTL_SECONDS", str(8 * 24 * 60 * 60))
)

MAX_CLIENT_ID_LENGTH = 128

MAX_CHAT_PROMPT_CHARS = int(os.getenv("MAX_CHAT_PROMPT_CHARS", "120000"))
MAX_CHAT_PROMPT_BYTES = int(os.getenv("MAX_CHAT_PROMPT_BYTES", "500000"))

MAX_GROQ_PROMPT_CHARS = int(os.getenv("MAX_GROQ_PROMPT_CHARS", "90000"))
MAX_GROQ_PROMPT_BYTES = int(os.getenv("MAX_GROQ_PROMPT_BYTES", "350000"))

MAX_MEDIA_PROMPT_CHARS = int(os.getenv("MAX_MEDIA_PROMPT_CHARS", "4000"))
MAX_MEDIA_PROMPT_BYTES = int(os.getenv("MAX_MEDIA_PROMPT_BYTES", "16000"))


# -------------------------------------------------------------------
# In-memory state
# -------------------------------------------------------------------

identity_cache: Dict[str, Dict[str, Any]] = {}

client_subject_bindings: Dict[str, Dict[str, Any]] = {}


# -------------------------------------------------------------------
# Public Rate Limit APIs
# -------------------------------------------------------------------

async def resolve_rate_limit_identity(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
) -> tuple[str, str]:
    now = time.time()

    normalized_client_id = sanitize_client_id(client_id)
    default_subject = build_default_subject(request, normalized_client_id)

    if not authorization or not authorization.startswith("Bearer "):
        return "free", resolve_bound_subject(normalized_client_id, default_subject)

    token = authorization.split(" ", 1)[1].strip()
    if not token:
        return "free", resolve_bound_subject(normalized_client_id, default_subject)

    cached = identity_cache.get(token)
    if cached and cached.get("expires_at", 0) > now:
        plan_key = cached.get("plan_key", "free")
        subject = cached.get("subject", default_subject)

        bind_client_subject(normalized_client_id, subject, plan_key)
        return plan_key, subject

    try:
        sub = await fetch_subscription(token)
    except Exception:
        return "free", resolve_bound_subject(normalized_client_id, default_subject)

    if not isinstance(sub, dict) or sub.get("error"):
        return "free", resolve_bound_subject(normalized_client_id, default_subject)

    email = sub.get("email")
    if isinstance(email, str) and email.strip():
        subject = f"user:{email.strip().lower()}"
    else:
        subject = default_subject

    plan_key = normalize_plan_key(sub.get("plan_key"))

    identity_cache[token] = {
        "plan_key": plan_key,
        "subject": subject,
        "expires_at": now + IDENTITY_CACHE_TTL_SECONDS,
    }

    bind_client_subject(normalized_client_id, subject, plan_key)

    return plan_key, subject


async def enforce_rate_limit(
    request: Request,
    authorization: Optional[str],
    metric: str,
    client_id: Optional[str] = None,
) -> Dict[str, Optional[int | str]]:
    if metric not in usage_store:
        raise HTTPException(status_code=500, detail=f"Unknown limit metric: {metric}")

    plan_key, subject = await resolve_rate_limit_identity(
        request, authorization, client_id
    )

    plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"]
    plan_limits = plan.get("limits", {})
    limit = plan_limits.get(metric)

    window_key = get_usage_period_key(metric)

    lock = get_usage_lock(metric, subject)

    async with lock:
        bucket = usage_store[metric]

        entry = bucket.get(subject)

        if not entry or entry.get("window") != window_key:
            entry = {"window": window_key, "count": 0}
            bucket[subject] = entry

        if limit is not None and entry["count"] >= int(limit):
            raise HTTPException(
                status_code=429,
                detail=f"{metric} limit reached for {plan.get('name', 'current plan')}",
            )

        entry["count"] += 1

        remaining = None if limit is None else max(0, int(limit) - entry["count"])

        return {
            "plan_key": plan_key,
            "remaining": remaining,
            "used": entry["count"],
            "window": window_key,
        }


async def check_audio_rate_limit(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
):
    await enforce_rate_limit(request, authorization, "audioWeekly", client_id)


async def check_image_rate_limit(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
):
    await enforce_rate_limit(request, authorization, "imagesDaily", client_id)


async def check_video_rate_limit(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
):
    await enforce_rate_limit(request, authorization, "videosDaily", client_id)


# -------------------------------------------------------------------
# Prompt Utilities
# -------------------------------------------------------------------

def normalize_prompt_value(prompt: Optional[str], field_name: str = "prompt") -> str:
    if not isinstance(prompt, str):
        raise HTTPException(status_code=400, detail=f"{field_name} is required")

    normalized = prompt.strip()

    if not normalized:
        raise HTTPException(status_code=400, detail=f"{field_name} is required")

    return normalized


def enforce_prompt_size(prompt: str, max_chars: int, max_bytes: int, context: str):
    char_len = len(prompt)
    byte_len = len(prompt.encode("utf-8"))

    if char_len > max_chars or byte_len > max_bytes:
        raise HTTPException(
            status_code=413,
            detail=(
                f"{context} is too large ({char_len} chars, {byte_len} bytes). "
                f"Max allowed is {max_chars} chars or {max_bytes} bytes."
            ),
        )


def calculate_messages_size(messages: list) -> tuple[int, int]:
    total_chars = 0
    total_bytes = 0

    for message in messages:
        if not isinstance(message, dict):
            continue

        text = message_content_to_text(message.get("content"))

        if not text:
            continue

        total_chars += len(text)
        total_bytes += len(text.encode("utf-8"))

    return total_chars, total_bytes


def extract_user_text(messages: list) -> str:
    return " ".join(
        message_content_to_text(m.get("content"))
        for m in messages
        if isinstance(m, dict) and m.get("role") == "user"
    ).lower()


# -------------------------------------------------------------------
# Usage / Rate Limit Internals
# -------------------------------------------------------------------

def get_usage_period_key(metric: str) -> str:
    now = time.gmtime()

    period = USAGE_PERIODS.get(metric, "daily")

    if period == "weekly":
        iso_year, iso_week, _ = time.strftime("%G %V %u", now).split(" ")
        return f"{iso_year}-W{iso_week}"

    return time.strftime("%Y-%m-%d", now)


def get_usage_lock(metric: str, subject: str) -> asyncio.Lock:
    metric_locks = usage_locks.get(metric)

    if metric_locks is None:
        metric_locks = {}
        usage_locks[metric] = metric_locks

    lock = metric_locks.get(subject)

    if lock is None:
        lock = asyncio.Lock()
        metric_locks[subject] = lock

    return lock


def get_usage_snapshot_for_subject(
    plan_key: str, subject: str
) -> Dict[str, Dict[str, Any]]:
    plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"]

    plan_limits = plan.get("limits", {})

    snapshot: Dict[str, Dict[str, Any]] = {}

    for metric in usage_store.keys():
        limit = plan_limits.get(metric)

        window_key = get_usage_period_key(metric)

        entry = usage_store[metric].get(subject)

        used = 0

        if entry and entry.get("window") == window_key:
            used = max(0, int(entry.get("count", 0)))

        remaining = None if limit is None else max(0, int(limit) - used)

        snapshot[metric] = {
            "limit": limit,
            "used": used,
            "remaining": remaining,
            "window": window_key,
            "period": USAGE_PERIODS.get(metric, "daily"),
        }

    return snapshot


# -------------------------------------------------------------------
# Identity / Client helpers
# -------------------------------------------------------------------

def sanitize_client_id(raw_client_id: Optional[str]) -> Optional[str]:
    if not isinstance(raw_client_id, str):
        return None

    trimmed = raw_client_id.strip()

    if not trimmed or len(trimmed) > MAX_CLIENT_ID_LENGTH:
        return None

    if not re.match(r"^[A-Za-z0-9._:-]+$", trimmed):
        return None

    return trimmed


def build_default_subject(request: Request, client_id: Optional[str]) -> str:
    if client_id:
        client_hash = hashlib.sha256(client_id.encode("utf-8")).hexdigest()[:24]
        return f"client:{client_hash}"

    host = request.client.host if request.client else "unknown"

    user_agent = request.headers.get("user-agent", "")

    ua_hash = (
        hashlib.sha256(user_agent.encode("utf-8")).hexdigest()[:12]
        if user_agent
        else "noua"
    )

    return f"anon:{host}:{ua_hash}"


def bind_client_subject(client_id: Optional[str], subject: str, plan_key: str):
    if not client_id:
        return

    client_subject_bindings[client_id] = {
        "subject": subject,
        "plan_key": plan_key,
        "expires_at": time.time() + CLIENT_BIND_TTL_SECONDS,
    }


def resolve_bound_subject(client_id: Optional[str], fallback_subject: str) -> str:
    if not client_id:
        return fallback_subject

    bound = client_subject_bindings.get(client_id)

    if not bound:
        return fallback_subject

    if bound.get("expires_at", 0) <= time.time():
        client_subject_bindings.pop(client_id, None)
        return fallback_subject

    return bound.get("subject", fallback_subject)


# -------------------------------------------------------------------
# Message parsing helpers
# -------------------------------------------------------------------

def message_content_to_text(content: Any) -> str:
    if isinstance(content, str):
        return content

    if isinstance(content, list):
        parts: List[str] = []

        for item in content:
            if isinstance(item, str):
                parts.append(item)
                continue

            if isinstance(item, dict):
                text = item.get("text")

                if isinstance(text, str):
                    parts.append(text)

        return " ".join(parts)

    return ""