"""
Events of the batprotocol.

The following batprotocol v1.0.0 events have an experimental implementation.
They may change behavior in a breaking manner upon minor version update.

- JobsKilledEvent
- SimulationBeginsEvent


The following batprotocol v1.0.0 events are unsupported.
They may be implemented in a subsequent minor version.

- CreateProbeEvent
- StopProbeEvent
- TriggerProbeEvent
- ResetProbeEvent
- ProbeDataEmittedEvent
"""

from __future__ import annotations

from abc import abstractmethod
from enum import Enum
from typing import Any, Final, Self, final, override

from procset import ProcSet

from .batprotocol import (
    AllStaticExternalEventsHaveBeenInjectedEventPDict,
    AllStaticJobsHaveBeenSubmittedEventPDict,
    CallId,
    CallMeLaterEventPDict,
    ChangeHostsPStateEventPDict,
    EDCHelloEventPDict,
    EventPDict,
    ExecuteJobEventPDict,
    ExternalEventOccurredEventPDict,
    FinishRegistrationEventPDict,
    ForceSimulationStopEventPDict,
    HostsPStateChangedEventPDict,
    HostsTurnedOnOffEventPDict,
    JobCompletedEventPDict,
    JobId,
    JobsKilledEventPDict,
    JobSubmittedEventPDict,
    KillJobsEventPDict,
    RegisterJobEventPDict,
    RegisterProfileEventPDict,
    RejectJobEventPDict,
    RequestedCallEventPDict,
    SimulationBeginsEventPDict,
    SimulationEndsEventPDict,
    StopCallMeLaterEventPDict,
    Timestamp,
    TurnOnOffHostsEventPDict,
    _SimulationBeginsEventPayload,
)
from .core import SimulationMetadata
from .job import FinalState, Job, PlacementPolicy, Profile


def __experimental(*, reason: str | None = None):  # type: ignore[no-untyped-def]
    """Decorator to identify experimental features."""

    def decorator(obj):  # type: ignore[no-untyped-def]
        # Craft experimental note.
        entity_type = 'class' if isinstance(obj, type) else 'function'
        note = f'This {entity_type} is experimental and is subject to backward-incompatible changes upon minor releases.'  # noqa: E501 (reason: grepability)
        if reason:
            note += f'\nReason: {reason}'

        # Update docstring.
        if obj.__doc__:
            obj.__doc__ += f'\n\n{note}'
        else:
            obj.__doc__ = note

        return obj

    return decorator


# abstract base classes  -------------------------------------------------------


class Event:
    timestamp: Timestamp

    def __new__(cls, *_args: Any, **_kwargs: Any) -> Self:
        # Ensure Event, RxEvent or TxEvent are not instantiated, see:
        # - https://stackoverflow.com/a/7990308
        # - https://docs.python.org/3/reference/datamodel.html#object.__new__
        if cls in (Event, RxEvent, TxEvent):
            err_msg = f'Cannot instantiate abstract class {cls.__name__}'
            raise TypeError(err_msg)

        # We pass only cls to object.__new__() as we also define __init__().
        # However, we need to keep args and kwargs in the signature as they are
        # forwarded to __init__().
        return object.__new__(cls)

    def __init__(self, timestamp: Timestamp):
        self.timestamp = timestamp


class RxEvent(Event):
    _batprotocol_rx_events: Final[dict[str, type[RxEvent]]] = {}

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        cls._batprotocol_rx_events[cls.__name__] = cls

    # Ideally, we would like to only have a from_protocol_dict() method, and no
    # _deserialize_protocol_dict() method.
    # This is not achievable via overriding in the concrete subclasses as this
    # would violate the Liskov substitution principle: method arguments are
    # contravariant, and we want to further specify the type in our case.
    # _deserialize_protocol_dict() does not uphold the Liskov substitution
    # principle either but this is not exposed to user-facing API.

    @final
    @classmethod
    def from_protocol_dict(cls, payload: EventPDict) -> RxEvent:
        concrete_cls = cls._batprotocol_rx_events[payload['event_type']]
        return concrete_cls._deserialize_protocol_dict(payload)  # noqa: SLF001 (reason: private in module scope)

    @classmethod
    @abstractmethod
    def _deserialize_protocol_dict(cls, payload) -> Self: ...  # type: ignore[no-untyped-def]


class TxEvent(Event):
    @abstractmethod
    def to_protocol_dict(self) -> EventPDict: ...


