rbox/mywebdav/billing/usage_tracker.py

import uuid
from datetime import datetime, date
from decimal import Decimal
from typing import Optional
from tortoise.transactions import in_transaction
from .models import UsageRecord, UsageAggregate
from ..models import User

class UsageTracker:
    @staticmethod
    async def track_storage(
        user: User,
        amount_bytes: int,
        resource_type: str = None,
        resource_id: int = None,
        metadata: dict = None
    ):
        idempotency_key = f"storage_{user.id}_{datetime.utcnow().timestamp()}_{uuid.uuid4().hex[:8]}"

        await UsageRecord.create(
            user=user,
            record_type="storage",
            amount_bytes=amount_bytes,
            resource_type=resource_type,
            resource_id=resource_id,
            idempotency_key=idempotency_key,
            metadata=metadata
        )

    @staticmethod
    async def track_bandwidth(
        user: User,
        amount_bytes: int,
        direction: str = "down",
        resource_type: str = None,
        resource_id: int = None,
        metadata: dict = None
    ):
        record_type = f"bandwidth_{direction}"
        idempotency_key = f"{record_type}_{user.id}_{datetime.utcnow().timestamp()}_{uuid.uuid4().hex[:8]}"

        await UsageRecord.create(
            user=user,
            record_type=record_type,
            amount_bytes=amount_bytes,
            resource_type=resource_type,
            resource_id=resource_id,
            idempotency_key=idempotency_key,
            metadata=metadata
        )

    @staticmethod
    async def aggregate_daily_usage(user: User, target_date: date = None):
        if target_date is None:
            target_date = date.today()

        start_of_day = datetime.combine(target_date, datetime.min.time())
        end_of_day = datetime.combine(target_date, datetime.max.time())

        storage_records = await UsageRecord.filter(
            user=user,
            record_type="storage",
            timestamp__gte=start_of_day,
            timestamp__lte=end_of_day
        ).all()

        storage_avg = sum(r.amount_bytes for r in storage_records) // max(len(storage_records), 1)
        storage_peak = max((r.amount_bytes for r in storage_records), default=0)

        bandwidth_up = await UsageRecord.filter(
            user=user,
            record_type="bandwidth_up",
            timestamp__gte=start_of_day,
            timestamp__lte=end_of_day
        ).all()

        bandwidth_down = await UsageRecord.filter(
            user=user,
            record_type="bandwidth_down",
            timestamp__gte=start_of_day,
            timestamp__lte=end_of_day
        ).all()

        total_up = sum(r.amount_bytes for r in bandwidth_up)
        total_down = sum(r.amount_bytes for r in bandwidth_down)

        async with in_transaction():
            aggregate, created = await UsageAggregate.get_or_create(
                user=user,
                date=target_date,
                defaults={
                    "storage_bytes_avg": storage_avg,
                    "storage_bytes_peak": storage_peak,
                    "bandwidth_up_bytes": total_up,
                    "bandwidth_down_bytes": total_down
                }
            )

            if not created:
                aggregate.storage_bytes_avg = storage_avg
                aggregate.storage_bytes_peak = storage_peak
                aggregate.bandwidth_up_bytes = total_up
                aggregate.bandwidth_down_bytes = total_down
                await aggregate.save()

        return aggregate

    @staticmethod
    async def get_current_storage(user: User) -> int:
        from ..models import File
        files = await File.filter(owner=user, is_deleted=False).all()
        return sum(f.size for f in files)

    @staticmethod
    async def get_monthly_usage(user: User, year: int, month: int) -> dict:
        from datetime import date
        from calendar import monthrange

        start_date = date(year, month, 1)
        _, last_day = monthrange(year, month)
        end_date = date(year, month, last_day)

        aggregates = await UsageAggregate.filter(
            user=user,
            date__gte=start_date,
            date__lte=end_date
        ).all()

        if not aggregates:
            return {
                "storage_gb_avg": 0,
                "storage_gb_peak": 0,
                "bandwidth_up_gb": 0,
                "bandwidth_down_gb": 0,
                "total_bandwidth_gb": 0
            }

        storage_avg = sum(a.storage_bytes_avg for a in aggregates) / len(aggregates)
        storage_peak = max(a.storage_bytes_peak for a in aggregates)
        bandwidth_up = sum(a.bandwidth_up_bytes for a in aggregates)
        bandwidth_down = sum(a.bandwidth_down_bytes for a in aggregates)

        return {
            "storage_gb_avg": round(storage_avg / (1024**3), 4),
            "storage_gb_peak": round(storage_peak / (1024**3), 4),
            "bandwidth_up_gb": round(bandwidth_up / (1024**3), 4),
            "bandwidth_down_gb": round(bandwidth_down / (1024**3), 4),
            "total_bandwidth_gb": round((bandwidth_up + bandwidth_down) / (1024**3), 4)
        }