import json
from functools import partial
from typing import Any
from typing import Callable
from typing import Dict
from typing import Optional
from typing import TypeVar
from typing import cast

from dramatiq.message import Message

from dramatiq.actor import Actor
from dramatiq.actor import actor
from pydantic import BaseModel
from pydantic import validate_arguments
from pydantic.decorator import V_DUPLICATE_KWARGS
from pydantic.decorator import V_POSITIONAL_ONLY_NAME

R = TypeVar('R')


class PydanticActor(Actor):
    def __init__(self, fn: Callable[..., R], **kwargs):
        super().__init__(validate_arguments(fn), **kwargs)

    def message_with_options(
        self,
        *,
        args: tuple = (),
        kwargs: Optional[Dict[str, Any]] = None,
        **options,
    ) -> Message[R]:
        assert not args, 'positional arguments are not allowed in pydantic model'
        instance = cast(BaseModel, self.fn.model(*args, **kwargs))
        # exclude extra fields created by `validate_arguments`
        exclude = {
            V_POSITIONAL_ONLY_NAME,
            V_DUPLICATE_KWARGS,
            self.fn.vd.v_args_name,
            self.fn.vd.v_kwargs_name,
        }
        json_kwargs = json.loads(instance.json(exclude=exclude))

        return super().message_with_options(args=args, kwargs=json_kwargs, **options)

    def send(self, /, **kwargs) -> Message[R]:
        # positional arguments are not allowed in pydantic model
        return super().send(**kwargs)


pydantic_actor = partial(actor, actor_class=PydanticActor)