"""Quota-safe monthly claims for organization lookups in Checko."""

from datetime import date, datetime

from apps.parsers.models import CheckoCollectionAttempt
from django.utils import timezone


def collection_period_month(at: date | datetime | None = None) -> date:
    if at is None:
        local_date = timezone.localdate()
    elif isinstance(at, datetime):
        local_date = timezone.localtime(at).date() if timezone.is_aware(at) else at.date()
    else:
        local_date = at
    return local_date.replace(day=1)


def claim_monthly_collection(
    *,
    organization_id: str,
    source: str,
    at: date | datetime | None = None,
) -> CheckoCollectionAttempt | None:
    """Atomically claim this month's only allowed Checko collection attempt."""
    attempt, created = CheckoCollectionAttempt.objects.get_or_create(
        organization_id=organization_id,
        source=source,
        period_month=collection_period_month(at),
        defaults={"status": CheckoCollectionAttempt.Status.IN_PROGRESS},
    )
    return attempt if created else None


def finish_collection(
    attempt: CheckoCollectionAttempt,
    *,
    records_count: int,
    error: Exception | str | None = None,
) -> None:
    """Persist the outcome without releasing the monthly claim."""
    attempt.records_count = max(int(records_count), 0)
    attempt.error_message = str(error or "")
    attempt.status = (
        CheckoCollectionAttempt.Status.FAILED
        if error is not None
        else CheckoCollectionAttempt.Status.SUCCESS
    )
    attempt.save(
        update_fields=[
            "records_count",
            "error_message",
            "status",
            "updated_at",
        ]
    )