"""
Definitions of types describing the structure of the messages exchanged with
batsim.
These descriptions define the pybatsim interpretation of the batprocotol.
"""

from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Final,
    Literal,
    LiteralString,
    NotRequired,
    TypedDict,
)

if TYPE_CHECKING:
    from .job import Job

__batprotocol_version__: Final[str] = '1.0.0'


type JobId = str
type ProfileId = str

type HostId = int
type StorageId = str

type CallId = str

type Timestamp = float
type Duration = float


# The wire format of the batprotocol is defined via TypedDict types.
# Keep in mind "a TypedDict is a structural type" and that "TypedDict types […]
# inherit from other TypedDict types to share common items".


# ----- simulation metadata wire format ----


class SimulationFeaturesPDict(TypedDict):
    dynamic_registration: bool
    profile_reuse: bool
    acknowledge_dynamic_jobs: bool
    forward_profiles_on_job_submission: bool
    forward_profiles_on_jobs_killed: bool
    forward_profiles_on_simulation_begins: bool


# XXX(rb): consider merging with related classes for EDCHelloEvent as this is a
# strict subset of the content of the EDCHelloEvent protocol table.
class SimulationMetadataPDict(TypedDict):
    batprotocol_version: str
    batprotocol_commit: NotRequired[str]
    requested_simulation_features: SimulationFeaturesPDict


# ----- events wire format -----


# The Event union of the batprotocol is never directly used: it is always
# wrapped in the EventAndTimestamp table.
# Hence, the _BaseEventPDict describes the structure of the EventAndTimestamp
# table.
class _BaseEventPDict[T: LiteralString, P](TypedDict):
    timestamp: Timestamp
    event_type: T
    event: P


class _EmptyEventPayload(TypedDict):
    pass


class _JobSubmittedEventPayload(TypedDict):
    job_id: JobId
    job: JobPDict
    submission_time: Timestamp
    profile: NotRequired[ProfilePDict]


class _JobCompletedEventPayload(TypedDict):
    job_id: JobId
    state: Literal[
        'COMPLETED_SUCCESSFULLY',
        'COMPLETED_FAILED',
        'COMPLETED_WALLTIME_REACHED',
        'COMPLETED_KILLED',
        'REJECTED',
    ]
    return_code: int


class _RejectJobEventPayload(TypedDict):
    job_id: JobId


class _ExecuteJobEventPayload(TypedDict):
    job_id: JobId
    allocation: JobAllocationPDict
    # profile_allocation_override and storage_placement are not supported yet
    profile_allocation_override: NotRequired[list[ProfileAllocationPDict]]
    storage_placement: NotRequired[list[StorageHostMappingPDict]]


class _KillJobsEventPayload(TypedDict):
    job_ids: list[JobId]


class _JobsKilledEventPayload(TypedDict):
    job_ids: list[JobId]
    progresses: list[Any]  # XXX: replace with semantic type
    profiles: NotRequired[list[ProfilePDict]]


class _RegisterProfileEventPayload(TypedDict):
    profile: ProfilePDict


class _RegisterJobEventPayload(TypedDict):
    job_id: JobId
    job: JobPDict


# CreateProbeEvent is not supported.

# StopProbeEvent is not supported.

# TriggerProbeEvent is not supported.

# ResetProbeEvent is not supported.

# ProbeDataEmittedEvent is not supported.


class _CallMeLaterEventPayload(TypedDict):
    call_me_later_id: CallId
    when_type: Any  # XXX: replace with semantic type
    when: Any  # XXX: replace with semantic type


class _RequestedCallEventPayload(TypedDict):
    call_me_later_id: CallId
    last_periodic_call: bool


class _StopCallMeLaterEventPayload(TypedDict):
    call_me_later_id: CallId


class _EDCHelloEventPayload(TypedDict):
    batprotocol_version: str
    batprotocol_commit: NotRequired[str]
    decision_component_name: str
    decision_component_version: str
    decision_component_commit: NotRequired[str]
    requested_simulation_features: SimulationFeaturesPDict


