from marshmallow import Schema, fields, post_dump, post_load

from models import Poll

from models.schemas import PollOptionSchema


class PollSchema(Schema):
    id = fields.Integer(required=True)
    name = fields.Str(required=True)
    description = fields.Str(required=True)
    start_date = fields.DateTime(required=True)
    end_date = fields.DateTime(required=True)
    project_id = fields.Integer(required=True)
    options = fields.List(fields.Nested(PollOptionSchema))

    @post_load
    def make_poll(self, data):
        return Poll(**data)

    @post_dump(pass_many=True)
    def add_envelope(self, data, many):
        if many:
            return {"polls": data}

        return
