backend/core/subscription/tracker.py

from datetime import datetime, timezone
from sqlalchemy.orm import Session
from core.subscription.db import SessionLocal
from core.subscription.models import UserUsage, User, UsageLog


def get_or_create_usage(db: Session, user_id: str) -> UserUsage:
    user = db.query(User).filter(User.clerk_id == user_id).first()
    if not user:
        user = User(clerk_id=user_id, tier="free")
        db.add(user)
        db.commit()

    usage = db.query(UserUsage).filter(UserUsage.user_id == user_id).first()
    if not usage:
        usage = UserUsage(user_id=user_id)
        db.add(usage)
        db.commit()
        db.refresh(usage)
    return usage


def check_and_reset_limits(usage: UserUsage, db: Session):
    now = datetime.now(timezone.utc)
    changed = False

    # Reset dzienny wizarda
    if usage.last_reset_daily.date() < now.date():
        usage.wizard_iterations_today = 0
        usage.last_reset_daily = now
        changed = True

    # Reset miesięczny tokenów
    if (
        usage.last_reset_monthly.month != now.month
        or usage.last_reset_monthly.year != now.year
    ):
        usage.tokens_used_month = 0
        usage.last_reset_monthly = now
        changed = True

    if changed:
        db.commit()


def increment_wizard_iteration(user_id: str) -> int:
    db = SessionLocal()
    try:
        usage = get_or_create_usage(db, user_id)
        check_and_reset_limits(usage, db)

        usage.wizard_iterations_today += 1

        # Log audytowy
        log = UsageLog(user_id=user_id, action_type="wizard_iteration", tokens_cost=0)
        db.add(log)

        db.commit()
        return usage.wizard_iterations_today
    finally:
        db.close()


def get_wizard_iterations(user_id: str) -> int:
    db = SessionLocal()
    try:
        usage = get_or_create_usage(db, user_id)
        check_and_reset_limits(usage, db)
        return usage.wizard_iterations_today
    finally:
        db.close()


def increment_tokens(user_id: str, tokens: int, action_type: str = "llm_call") -> int:
    db = SessionLocal()
    try:
        usage = get_or_create_usage(db, user_id)
        check_and_reset_limits(usage, db)

        usage.tokens_used_month += tokens

        # Log audytowy
        log = UsageLog(user_id=user_id, action_type=action_type, tokens_cost=tokens)
        db.add(log)

        db.commit()
        return usage.tokens_used_month
    finally:
        db.close()


def get_used_tokens(user_id: str) -> int:
    db = SessionLocal()
    try:
        usage = get_or_create_usage(db, user_id)
        check_and_reset_limits(usage, db)
        return usage.tokens_used_month
    finally:
        db.close()