class _SimulationBeginsEventPayload(TypedDict):
    batprotocol_version: str
    batsim_version: str
    batsim_commit: NotRequired[str]

    host_number: int
    computation_host_number: int
    computation_hosts: list[Any]  # XXX: replace with semantic type
    storage_host_number: int
    storage_hosts: list[Any]  # XXX: replace with semantic type

    batsim_arguments: list[str]
    batsim_execution_context_json: str

    workloads: list[Any]  # XXX: replace with semantic type
    external_event_lists: list[Any]  # XXX: replace with semantic type
    profiles: NotRequired[list[ProfilePDict]]


# SimulationEndsEvent has an empty payload.

# SimulationErrorEvent is not supported.

# AllStaticJobsHaveBeenSubmittedEvent has an empty payload.

# AllStaticExternalEventsHaveBeenInjectedEvent has an empty payload.

# FinishRegistrationEvent has an empty payload.

# ForceSimulationStopEvent has an empty payload.


class _ChangeHostsPStateEventPayload(TypedDict):
    host_ids: str
    pstate: int


class _HostsPStateChangedEventPayload(TypedDict):
    host_ids: str
    pstate: int


class _TurnOnOffHostsEventPayload(TypedDict):
    host_ids: str
    state: int


class _HostsTurnedOnOffEventPayload(TypedDict):
    host_ids: str
    state: int


class _ExternalEventOccurredEventPayload(TypedDict):
    id: str
    external_event_type: str  # XXX: replace with semantic type
    external_event: dict[Any, Any]  # XXX: replace with semantic type


type JobSubmittedEventPDict = _BaseEventPDict[
    Literal['JobSubmittedEvent'],
    _JobSubmittedEventPayload,
]


class JobCompletedEventPDict(
    _BaseEventPDict[
        Literal['JobCompletedEvent'],
        _JobCompletedEventPayload,
    ]
):
    __pybatsim_job: Job


type RejectJobEventPDict = _BaseEventPDict[
    Literal['RejectJobEvent'],
    _RejectJobEventPayload,
]
type ExecuteJobEventPDict = _BaseEventPDict[
    Literal['ExecuteJobEvent'],
    _ExecuteJobEventPayload,
]
type KillJobsEventPDict = _BaseEventPDict[
    Literal['KillJobsEvent'],
    _KillJobsEventPayload,
]


class JobsKilledEventPDict(
    _BaseEventPDict[
        Literal['JobsKilledEvent'],
        _JobsKilledEventPayload,
    ]
):
    __pybatsim_dead_jobs: list[Job]


type RegisterProfileEventPDict = _BaseEventPDict[
    Literal['RegisterProfileEvent'],
    _RegisterProfileEventPayload,
]
type RegisterJobEventPDict = _BaseEventPDict[
    Literal['RegisterJobEvent'],
    _RegisterJobEventPayload,
]
# CreateProbeEventPDict is not supported.
# StopProbeEventPDict is not supported.
# TriggerProbeEventPDict is not supported.
# ResetProbeEventPDict is not supported.
# ProbeDataEmittedEventPDict is not supported.
type CallMeLaterEventPDict = _BaseEventPDict[
    Literal['CallMeLaterEvent'],
    _CallMeLaterEventPayload,
]
type RequestedCallEventPDict = _BaseEventPDict[
    Literal['RequestedCallEvent'],
    _RequestedCallEventPayload,
]
type StopCallMeLaterEventPDict = _BaseEventPDict[
    Literal['StopCallMeLaterEvent'],
    _StopCallMeLaterEventPayload,
]
type EDCHelloEventPDict = _BaseEventPDict[
    Literal['EDCHelloEvent'],
    _EDCHelloEventPayload,
]
type SimulationBeginsEventPDict = _BaseEventPDict[
    Literal['SimulationBeginsEvent'],
    _SimulationBeginsEventPayload,
]
type SimulationEndsEventPDict = _BaseEventPDict[
    Literal['SimulationEndsEvent'],
    _EmptyEventPayload,
]
# SimulationErrorEvent is not supported.
type AllStaticJobsHaveBeenSubmittedEventPDict = _BaseEventPDict[
    Literal['AllStaticJobsHaveBeenSubmittedEvent'],
    _EmptyEventPayload,
]
type AllStaticExternalEventsHaveBeenInjectedEventPDict = _BaseEventPDict[
    Literal['AllStaticExternalEventsHaveBeenInjectedEvent'],
    _EmptyEventPayload,
]
type FinishRegistrationEventPDict = _BaseEventPDict[
    Literal['FinishRegistrationEvent'],
    _EmptyEventPayload,
]
type ForceSimulationStopEventPDict = _BaseEventPDict[
    Literal['ForceSimulationStopEvent'],
    _EmptyEventPayload,
]
type ChangeHostsPStateEventPDict = _BaseEventPDict[
    Literal['ChangeHostsPStateEvent'],
    _ChangeHostsPStateEventPayload,
]
type HostsPStateChangedEventPDict = _BaseEventPDict[
    Literal['HostsPStateChangedEvent'],
    _HostsPStateChangedEventPayload,
]
type TurnOnOffHostsEventPDict = _BaseEventPDict[
    Literal['TurnOnOffHostsEvent'],
    _TurnOnOffHostsEventPayload,
]
type HostsTurnedOnOffEventPDict = _BaseEventPDict[
    Literal['HostsTurnedOnOffEvent'],
    _HostsTurnedOnOffEventPayload,
]
type ExternalEventOccurredEventPDict = _BaseEventPDict[
    Literal['ExternalEventOccurredEvent'],
    _ExternalEventOccurredEventPayload,
]


