# MODELS

class User(models.Model):
    name = models.CharField(max_length=50)
    # ...etc...

class Assignment(models.Model):
    user = models.ForeignKey(
        settings.AUTH_USER_MODEL, on_delete=models.PROTECT
    )
    job = models.ForeignKey(Job, on_delete=models.PROTECT)
    location = models.ForeignKey(Location, on_delete=models.PROTECT)

# SERIALIZER

class UserSerializer(serializers.HyperlinkedModelSerializer):
    assignment_job = serializers.SerializerMethodField()
    multiple_locations = serializers.SerializerMethodField()
    multiple_jobs = serializers.SerializerMethodField()

    class Meta:
        model = User
        fields = [
            "url",
            "id",
            "username",
            "first_name",
            "last_name",
            "full_name",
            "email",
            "is_staff",
            "assignment_job",
            "multiple_locations",
            "multiple_jobs",
        ]

    def get_assignment_job(self, obj):
        assignment = obj.assignment_set.get(primary=True)
        return assignment.job.description

    def get_multiple_locations(self, obj):
        location_count = obj.assignment_set.filter(
            end_date__isnull=True
        ).aggregate(total_locations=Count("location"))
        if location_count["total_locations"] > 1:
            return True
        return False

    def get_multiple_jobs(self, obj):
        job_count = obj.assignment_set.filter(end_date__isnull=True).aggregate(
            total_jobs=Count("job")
        )
        if job_count["total_jobs"] > 1:
            return True
        return False


# SERIALIZER VIEW

class UserViewSet(viewsets.ModelViewSet):
    serializer_class = UserSerializer

    def get_queryset(self):
        users = get_observable_users(self.request.user)
        return users


# get_observable_users Method
# this returns the correct users - I've tried adding the prefetch_related here

def get_observable_users(user):
    user_assignments = Assignment.active.filter(user=user)
    user_jobs = user_assignments.values_list("job", flat=True)
    locations = user_assignments.values_list("location", flat=True)
    administrative_job_observations = JobObservation.objects.filter(
        observer_job__in=user_jobs,
        restrict_to_location=False,
    )
    observee_jobs = JobObservation.objects.filter(
        observer_job__in=user_jobs
    ).values_list("observee_job", flat=True)
    if administrative_job_observations:
        people_assignments = Assignment.objects.filter(job__in=observee_jobs)
    else:
        people_assignments = Assignment.objects.filter(
            job__in=observee_jobs, location__in=locations
        )
    all_assignments = user_assignments | people_assignments
    all_user_ids = all_assignments.values_list("user", flat=True)
    user_list = User.objects.prefetch_related("assignment_set").filter(
        id__in=all_user_ids
    )
    return user_list