import abc
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Tuple, Union

import csp
from csp import ts
from csp.adapters.status import Status
from pydantic import Field

from csp_gateway.server.gateway.csp import ChannelsType
from csp_gateway.server.gateway.gateway import GatewayModule
from csp_gateway.utils import ReadWriteMode, set_alarm_and_fetch_alarm_time

from . import ChannelSelection, EncodedEngineCycle, JSONConverter

__all__ = ("EngineReplay",)


class EngineReplay(GatewayModule, abc.ABC):
    """
    This module defines an abstract class that allows reading and writing
    csp engine cycle ticks of the channels within a gateway.

    **WRITE**
    When writing, every tick of the specified channels are encoded and written to some source.

    **READ**
    From a specified source, the encodings for the ticks of the specified channels are decoded, then propagated through the underlying csp graph of the gateway.
    """

    selection: ChannelSelection = Field(default_factory=ChannelSelection)
    encode_selection: Optional[ChannelSelection] = Field(
        default=None,
        description=("Optional selection that can be specified to override the selection to specify channels only for encoding."),
    )
    decode_selection: Optional[ChannelSelection] = Field(
        default=None,
        description=("Optional selection that can be specified to override the selection to specify channels only for decoding."),
    )
    read_write_mode: ReadWriteMode = ReadWriteMode.READ_AND_WRITE
    requires: Optional[ChannelSelection] = []
    start_writing: Union[datetime, timedelta] = Field(
        default=timedelta(),
        description="""
        This determines the csp time at which the module stops reading. If in WRITE mode or
        READ_AND_WRITE mode, then the module will start writing the
        cycle encodings to Kafka at that time.
        If none is supplied, the module will start writing from the start of the graph.
        """,
    )
    flag_updates: Dict[str, List[Tuple[str, bool]]] = Field(
        default_factory=dict,
        description="""
        A mapping of channels to a list of tuples each containing:
        A field name and the flag it should be updated to, when in Read mode.
        Only works for non-dict basket channels
        """,
    )
    log_lagging_engine_cycles: bool = Field(
        default=False,
        description="Whether we log when engine cycles are lagging from the current engine time.",
    )
    subscribe_with_struct_id: bool = Field(
        default=True,
        description=("If False, replaces the id field on GatewayStructs from the source with one autogenerated by the current Gateway."),
    )
    subscribe_with_struct_timestamp: bool = Field(
        default=True,
        description=("If False, replaces the timestamp field on the GatewayStruct with a timestamp autogenerated by the current Gateway."),
    )

    @abc.abstractclassmethod
    def subscribe(self) -> Union[ts[str], ts[Dict[str, Any]]]:
        """Subscribe to engine ticks from some source.
        Returns ts that will be decoded into an engine cycle and replayed in the current gateway."""
        ...

    @abc.abstractclassmethod
    def publish(self, encoded_channels: ts[EncodedEngineCycle]) -> None:
        """Publish to some source an ts of an EncodedEngineCycle representing an encoding of a csp engine cycle."""
        ...

    def connect(self, channels: ChannelsType):
        """:meta private:"""
        decode_channels_list = encode_channels_list = self.selection.select_from(channels)
        if self.decode_selection is not None:
            decode_channels_list = self.decode_selection.select_from(channels)

        if self.encode_selection is not None:
            encode_channels_list = self.encode_selection.select_from(channels)

        exclude = set()
        if not self.subscribe_with_struct_id:
            exclude.add("id")
        if not self.subscribe_with_struct_timestamp:
            exclude.add("timestamp")

        json_channel_converter = JSONConverter(
            channels=channels,
            decode_exclude_fields=exclude,
            decode_channels=decode_channels_list,
            encode_channels=encode_channels_list,
            flag_updates=self.flag_updates,
            log_lagging_engine_cycles=self.log_lagging_engine_cycles,
        )

        writing_time_info = set_alarm_and_fetch_alarm_time(self.start_writing)

        # DECODING
        if self.read_write_mode in [ReadWriteMode.READ, ReadWriteMode.READ_AND_WRITE]:
            fat_pipe = self.subscribe()
            read = csp.times(fat_pipe) < writing_time_info.alarm_time
            fat_pipe_filtered = csp.filter(read, fat_pipe)
            json_channel_converter.decode(fat_pipe_filtered)

        # ENCODING
        if self.read_write_mode in [ReadWriteMode.WRITE, ReadWriteMode.READ_AND_WRITE]:
            encoded_channels = json_channel_converter.encode(writing_time_info.alarm_ticked)
            self.publish(encoded_channels=encoded_channels)

    def status(self) -> ts[Status]:
        "Status must be expicitly defined. Should return a ts of csp.Status structs that represent the status of this module."
        raise NotImplementedError()
