from typing import Union
from dataclasses import dataclass

from marl_factory_grid.environment.entity.entity import Entity

TYPE_VALUE  = 'value'
TYPE_REWARD = 'reward'
types = [TYPE_VALUE, TYPE_REWARD]

@dataclass
class InfoObject:
    identifier: str
    val_type: str
    value: Union[float, int]


@dataclass
class Result:
    identifier: str
    validity: bool
    reward: Union[float, None] = None
    value: Union[float, None] = None
    entity: Union[Entity, None] = None

    def get_infos(self):
        n = self.entity.name if self.entity is not None else "Global"
        return [InfoObject(identifier=f'{n}_{self.identifier}_{t}',
                           val_type=t, value=self.__getattribute__(t)) for t in types
                if self.__getattribute__(t) is not None]

    def __repr__(self):
        valid = "not " if not self.validity else ""
        return f'{self.__class__.__name__}({self.identifier.capitalize()} {valid}valid: {self.reward})'


@dataclass
class TickResult(Result):
    pass


@dataclass
class ActionResult(Result):
    pass


@dataclass
class DoneResult(Result):
    pass