# concrete classes  ------------------------------------------------------------


# TODO(rb):
#   dispatch elements of _payload to attributes with semantic types:
#     - SimulationMetadata
#     - (to create/define) Simulation(Args|Params|Context)
@__experimental(reason='_payload is yet to be mapped to semantic types.')
class SimulationBeginsEvent(RxEvent):
    _payload: _SimulationBeginsEventPayload  # XXX: replace with semantic attributes

    @override
    def __init__(self, timestamp: Timestamp, payload: _SimulationBeginsEventPayload):
        super().__init__(timestamp)
        self._payload = payload

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: SimulationBeginsEventPDict,
    ) -> SimulationBeginsEvent:
        return cls(
            timestamp=payload['timestamp'],
            payload=payload['event'],
        )

    @property
    def computation_host_number(self) -> int:
        return self._payload['computation_host_number']


class SimulationEndsEvent(RxEvent):
    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: SimulationEndsEventPDict,
    ) -> SimulationEndsEvent:
        return cls(
            timestamp=payload['timestamp'],
        )


class JobSubmittedEvent(RxEvent):
    job: Job
    fwd_profile: Profile | None

    @override
    def __init__(
        self,
        timestamp: Timestamp,
        job: Job,
        fwd_profile: Profile | None = None,
    ):
        super().__init__(timestamp)
        self.job = job
        self.fwd_profile = fwd_profile

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: JobSubmittedEventPDict,
    ) -> JobSubmittedEvent:
        fwd_profile_dict = payload['event'].get('profile')
        fwd_profile = (
            Profile.from_protocol_dict(fwd_profile_dict)
            if fwd_profile_dict is not None
            else None
        )

        job = Job(
            id_=payload['event']['job_id'],
            resource_request=payload['event']['job']['resource_request'],
            walltime=payload['event']['job'].get('walltime'),
            profile_id=payload['event']['job']['profile_id'],
            extra_data=payload['event']['job'].get('extra_data'),
        )
        job.submission_time = payload['event']['submission_time']

        assert fwd_profile is None or fwd_profile.id_ == job.profile_id

        return cls(
            timestamp=payload['timestamp'],
            job=job,
            fwd_profile=fwd_profile,
        )


class JobCompletedEvent(RxEvent):
    job: Job
    state: FinalState
    return_code: int

    @override
    def __init__(
        self,
        timestamp: Timestamp,
        job: Job,
        state: FinalState,
        return_code: int,
    ):
        super().__init__(timestamp)
        self.job = job
        self.state = state
        self.return_code = return_code

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: JobCompletedEventPDict,
    ) -> JobCompletedEvent:
        assert payload['__pybatsim_job'].id_ == payload['event']['job_id']
        return cls(
            timestamp=payload['timestamp'],
            job=payload['__pybatsim_job'],  # injected by deserialisation
            state=FinalState(payload['event']['state']),
            return_code=payload['event']['return_code'],
        )


@__experimental(reason='progress is yet to be mapped to a semantic type.')
class JobsKilledEvent(RxEvent):
    # TODO(rb): consider using a list of tuples
    jobs: list[Job]
    progresses: dict[JobId, dict[Any, Any]]

    @override
    def __init__(
        self,
        timestamp: Timestamp,
        jobs: list[Job],
        progresses: dict[JobId, dict[Any, Any]],
    ):
        super().__init__(timestamp)
        self.jobs = jobs
        self.progresses = progresses

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: JobsKilledEventPDict,
    ) -> JobsKilledEvent:
        progresses: dict[JobId, Any] = {}  # XXX: replace with semantic type
        for progress_dict in payload['event']['progresses']:
            job_id = progress_dict['job_id']
            # TODO: correctly deserialise the wrapped progress in a KillProgress object
            progresses[job_id] = progress_dict['wrapper']

        return cls(
            timestamp=payload['timestamp'],
            jobs=payload['__pybatsim_dead_jobs'],  # injected by deserialisation
            progresses=progresses,
        )


class RequestedCallEvent(RxEvent):
    call_id: CallId
    is_last_call: bool

    @override
    def __init__(self, timestamp: Timestamp, call_id: CallId, is_last_call: bool):
        super().__init__(timestamp)
        self.call_id = call_id
        self.is_last_call = is_last_call

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: RequestedCallEventPDict,
    ) -> RequestedCallEvent:
        return cls(
            timestamp=payload['timestamp'],
            call_id=payload['event']['call_me_later_id'],
            is_last_call=payload['event']['last_periodic_call'],
        )


