from marshmallow import Schema, fields
from typing import Type

from server import config
from server.publishers.clients.sns.utils import format_attributes


class DumpMessageMixin:
    _message_schema: Type[Schema]

    Message = fields.Method("get_message")

    def get_message(self, obj, **kwargs):
        return self._message_schema(context=self.context).dumps(obj)


class BaseMessageBodySchema(Schema):
    """Basic schema for sns-message 'Message' attribute."""

    id = fields.Int()
    app_slug = fields.Str(data_key="app")
    ttl = fields.Int()
    created_at = fields.DateTime()
    meta = fields.Dict(default=None)
    publisher = fields.Constant(
        {"id": f"{config.APP}_{config.SERVICE_ENV}", "owner": "apollo", "env": config.SERVICE_ENV}
    )


def sns_common_messages_schema(message_body_schema: Schema = BaseMessageBodySchema):
    class MessageBodyWithCodeSchema(message_body_schema):
        code = fields.Function(lambda item, context: context["code"])

    class SNSCommonMessageSchema(DumpMessageMixin, Schema):
        """Basic schema for sns-message."""

        _message_schema = MessageBodyWithCodeSchema

        Id = fields.Function(lambda obj: str(obj.id))
        MessageAttributes = fields.Function(
            lambda item, context: format_attributes(
                code=context["code"], **({"subject": item.meta["subject"]} if item.meta.get("subject") else {})
            )
        )

    return SNSCommonMessageSchema