from marshmallow import fields
from marshmallow_enum import EnumField

from maps_adv.stat_controller.server.lib.domains.charger import Domain, TaskStatus

from .schemas import (
    FindNewTaskSchema,
    TaskDetailsSchema,
    UpdateTaskSchema as BaseUpdateTaskSchema,
    with_schemas,
)

__all__ = ["ApiProvider"]


class ChargerTaskDetailsSchema(TaskDetailsSchema):
    status = EnumField(TaskStatus, required=True)
    execution_state = fields.Raw(required=False)


class UpdateTaskSchema(BaseUpdateTaskSchema):
    status = EnumField(TaskStatus, required=True)
    execution_state = fields.Raw(required=True)


class ApiProvider:
    __slots__ = ("_domain",)

    _domain: Domain

    def __init__(self, domain: Domain):
        self._domain = domain

    @with_schemas(FindNewTaskSchema, ChargerTaskDetailsSchema)
    async def find_new(self, **kwargs) -> dict:
        return await self._domain.find_new(**kwargs)

    @with_schemas(UpdateTaskSchema, TaskDetailsSchema)
    async def update(self, **kwargs) -> dict:
        return await self._domain.update(**kwargs)