type EventPDict = (
    JobSubmittedEventPDict
    | JobCompletedEventPDict
    | RejectJobEventPDict
    | ExecuteJobEventPDict
    | KillJobsEventPDict
    | JobsKilledEventPDict
    | RegisterProfileEventPDict
    | RegisterJobEventPDict
    | CallMeLaterEventPDict
    | RequestedCallEventPDict
    | StopCallMeLaterEventPDict
    | EDCHelloEventPDict
    | SimulationBeginsEventPDict
    | SimulationEndsEventPDict
    | AllStaticJobsHaveBeenSubmittedEventPDict
    | AllStaticExternalEventsHaveBeenInjectedEventPDict
    | FinishRegistrationEventPDict
    | ForceSimulationStopEventPDict
    | ChangeHostsPStateEventPDict
    | HostsPStateChangedEventPDict
    | TurnOnOffHostsEventPDict
    | HostsTurnedOnOffEventPDict
    | ExternalEventOccurredEventPDict
)


# ----- message wire format -----


class MessagePDict(TypedDict):
    now: Timestamp
    events: list[EventPDict]


# ----- host mappings wire format -----


class _BaseMappingStrategyPDict[T: LiteralString, P](TypedDict):
    executor_placement_type: T
    executor_placement: P


class _PredefinedMappingPayload(TypedDict):
    strategy: Literal[
        'SpreadOverHostsFirst',
        'FillOneHostCoresFirst',
    ]


class _CustomMappingPayload(TypedDict):
    mapping: list[int]


type PredefinedMappingPDict = _BaseMappingStrategyPDict[
    Literal['PredefinedExecutorPlacementStrategyWrapper'],
    _PredefinedMappingPayload,
]
type CustomMappingPDict = _BaseMappingStrategyPDict[
    Literal['CustomExecutorToHostMapping'],
    _CustomMappingPayload,
]
type MappingStrategyPDict = PredefinedMappingPDict | CustomMappingPDict


# It would be nice if we could automatically define JobAllocationPDict and
# ProfileAllocationPDict variant compatible with each corresponding
# MappingStrategyPDict variant.
# Here, we only check the literal used: this is not ideal, but this should let
# the type checker catch most incoherences.


class _BaseJobAllocationPDict[T: LiteralString, P](_BaseMappingStrategyPDict[T, P]):
    host_allocation: str


class _BaseProfileAllocationPDict[T: LiteralString, P](_BaseJobAllocationPDict[T, P]):
    profile_id: ProfileId


type JobAllocationPDict = _BaseJobAllocationPDict[
    Literal[
        'PredefinedExecutorPlacementStrategyWrapper',
        'CustomExecutorToHostMapping',
    ],
    Any,
]

type ProfileAllocationPDict = _BaseJobAllocationPDict[
    Literal[
        'PredefinedExecutorPlacementStrategyWrapper',
        'CustomExecutorToHostMapping',
    ],
    Any,
]


class StorageHostMappingPDict(TypedDict):
    storage_name: StorageId
    host_id: HostId


# ----- jobs and profiles wire format -----