class ExternalEventOccurredEvent(RxEvent):
    external_event_id: str
    # TODO: consider using an enum, might be overkill as there exists a single
    # type as of now (GenericExternalEvent)
    external_event_type: str
    external_event: dict[Any, Any]  # # XXX: replace with semantic type

    @override
    def __init__(
        self,
        timestamp: Timestamp,
        external_event_id: str,
        external_event_type: str,
        external_event: dict[Any, Any],
    ):
        super().__init__(timestamp)
        self.external_event_id = external_event_id
        self.external_event_type = external_event_type
        self.external_event = external_event

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: ExternalEventOccurredEventPDict,
    ) -> ExternalEventOccurredEvent:
        return cls(
            timestamp=payload['timestamp'],
            external_event_id=payload['event']['id'],
            external_event_type=payload['event']['external_event_type'],
            external_event=payload['event']['external_event'],
        )


class AllStaticJobsHaveBeenSubmittedEvent(RxEvent):
    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: AllStaticJobsHaveBeenSubmittedEventPDict,
    ) -> AllStaticJobsHaveBeenSubmittedEvent:
        return cls(
            timestamp=payload['timestamp'],
        )


class AllStaticExternalEventsHaveBeenInjectedEvent(RxEvent):
    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: AllStaticExternalEventsHaveBeenInjectedEventPDict,
    ) -> AllStaticExternalEventsHaveBeenInjectedEvent:
        return cls(
            timestamp=payload['timestamp'],
        )


class HostsPStateChangedEvent(RxEvent):
    host_ids: ProcSet
    pstate: int

    @override
    def __init__(self, timestamp: Timestamp, host_ids: ProcSet, pstate: int):
        super().__init__(timestamp)
        self.host_ids = host_ids
        self.pstate = pstate

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: HostsPStateChangedEventPDict,
    ) -> HostsPStateChangedEvent:
        return cls(
            timestamp=payload['timestamp'],
            host_ids=ProcSet.from_str(payload['event']['host_ids']),
            pstate=payload['event']['pstate'],
        )


class HostsTurnedOnOffEvent(RxEvent):
    host_ids: ProcSet
    state: int

    @override
    def __init__(self, timestamp: Timestamp, host_ids: ProcSet, state: int):
        super().__init__(timestamp)
        self.host_ids = host_ids
        self.state = state

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: HostsTurnedOnOffEventPDict,
    ) -> HostsTurnedOnOffEvent:
        return cls(
            timestamp=payload['timestamp'],
            host_ids=ProcSet.from_str(payload['event']['host_ids']),
            state=payload['event']['state'],
        )


class EDCHelloEvent(TxEvent):
    simulation_metadata: SimulationMetadata
    # TODO: consider integrating edc_* in simulation_metadata
    edc_name: str
    edc_version: str
    edc_commit: str

    @override
    def __init__(
        self,
        timestamp: Timestamp,
        simulation_metadata: SimulationMetadata,
        edc_name: str,
        edc_version: str,
        edc_commit: str,
    ):
        super().__init__(timestamp)
        self.simulation_metadata = simulation_metadata
        self.edc_name = edc_name
        self.edc_version = edc_version
        self.edc_commit = edc_commit

    @override
    def to_protocol_dict(self) -> EDCHelloEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'EDCHelloEvent',
            'event': {
                **self.simulation_metadata.to_protocol_dict(),
                'decision_component_name': self.edc_name,
                'decision_component_version': self.edc_version,
                'decision_component_commit': self.edc_commit,
            },
        }


class RejectJobEvent(TxEvent):
    job: Job

    @override
    def __init__(self, timestamp: Timestamp, job: Job):
        super().__init__(timestamp)
        self.job = job

    @override
    def to_protocol_dict(self) -> RejectJobEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'RejectJobEvent',
            'event': {
                'job_id': self.job.id_,
            },
        }


class ExecuteJobEvent(TxEvent):
    job: Job
    placement_policy: PlacementPolicy

    @override
    def __init__(
        self,
        timestamp: Timestamp,
        job: Job,
        placement_policy: PlacementPolicy | None = None,
    ):
        super().__init__(timestamp)
        self.job = job
        self.placement_policy = (
            PlacementPolicy.default() if placement_policy is None else placement_policy
        )

    @override
    def to_protocol_dict(self) -> ExecuteJobEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'ExecuteJobEvent',
            'event': {
                'job_id': self.job.id_,
                'allocation': {
                    **self.placement_policy.mapping_strategy.to_protocol_dict(),
                    'host_allocation': str(self.job.allocation),
                },
            },
        }


