rate_limiter.py

"""Rate limiting service for chat API.

[Task]: T021
[From]: specs/004-ai-chatbot/tasks.md

This service enforces the 100 messages/day limit per user (NFR-011).
"""
import uuid
from datetime import datetime, timedelta
from typing import Optional
from sqlmodel import Session, select
from sqlalchemy import func

from models.message import Message


# Rate limit constants
DAILY_MESSAGE_LIMIT = 100  # NFR-011: Maximum messages per user per day


def check_rate_limit(
    db: Session,
    user_id: uuid.UUID
) -> tuple[bool, int, Optional[datetime]]:
    """Check if user has exceeded their daily message limit.

    [From]: specs/004-ai-chatbot/spec.md - NFR-011

    Args:
        db: Database session (synchronous)
        user_id: User ID to check

    Returns:
        Tuple of (allowed, remaining_count, reset_time)
        - allowed: True if user can send message, False if limit exceeded
        - remaining_count: Number of messages remaining today
        - reset_time: When the limit resets (midnight UTC), or None if allowed

    Example:
        >>> allowed, remaining, reset = await check_rate_limit(db, user_id)
        >>> if not allowed:
        ...     print(f"Rate limited. Resets at {reset}")
    """
    # Calculate today's date range (UTC)
    now = datetime.utcnow()
    today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
    today_end = today_start + timedelta(days=1)

    # Count messages sent by user today
    # [From]: specs/004-ai-chatbot/spec.md - NFR-011
    # Count both user and assistant messages (all messages in conversation)
    statement = select(func.count(Message.id)).where(
        Message.user_id == user_id,
        Message.created_at >= today_start,
        Message.created_at < today_end
    )

    message_count = db.exec(statement).one() or 0

    # Calculate remaining messages
    remaining = DAILY_MESSAGE_LIMIT - message_count

    if remaining <= 0:
        # Rate limit exceeded
        return False, 0, today_end
    else:
        # User can send message
        return True, remaining - 1, None


def record_message(
    db: Session,
    user_id: uuid.UUID,
    conversation_id: uuid.UUID,
    role: str,
    content: str
) -> Message:
    """Record a message in the database (for rate limit tracking).

    [From]: specs/004-ai-chatbot/plan.md - Message Persistence

    Note: This function is primarily for rate limit tracking.
    The actual message persistence should happen in the chat API endpoint
    before AI processing (T017) and after AI response (T018).

    Args:
        db: Database session
        user_id: User ID who sent the message
        conversation_id: Conversation ID
        role: Message role ("user" or "assistant")
        content: Message content

    Returns:
        Created message object
    """
    message = Message(
        id=uuid.uuid4(),
        conversation_id=conversation_id,
        user_id=user_id,
        role=role,
        content=content,
        created_at=datetime.utcnow()
    )

    db.add(message)
    db.commit()
    db.refresh(message)

    return message


def get_message_count_today(
    db: Session,
    user_id: uuid.UUID
) -> int:
    """Get the number of messages sent by user today.

    [From]: specs/004-ai-chatbot/spec.md - NFR-011

    Args:
        db: Database session
        user_id: User ID to check

    Returns:
        Number of messages sent today (both user and assistant)
    """
    now = datetime.utcnow()
    today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
    today_end = today_start + timedelta(days=1)

    statement = select(func.count(Message.id)).where(
        Message.user_id == user_id,
        Message.created_at >= today_start,
        Message.created_at < today_end
    )

    return db.exec(statement).one() or 0


def get_rate_limit_status(
    db: Session,
    user_id: uuid.UUID
) -> dict:
    """Get comprehensive rate limit status for a user.

    [From]: specs/004-ai-chatbot/spec.md - NFR-011

    Args:
        db: Database session
        user_id: User ID to check

    Returns:
        Dictionary with rate limit information:
        {
            "limit": 100,
            "used": 45,
            "remaining": 55,
            "resets_at": "2025-01-16T00:00:00Z"
        }
    """
    now = datetime.utcnow()
    today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
    today_end = today_start + timedelta(days=1)

    # Count messages sent today
    statement = select(func.count(Message.id)).where(
        Message.user_id == user_id,
        Message.created_at >= today_start,
        Message.created_at < today_end
    )

    message_count = db.exec(statement).one() or 0

    remaining = max(0, DAILY_MESSAGE_LIMIT - message_count)

    return {
        "limit": DAILY_MESSAGE_LIMIT,
        "used": message_count,
        "remaining": remaining,
        "resets_at": today_end.isoformat() + "Z"
    }