import graphene
import json
from graphene import relay

from api.utils import get_object
from surveys.inputs import UpdateAnswerArgument
from surveys.models import Survey, Answer
from surveys.schema import AnswerNode


class UpdateAnswer(relay.ClientIDMutation):
    class Input:
        answer = graphene.Argument(UpdateAnswerArgument)

    answer = graphene.Field(AnswerNode)

    @classmethod
    def mutate_and_get_payload(cls, root, info, **kwargs):
        user = info.context.user

        answer = kwargs.get('answer')
        survey_id = answer.get('survey_id')
        data = json.loads(answer.get('data'))
        survey = get_object(Survey, survey_id)

        try:
            answer = survey.answers.get(owner=user)
            answer.data = data
            answer.save()
        except Answer.DoesNotExist:
            answer = Answer.objects.create(owner=user, survey=survey, data=data)

        return cls(answer=answer)


class SurveyMutations:
    update_answer = UpdateAnswer.Field()