from django.db.models import Q

from apps.core.models import AiDataset, AiDatasetSample


def get_dataset_sample_counters(dataset_id: str) -> tuple[int, int]:
    queryset = AiDatasetSample.objects.filter(dataset_id=dataset_id)
    total_count = queryset.count()
    annotated_count = queryset.filter(
        (Q(rectangle__isnull=False) & ~Q(rectangle=""))
        | (Q(polygon__isnull=False) & ~Q(polygon=""))
    ).count()
    return total_count, annotated_count


def refresh_dataset_sample_counters(dataset_id: str) -> tuple[int, int]:
    if not dataset_id:
        return 0, 0
    total_count, annotated_count = get_dataset_sample_counters(dataset_id)
    AiDataset.objects.filter(id=dataset_id).update(
        dataset_count=total_count,
        annotated_count=annotated_count,
    )
    return total_count, annotated_count