class KillJobsEvent(TxEvent):
    jobs: list[Job]

    @override
    def __init__(self, timestamp: Timestamp, jobs: list[Job]):
        super().__init__(timestamp)
        self.jobs = jobs

    @override
    def to_protocol_dict(self) -> KillJobsEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'KillJobsEvent',
            'event': {
                'job_ids': [job.id_ for job in self.jobs],
            },
        }


class CallMeLaterEvent(TxEvent):
    # TODO: nesting class is not Pythonic
    class TemporalTriggerType(Enum):
        OneShot = 0
        Periodic = 1

    call_id: CallId
    # XXX: replace when_type and when by a semantic type
    when_type: TemporalTriggerType
    when: dict[Any, Any]

    @override
    def __init__(
        self,
        timestamp: Timestamp,
        call_id: CallId,
        when_type: TemporalTriggerType,
        when_dict: dict[Any, Any],
    ):
        super().__init__(timestamp)
        self.call_id = call_id
        self.when_type = when_type
        self.when = when_dict

    @override
    def to_protocol_dict(self) -> CallMeLaterEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'CallMeLaterEvent',
            'event': {
                'call_me_later_id': self.call_id,
                'when_type': self.when_type.name,
                'when': self.when,
            },
        }


class StopCallMeLaterEvent(TxEvent):
    call_id: CallId

    @override
    def __init__(self, timestamp: Timestamp, call_id: CallId):
        super().__init__(timestamp)
        self.call_id = call_id

    @override
    def to_protocol_dict(self) -> StopCallMeLaterEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'StopCallMeLaterEvent',
            'event': {
                'call_me_later_id': self.call_id,
            },
        }


class ForceSimulationStopEvent(TxEvent):
    @override
    def to_protocol_dict(self) -> ForceSimulationStopEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'ForceSimulationStopEvent',
            'event': {},
        }


class FinishRegistrationEvent(TxEvent):
    @override
    def to_protocol_dict(self) -> FinishRegistrationEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'FinishRegistrationEvent',
            'event': {},
        }


class RegisterJobEvent(TxEvent):
    job: Job

    @override
    def __init__(self, timestamp: Timestamp, job: Job):
        super().__init__(timestamp)
        self.job = job
        # XXX: job.submission_time will never be set if ack of dynamic jobs is
        # disabled: consider injecting timestamp as job.submission_time

    @override
    def to_protocol_dict(self) -> RegisterJobEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'RegisterJobEvent',
            'event': {
                'job_id': self.job.id_,
                'job': self.job.to_protocol_dict(),
            },
        }


class RegisterProfileEvent(TxEvent):
    profile: Profile

    @override
    def __init__(self, timestamp: Timestamp, profile: Profile):
        super().__init__(timestamp)
        self.profile = profile

    @override
    def to_protocol_dict(self) -> RegisterProfileEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'RegisterProfileEvent',
            'event': {
                'profile': self.profile.to_protocol_dict(),
            },
        }


class ChangeHostsPStateEvent(TxEvent):
    host_ids: ProcSet
    pstate: int

    @override
    def __init__(self, timestamp: Timestamp, host_ids: ProcSet, pstate: int):
        super().__init__(timestamp)
        self.host_ids = host_ids
        self.pstate = pstate

    @override
    def to_protocol_dict(self) -> ChangeHostsPStateEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'ChangeHostsPStateEvent',
            'event': {
                'host_ids': str(self.host_ids),
                'pstate': self.pstate,
            },
        }


class TurnOnOffHostsEvent(TxEvent):
    host_ids: ProcSet
    state: int

    @override
    def __init__(self, timestamp: Timestamp, host_ids: ProcSet, state: int):
        super().__init__(timestamp)
        self.host_ids = host_ids
        self.state = state

    @override
    def to_protocol_dict(self) -> TurnOnOffHostsEventPDict:
        return {
            'timestamp': self.timestamp,
            'event_type': 'TurnOnOffHostsEvent',
            'event': {
                'host_ids': str(self.host_ids),
                'state': self.state,
            },
        }