class JobPDict(TypedDict):
    profile_id: ProfileId
    resource_request: int
    walltime: NotRequired[Duration]
    extra_data: NotRequired[str]


# The Profile union of the batprotocol is never directly used: it is always
# wrapped in the ProfileAndId table.
# Hence, the _BaseProfilePDict describes the structure of the ProfileAndId
# table.
class _BaseProfilePDict[T: LiteralString, P](TypedDict):
    id: ProfileId
    profile_type: T
    profile: P
    extra_data: NotRequired[str]


type _HomogeneousProfileGenerationStrategyLiteral = Literal[
    'DefinedAmountsUsedForEachValue',
    'DefinedAmountsSpreadUniformly',
]


class _DelayProfilePayload(TypedDict):
    delay: Duration


class _ParallelTaskProfilePayload(TypedDict):
    computation_vector: NotRequired[list[float]]
    communication_matrix: NotRequired[list[float]]


class _SequentialCompositionProfilePayload(TypedDict):
    repetition_count: int
    profile_ids: list[ProfileId]


class _ForkJoinCompositionProfilePayload(TypedDict):
    profile_ids: list[ProfileId]


class _ParallelTaskMergeCompositionProfilePayload(TypedDict):
    profile_ids: list[ProfileId]


class _ParallelTaskHomogeneousProfilePayload(TypedDict):
    computation_amount: float
    communication_amount: float
    generation_strategy: _HomogeneousProfileGenerationStrategyLiteral


class _ParallelTaskOnStorageHomogeneousProfilePayload(TypedDict):
    storage_name: StorageId
    bytes_to_read: float
    bytes_to_write: float
    generation_strategy: _HomogeneousProfileGenerationStrategyLiteral


class _ParallelTaskDataStagingBetweenStoragesProfilePayload(TypedDict):
    bytes_to_transfer: float
    emitter_storage_name: StorageId
    receiver_storage_name: StorageId


class _TraceReplayProfilePayload(TypedDict):
    trace_type: Literal[
        'SMPI',
        'FractionalComputation',
    ]
    filename: str


type DelayProfilePDict = _BaseProfilePDict[
    Literal['DelayProfile'],
    _DelayProfilePayload,
]
type ParallelTaskProfilePDict = _BaseProfilePDict[
    Literal['ParallelTaskProfile'],
    _ParallelTaskProfilePayload,
]
type SequentialCompositionProfilePDict = _BaseProfilePDict[
    Literal['SequentialCompositionProfile'],
    _SequentialCompositionProfilePayload,
]
type ForkJoinCompositionProfilePDict = _BaseProfilePDict[
    Literal['ForkJoinCompositionProfile'],
    _ForkJoinCompositionProfilePayload,
]
type ParallelTaskMergeCompositionProfilePDict = _BaseProfilePDict[
    Literal['ParallelTaskMergeCompositionProfile'],
    _ParallelTaskMergeCompositionProfilePayload,
]
type ParallelTaskHomogeneousProfilePDict = _BaseProfilePDict[
    Literal['ParallelTaskHomogeneousProfile'],
    _ParallelTaskHomogeneousProfilePayload,
]
type ParallelTaskOnStorageHomogeneousProfilePDict = _BaseProfilePDict[
    Literal['ParallelTaskOnStorageHomogeneousProfile'],
    _ParallelTaskOnStorageHomogeneousProfilePayload,
]
type ParallelTaskDataStagingBetweenStoragesProfilePDict = _BaseProfilePDict[
    Literal['ParallelTaskDataStagingBetweenStoragesProfile'],
    _ParallelTaskDataStagingBetweenStoragesProfilePayload,
]
type TraceReplayProfilePDict = _BaseProfilePDict[
    Literal['TraceReplayProfile'],
    _TraceReplayProfilePayload,
]


type ProfilePDict = (
    DelayProfilePDict
    | ParallelTaskProfilePDict
    | SequentialCompositionProfilePDict
    | ForkJoinCompositionProfilePDict
    | ParallelTaskMergeCompositionProfilePDict
    | ParallelTaskHomogeneousProfilePDict
    | ParallelTaskOnStorageHomogeneousProfilePDict
    | ParallelTaskDataStagingBetweenStoragesProfilePDict
    | TraceReplayProfilePDict
)
