from typing import Iterable, Optional

from django.contrib.auth import get_user_model

from lms.courses.models import Course, CourseVisibility

from ..ammo.enrollment_detail import EnrollmentDetailAmmo
from .base_ammo import UsersWithUidBaseAmmoFactory

User = get_user_model()


class EnrollmentDetailAmmoFactory(UsersWithUidBaseAmmoFactory):
    ammo_class = EnrollmentDetailAmmo

    def generate_payload(self, user: Optional[User] = None) -> Iterable[dict]:
        for course in (
            Course.objects
            .active(user)
            .exclude(id__in=CourseVisibility.objects.unavailable_for(user))
            .prefetch_related('enrollments')
        ):
            for enrollment in course.enrollments.all():
                yield {'enrollment_id': enrollment.id}
