"""
Workload-related data structures.
"""

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 (
    CustomMappingPDict,
    DelayProfilePDict,
    Duration,
    ForkJoinCompositionProfilePDict,
    JobId,
    JobPDict,
    MappingStrategyPDict,
    ParallelTaskDataStagingBetweenStoragesProfilePDict,
    ParallelTaskHomogeneousProfilePDict,
    ParallelTaskMergeCompositionProfilePDict,
    ParallelTaskOnStorageHomogeneousProfilePDict,
    ParallelTaskProfilePDict,
    PredefinedMappingPDict,
    ProfileId,
    ProfilePDict,
    SequentialCompositionProfilePDict,
    StorageId,
    Timestamp,
    TraceReplayProfilePDict,
)


class Job:
    # attributes defined at initialization
    id_: JobId
    resource_request: int
    walltime: Duration | None
    profile_id: ProfileId
    extra_data: str | None

    # attributes defined later
    submission_time: Timestamp | None = None  # set by Batsim
    allocation: ProcSet | None = None  # set by the EDC

    def __init__(
        self,
        id_: JobId,
        resource_request: int,
        profile_id: ProfileId,
        walltime: Duration | None = None,
        extra_data: str | None = None,
    ):
        self.id_ = id_
        self.resource_request = resource_request
        self.walltime = walltime
        self.profile_id = profile_id
        self.extra_data = extra_data

    def to_protocol_dict(self) -> JobPDict:
        protocol_dict: JobPDict = {
            'resource_request': self.resource_request,
            'profile_id': self.profile_id,
        }
        if self.walltime is not None:
            protocol_dict['walltime'] = self.walltime
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict


class PlacementPolicy:
    mapping_strategy: MappingStrategy
    # TODO: integrate profile_allocation_override and storage_placement

    def __init__(self, mapping_strategy: MappingStrategy):
        self.mapping_strategy = mapping_strategy

    @classmethod
    def default(cls) -> Self:
        """Return the default placement policy."""
        return cls(mapping_strategy=MappingStrategy.default())


# pybatsim representation of ExecutorPlacement
class MappingStrategy:
    @abstractmethod
    def to_protocol_dict(self) -> MappingStrategyPDict: ...

    @staticmethod
    def default() -> MappingStrategy:
        """Return the default mapping strategy."""
        return SpreadHostsMapping()


class SpreadHostsMapping(MappingStrategy):
    @override
    def to_protocol_dict(self) -> PredefinedMappingPDict:
        return {
            'executor_placement_type': 'PredefinedExecutorPlacementStrategyWrapper',
            'executor_placement': {
                'strategy': 'SpreadOverHostsFirst',
            },
        }


class FillHostsMapping(MappingStrategy):
    @override
    def to_protocol_dict(self) -> PredefinedMappingPDict:
        return {
            'executor_placement_type': 'PredefinedExecutorPlacementStrategyWrapper',
            'executor_placement': {
                'strategy': 'FillOneHostCoresFirst',
            },
        }


class CustomMapping(MappingStrategy):
    mapping: list[int]

    def __init__(self, mapping: list[int]):
        self.mapping = mapping

    @override
    def to_protocol_dict(self) -> CustomMappingPDict:
        return {
            'executor_placement_type': 'CustomExecutorToHostMapping',
            'executor_placement': {
                'mapping': self.mapping,
            },
        }


# class KillProgress
# TODO: implement me


class FinalState(Enum):
    SUCCESS = 'COMPLETED_SUCCESSFULLY'
    FAILED = 'COMPLETED_FAILED'
    WALLTIME_REACHED = 'COMPLETED_WALLTIME_REACHED'
    KILLED = 'COMPLETED_KILLED'
    REJECTED = 'REJECTED'


