forms.py

__copyright__ = "Copyright © Stichting SciPost (SciPost Foundation)"
__license__ = "AGPL v3"

import datetime

from django import forms
from django.contrib.auth import get_user_model
from django.utils.dates import MONTHS
from django.db.models import Sum
from django.utils import timezone

from ajax_select.fields import AutoCompleteSelectField
from dateutil.rrule import rrule, MONTHLY

from common.forms import MonthYearWidget
from scipost.fields import UserModelChoiceField

from .models import Subsidy, SubsidyAttachment, WorkLog


class SubsidyForm(forms.ModelForm):
    organization = AutoCompleteSelectField('organization_lookup')

    class Meta:
        model = Subsidy
        fields = ['organization', 'subsidy_type', 'description',
                  'amount', 'amount_publicly_shown', 'status',
                  'date', 'date_until']


class SubsidyAttachmentForm(forms.ModelForm):
    class Meta:
        model = SubsidyAttachment
        fields = (
            'name',
            'attachment',
            'publicly_visible',
        )

    def save(self, to_object, commit=True):
        """
        This custom save method will automatically assign the file to the object
        given when it is a valid instance type.
        """
        attachment = super().save(commit=False)

        # Formset might save an empty Instance
        if not attachment.name or not attachment.attachment:
            return None

        if isinstance(to_object, Subsidy):
            attachment.subsidy = to_object
        else:
            raise forms.ValidationError('You cannot save Attachments to this type of object.')
        if commit:
            attachment.save()
        return attachment


class SubsidyAttachmentFormSet(forms.BaseModelFormSet):
    def save(self, to_object, commit=True):
        """
        This custom save method will automatically assign the file to the object
        given when it is a valid instance type.
        """
        returns = []
        for form in self.forms:
            returns.append(form.save(to_object))
        return returns



class WorkLogForm(forms.ModelForm):
    def __init__(self, *args, **kwargs):
        self.types = kwargs.pop('log_types', False)
        super().__init__(*args, **kwargs)
        if self.types:
            self.fields['log_type'] = forms.ChoiceField(choices=self.types)

    class Meta:
        model = WorkLog
        fields = (
            'comments',
            'log_type',
            'duration',
        )
        widgets = {
            'comments': forms.Textarea(attrs={'rows': 4}),
            'duration': forms.TextInput(attrs={'placeholder': 'HH:MM:SS'})
        }


class LogsFilter(forms.Form):
    """
    Filter work logs given the requested date range and users.
    """

    employee = UserModelChoiceField(
        queryset=get_user_model().objects.filter(work_logs__isnull=False).distinct(),
        required=False, empty_label='Show all')
    start = forms.DateField(widget=MonthYearWidget(required=True), required=True)  # Month
    end = forms.DateField(widget=MonthYearWidget(required=True, end=True), required=True)  # Month

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        today = timezone.now().date()
        self.initial['start'] = today.today()
        self.initial['end'] = today.today()

    def clean(self):
        if self.is_valid():
            self.cleaned_data['months'] = [
                dt for dt in rrule(
                    MONTHLY,
                    dtstart=self.cleaned_data['start'],
                    until=self.cleaned_data['end'])]
        return self.cleaned_data

    def get_months(self):
        if self.is_valid():
            return self.cleaned_data.get('months', [])
        return []

    def filter(self):
        """Filter work logs and return in user-grouped format."""
        output = []
        if self.is_valid():
            if self.cleaned_data['employee']:
                user_qs = get_user_model().objects.filter(id=self.cleaned_data['employee'].id)
            else:
                user_qs = get_user_model().objects.filter(work_logs__isnull=False)
            user_qs = user_qs.filter(
                work_logs__work_date__gte=self.cleaned_data['start'],
                work_logs__work_date__lte=self.cleaned_data['end']).distinct()

            output = []
            for user in user_qs:
                logs = user.work_logs.filter(
                    work_date__gte=self.cleaned_data['start'],
                    work_date__lte=self.cleaned_data['end']).distinct()

                output.append({
                    'logs': logs,
                    'duration': logs.aggregate(total=Sum('duration')),
                    'user': user,
                })
        return output

    def filter_per_month(self):
        """Filter work logs and return in per-month format."""
        output = []
        if self.is_valid():
            if self.cleaned_data['employee']:
                user_qs = get_user_model().objects.filter(id=self.cleaned_data['employee'].id)
            else:
                user_qs = get_user_model().objects.filter(work_logs__isnull=False)
            user_qs = user_qs.filter(
                work_logs__work_date__gte=self.cleaned_data['start'],
                work_logs__work_date__lte=self.cleaned_data['end']).distinct()

            output = []
            for user in user_qs:
                # If logs exists for given filters
                output.append({
                    'logs': [],
                    'user': user,
                })
                for dt in self.get_months():
                    output[-1]['logs'].append(
                        user.work_logs.filter(
                            work_date__year=dt.year, work_date__month=dt.month).aggregate(total=Sum('duration'))['total'])
        return output