from decimal import Decimal, InvalidOperation

from django.views.generic import TemplateView

from core.models import Subject
from core.prediction import predict_all, predict_balance


class BalanceView(TemplateView):
    template_name = "core/balance.html"

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        subjects = Subject.objects.filter(user=self.request.user).order_by("name")
        amount, amount_error = Decimal(), False
        if amount_query := self.request.GET.get("amount"):
            try:
                amount = Decimal(amount_query)
            except InvalidOperation:
                amount_error = True
        context["amount"] = round(amount, 2)
        context["amount_error"] = amount_error
        context["future_transactions"] = predict_balance(subjects, amount)
        return context


class SubjectsView(TemplateView):
    template_name = "core/subjects.html"

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        subjects = Subject.objects.filter(user=self.request.user).order_by("name")
        context["prediction_list"] = predict_all(subjects)
        return context