class Profile:
    _batprotocol_profiles: Final[dict[str, type[Profile]]] = {}

    id_: ProfileId
    extra_data: str | None

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

    def __init__(self, id_: ProfileId, extra_data: str | None = None):
        self.id_ = id_
        self.extra_data = extra_data

    # Serialization and deserialization are made against ProfileAndId structure.

    @abstractmethod
    def to_protocol_dict(self) -> ProfilePDict: ...

    # 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: ProfilePDict) -> Profile:
        concrete_cls = cls._batprotocol_profiles[payload['profile_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 DelayProfile(Profile):
    delay: Duration

    @override
    def __init__(self, id_: ProfileId, delay: Duration, extra_data: str | None = None):
        super().__init__(id_, extra_data)
        self.delay = delay

    @override
    def to_protocol_dict(self) -> DelayProfilePDict:
        protocol_dict: DelayProfilePDict = {
            'id': self.id_,
            'profile_type': 'DelayProfile',
            'profile': {
                'delay': self.delay,
            },
        }
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(cls, payload: DelayProfilePDict) -> Self:
        return cls(
            id_=payload['id'],
            delay=payload['profile']['delay'],
            extra_data=payload.get('extra_data'),
        )


class ParallelTaskProfile(Profile):
    computation_vector: list[float] | None
    communication_matrix: list[float] | None

    @override
    def __init__(
        self,
        id_: ProfileId,
        computation_vector: list[float] | None,
        communication_matrix: list[float] | None,
        extra_data: str | None = None,
    ):
        super().__init__(id_, extra_data)
        self.computation_vector = computation_vector
        self.communication_matrix = communication_matrix

    @override
    def to_protocol_dict(self) -> ParallelTaskProfilePDict:
        protocol_dict: ParallelTaskProfilePDict = {
            'id': self.id_,
            'profile_type': 'ParallelTaskProfile',
            'profile': {},
        }
        if self.computation_vector is not None:
            protocol_dict['profile']['computation_vector'] = self.computation_vector
        if self.communication_matrix is not None:
            protocol_dict['profile']['communication_matrix'] = self.communication_matrix
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(cls, payload: ParallelTaskProfilePDict) -> Self:
        return cls(
            id_=payload['id'],
            computation_vector=payload['profile'].get('computation_vector'),
            communication_matrix=payload['profile'].get('communication_matrix'),
            extra_data=payload.get('extra_data'),
        )


class SequentialCompositionProfile(Profile):
    repetition_count: int = 1
    profile_ids: list[ProfileId]

    @override
    def __init__(
        self,
        id_: ProfileId,
        profile_ids: list[ProfileId],
        repetition_count: int = 1,
        extra_data: str | None = None,
    ):
        super().__init__(id_, extra_data)
        self.repetition_count = repetition_count
        self.profile_ids = profile_ids

    @override
    def to_protocol_dict(self) -> SequentialCompositionProfilePDict:
        protocol_dict: SequentialCompositionProfilePDict = {
            'id': self.id_,
            'profile_type': 'SequentialCompositionProfile',
            'profile': {
                'repetition_count': 1,
                'profile_ids': self.profile_ids,
            },
        }
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: SequentialCompositionProfilePDict,
    ) -> Self:
        return cls(
            id_=payload['id'],
            repetition_count=payload['profile']['repetition_count'],
            profile_ids=payload['profile']['profile_ids'],
            extra_data=payload.get('extra_data'),
        )


class ForkJoinCompositionProfile(Profile):
    profile_ids: list[ProfileId]

    @override
    def __init__(
        self,
        id_: ProfileId,
        profile_ids: list[ProfileId],
        extra_data: str | None = None,
    ):
        super().__init__(id_, extra_data)
        self.profile_ids = profile_ids

    @override
    def to_protocol_dict(self) -> ForkJoinCompositionProfilePDict:
        protocol_dict: ForkJoinCompositionProfilePDict = {
            'id': self.id_,
            'profile_type': 'ForkJoinCompositionProfile',
            'profile': {
                'profile_ids': self.profile_ids,
            },
        }
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: ForkJoinCompositionProfilePDict,
    ) -> Self:
        return cls(
            id_=payload['id'],
            profile_ids=payload['profile']['profile_ids'],
            extra_data=payload.get('extra_data'),
        )


class ParallelTaskMergeCompositionProfile(Profile):
    profile_ids: list[ProfileId]

    @override
    def __init__(
        self,
        id_: ProfileId,
        profile_ids: list[ProfileId],
        extra_data: str | None = None,
    ):
        super().__init__(id_, extra_data)
        self.profile_ids = profile_ids

    @override
    def to_protocol_dict(self) -> ParallelTaskMergeCompositionProfilePDict:
        protocol_dict: ParallelTaskMergeCompositionProfilePDict = {
            'id': self.id_,
            'profile_type': 'ParallelTaskMergeCompositionProfile',
            'profile': {
                'profile_ids': self.profile_ids,
            },
        }
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: ParallelTaskMergeCompositionProfilePDict,
    ) -> Self:
        return cls(
            id_=payload['id'],
            profile_ids=payload['profile']['profile_ids'],
            extra_data=payload.get('extra_data'),
        )


class HomogeneousProfileGenerationStrategy(Enum):
    AMOUNT_UNIT = 'DefinedAmountsUsedForEachValue'
    AMOUNT_GLOBAL = 'DefinedAmountsSpreadUniformly'


class ParallelTaskHomogeneousProfile(Profile):
    computation_amount: float
    communication_amount: float
    generation_strategy: HomogeneousProfileGenerationStrategy

    @override
    def __init__(
        self,
        id_: ProfileId,
        computation_amount: float,
        communication_amount: float,
        generation_strategy: HomogeneousProfileGenerationStrategy,
        extra_data: str | None = None,
    ):
        super().__init__(id_, extra_data)
        self.computation_amount = computation_amount
        self.communication_amount = communication_amount
        self.generation_strategy = generation_strategy

    @override
    def to_protocol_dict(self) -> ParallelTaskHomogeneousProfilePDict:
        protocol_dict: ParallelTaskHomogeneousProfilePDict = {
            'id': self.id_,
            'profile_type': 'ParallelTaskHomogeneousProfile',
            'profile': {
                'computation_amount': self.computation_amount,
                'communication_amount': self.communication_amount,
                'generation_strategy': self.generation_strategy.value,
            },
        }
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: ParallelTaskHomogeneousProfilePDict,
    ) -> Self:
        return cls(
            id_=payload['id'],
            computation_amount=payload['profile'].get('computation_amount'),
            communication_amount=payload['profile'].get('communication_amount'),
            generation_strategy=HomogeneousProfileGenerationStrategy(
                payload['profile']['generation_strategy']
            ),
            extra_data=payload.get('extra_data'),
        )


