forms.py

import zipfile

from django import forms
from django.utils.text import slugify
from wagtail.admin.forms import WagtailAdminPageForm

from .constants import (
    AGE_30_49,
    AGE_CHOICES,
    OCCUPATION_CHOICES,
    OCCUPATION_WORKING,
    TOPIC_CHOICES,
    WEALTH_AVERAGE,
    WEALTH_CHOICES,
)
from .parser import (
    parse_program_html,
    prepare_benefit_for_all,
    prepare_benefits,
    prepare_faq,
    prepare_horizon,
    prepare_main_benefits,
    prepare_point,
)


class ProgramPointPageForm(WagtailAdminPageForm):
    import_file = forms.FileField(label="soubor s programem", required=False)

    def clean(self):
        cleaned_data = super().clean()
        cleaned_data["_import_data"] = None

        # extract data from ZIP file with HTML export
        import_file = cleaned_data["import_file"]
        if import_file:
            try:
                with zipfile.ZipFile(import_file) as archive:
                    name = archive.namelist()[0]
                    data = parse_program_html(archive.read(name))
                cleaned_data["_import_data"] = data
            except zipfile.BadZipFile:
                self.add_error("import_file", "Vadný ZIP soubor. Nelze rozbalit.")

        return cleaned_data

    def save(self, commit=True):
        page = super().save(commit=False)

        if self.cleaned_data["_import_data"]:
            point = prepare_point(self.cleaned_data["_import_data"]["sekce"])
            page.title = point["nadpis"]
            page.slug = slugify(page.title)
            page.annotation = point.get("anotace")
            page.problem = point.get("problem")
            page.context = point.get("kontext-problemu")
            page.ideal = point.get("ideal")
            page.proposal = point.get("navrhovana-opatreni")
            page.already_done = point.get("co-jsme-uz-udelali")
            page.sources = point.get("zdroje")
            text, number, unit = prepare_horizon(point.get("casovy-horizont", ""))
            page.time_horizon_text = text
            page.time_horizon_number = number
            page.time_horizon_unit = unit
            page.faq = prepare_faq(point.get("faq", ""))
            page.benefit_for_all = prepare_benefit_for_all(
                self.cleaned_data["_import_data"]["benefity"]
            )
            page.benefits_main = prepare_main_benefits(
                self.cleaned_data["_import_data"]["benefity"]
            )
            page.benefits = prepare_benefits(
                self.cleaned_data["_import_data"]["benefity"]
            )

        if commit:
            page.save()

        return page


class ProgramAppForm(forms.Form):
    age = forms.ChoiceField(choices=AGE_CHOICES, initial=AGE_30_49)
    kids = forms.BooleanField(required=False, initial=True)
    wealth = forms.ChoiceField(choices=WEALTH_CHOICES, initial=WEALTH_AVERAGE)
    occupation = forms.ChoiceField(
        choices=OCCUPATION_CHOICES, initial=OCCUPATION_WORKING
    )
    topics = forms.MultipleChoiceField(required=False, choices=TOPIC_CHOICES)