"""Folder quota helpers."""
from __future__ import annotations

from datetime import datetime, timezone

import pytest
from sqlalchemy import select

from app.billing.quota import (
    check_folder_quota,
    add_token_usage,
    QuotaExceeded,
)
from app.models import MonthlyTokenUsage


pytestmark = pytest.mark.asyncio


async def test_check_folder_quota_free_rejects_above_file_cap(db, test_user_free):
    with pytest.raises(QuotaExceeded) as exc:
        await check_folder_quota(
            user_id=test_user_free.id, tier="free", estimated_files=500, db=db
        )
    assert exc.value.reason == "max_files"


async def test_check_folder_quota_free_passes_under_cap(db, test_user_free):
    # No raise
    await check_folder_quota(
        user_id=test_user_free.id, tier="free", estimated_files=50, db=db
    )


async def test_check_folder_quota_rejects_when_monthly_exhausted(db, test_user_free):
    ym = datetime.now(timezone.utc).strftime("%Y-%m")
    db.add(MonthlyTokenUsage(
        user_id=test_user_free.id, year_month=ym, feature="folder_index", tokens_used=100_000
    ))
    await db.commit()
    with pytest.raises(QuotaExceeded) as exc:
        await check_folder_quota(
            user_id=test_user_free.id, tier="free", estimated_files=10, db=db
        )
    assert exc.value.reason == "monthly_tokens"


async def test_check_folder_quota_power_unlimited(db, test_user_power):
    await check_folder_quota(
        user_id=test_user_power.id, tier="power", estimated_files=999_999, db=db
    )


async def test_add_token_usage_atomic_increment(db, test_user_free):
    await add_token_usage(user_id=test_user_free.id, feature="folder_index", tokens=1500, db=db)
    await add_token_usage(user_id=test_user_free.id, feature="folder_index", tokens=2500, db=db)
    ym = datetime.now(timezone.utc).strftime("%Y-%m")
    row = (await db.execute(
        select(MonthlyTokenUsage).where(
            MonthlyTokenUsage.user_id == test_user_free.id,
            MonthlyTokenUsage.year_month == ym,
            MonthlyTokenUsage.feature == "folder_index",
        )
    )).scalar_one()
    assert row.tokens_used == 4000


async def test_add_token_usage_returns_exhausted_when_over_cap(db, test_user_free):
    result = await add_token_usage(
        user_id=test_user_free.id, feature="folder_index", tokens=150_000, db=db, cap=100_000
    )
    assert result.exhausted is True
    assert result.tokens_used == 150_000