class ParallelTaskOnStorageHomogeneousProfile(Profile):
    storage_name: StorageId
    bytes_to_read: float
    bytes_to_write: float
    generation_strategy: HomogeneousProfileGenerationStrategy

    @override
    def __init__(
        self,
        id_: ProfileId,
        storage_name: StorageId,
        bytes_to_read: float,
        bytes_to_write: float,
        generation_strategy: HomogeneousProfileGenerationStrategy,
        extra_data: str | None = None,
    ):
        super().__init__(id_, extra_data)
        self.storage_name = storage_name
        self.bytes_to_read = bytes_to_read
        self.bytes_to_write = bytes_to_write
        self.generation_strategy = generation_strategy

    @override
    def to_protocol_dict(self) -> ParallelTaskOnStorageHomogeneousProfilePDict:
        protocol_dict: ParallelTaskOnStorageHomogeneousProfilePDict = {
            'id': self.id_,
            'profile_type': 'ParallelTaskOnStorageHomogeneousProfile',
            'profile': {
                'storage_name': self.storage_name,
                'bytes_to_read': self.bytes_to_read,
                'bytes_to_write': self.bytes_to_write,
                'generation_strategy': self.generation_strategy.value,
            },
        }
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: ParallelTaskOnStorageHomogeneousProfilePDict,
    ) -> Self:
        return cls(
            id_=payload['id'],
            storage_name=payload['profile']['storage_name'],
            bytes_to_read=payload['profile']['bytes_to_read'],
            bytes_to_write=payload['profile']['bytes_to_write'],
            generation_strategy=HomogeneousProfileGenerationStrategy(
                payload['profile']['generation_strategy']
            ),
            extra_data=payload.get('extra_data'),
        )


class ParallelTaskDataStagingBetweenStoragesProfile(Profile):
    bytes_to_transfer: float
    emitter_storage_name: StorageId
    receiver_storage_name: StorageId

    @override
    def __init__(
        self,
        id_: ProfileId,
        bytes_to_transfer: float,
        emitter_storage_name: StorageId,
        receiver_storage_name: StorageId,
        extra_data: str | None = None,
    ):
        super().__init__(id_, extra_data)
        self.bytes_to_transfer = bytes_to_transfer
        self.emitter_storage_name = emitter_storage_name
        self.receiver_storage_name = receiver_storage_name

    @override
    def to_protocol_dict(self) -> ParallelTaskDataStagingBetweenStoragesProfilePDict:
        protocol_dict: ParallelTaskDataStagingBetweenStoragesProfilePDict = {
            'id': self.id_,
            'profile_type': 'ParallelTaskDataStagingBetweenStoragesProfile',
            'profile': {
                'bytes_to_transfer': self.bytes_to_transfer,
                'emitter_storage_name': self.emitter_storage_name,
                'receiver_storage_name': self.receiver_storage_name,
            },
        }
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(
        cls,
        payload: ParallelTaskDataStagingBetweenStoragesProfilePDict,
    ) -> Self:
        return cls(
            id_=payload['id'],
            bytes_to_transfer=payload['profile']['bytes_to_transfer'],
            emitter_storage_name=payload['profile']['emitter_storage_name'],
            receiver_storage_name=payload['profile']['receiver_storage_name'],
            extra_data=payload.get('extra_data'),
        )


class TraceType(Enum):
    SMPI = 'SMPI'
    FRACTIONAL_COMPUTATION = 'FractionalComputation'


class TraceReplayProfile(Profile):
    trace_type: TraceType
    filename: str

    @override
    def __init__(
        self,
        id_: ProfileId,
        trace_type: TraceType,
        filename: str,
        extra_data: str | None = None,
    ):
        super().__init__(id_, extra_data)
        self.trace_type = trace_type
        self.filename = filename

    @override
    def to_protocol_dict(self) -> TraceReplayProfilePDict:
        protocol_dict: TraceReplayProfilePDict = {
            'id': self.id_,
            'profile_type': 'TraceReplayProfile',
            'profile': {
                'trace_type': self.trace_type.value,
                'filename': self.filename,
            },
        }
        if self.extra_data is not None:
            protocol_dict['extra_data'] = self.extra_data
        return protocol_dict

    @override
    @classmethod
    def _deserialize_protocol_dict(cls, payload: TraceReplayProfilePDict) -> Self:
        return cls(
            id_=payload['id'],
            trace_type=TraceType(payload['profile']['trace_type']),
            filename=payload['profile']['filename'],
            extra_data=payload.get('extra_data'),
        )
