import functools
import logging
from typing import Any, Iterator, List, Optional

import flatbuffers
import numpy as np
from physiq_cloud.series_frame import fb_builder
from physiq_cloud.series_frame.fb.math import ceildiv
from physiq_cloud.series_frame.fb_to_numpy import from_channel_data
from physiq_cloud.series_frame.fb_wrapper import FlatBufferWrapper
from physiq_cloud.time import Instant, TimeUnit
from physiqfb.ChannelData import ChannelData
from physiqfb.ReadingType import ReadingType
from physiqfb.SamplingSetData import SamplingSetData
from physiqfb.SeriesFrame import SeriesFrame
from pydantic import BaseModel

from pojagi_dsp.series_frame import (
    ASeriesFrameEmitter,
    SeriesFrameFilter,
    SeriesFramePackage,
    SftInfo,
)

logger = logging.getLogger(__name__)


class Synchronize(BaseModel):
    hour: bool = False
    minute: bool = False
    second: bool = False
    microsecond: bool = False


class Replaying(SeriesFrameFilter):
    """
    Filter that replays input frames starting at a given :class:`Instant`
    in time.
    """

    def __init__(
        self,
        emit_start: Instant = Instant.now(),
        synchronize: Optional[Synchronize] = None,
        **kwargs,
    ) -> None:
        """
        :param:`emit_start` - instant in time to replay the frames.
        :param:`synchronize` - Whether and at what granularity to synchronize the Instant from the source
        data with the configured :param:`emit_start`. This will alter the precision of the
        emit start seconds depending on the level of granularity you choose,
        possibly drastically.
        """
        super().__init__(**kwargs)

        self.frame_idx = 0

        self.emit_start = emit_start
        self.synchronize = synchronize
        self._first_frame_id: int | None = None
        self._start_frame_id_inst: int | None = None

    @ASeriesFrameEmitter.sft_info.setter
    def sft_info(self, val: SftInfo):
        """
        This override's :class:`ASeriesFrameFilter`'s :meth:`sft_info`
        setter, so that we can initialize the `timestamped_map` property.
        """
        self._sft_info = val
        self.timestamped_map = {
            s.id: s.timestamped for s in self.sft_info.sampling_sets
        }

    @property
    def start_frame_id(self):
        """
        :meth:`frames` must be called before accessing this for the first time,
        because the calculation depends on the first frame.
        """
        if self._start_frame_id_inst is not None:
            return self._start_frame_id_inst

        if not self.synchronize or not any(v for _, v in self.synchronize):
            self.synchronized_emit_start = self.emit_start
        else:
            emit_start_dt = self.emit_start.to_datetime()
            source_start_dt = Instant.from_unix(
                TimeUnit.MICROSECONDS,
                self._first_frame_id * self.sft_info.frame_size_micros,
            ).to_datetime()

            self.synchronized_emit_start = Instant.from_datetime(
                emit_start_dt.replace(
                    **{
                        k: getattr(source_start_dt, k)
                        if getattr(self.synchronize, k)
                        else getattr(emit_start_dt, k)
                        for k in ["hour", "minute", "second", "microsecond"]
                    },
                ),
            )

        # We use ceildiv to calculate the first frame ID, because normal
        # division would likely yield a frame BEFORE the emit_start. We want to
        # start emitting frames at or after emit_start, and ceildiv does that
        # for us.
        self._start_frame_id_inst = ceildiv(
            self.synchronized_emit_start.to_unix(
                TimeUnit.MICROSECONDS,
            ),
            self.sft_info.frame_size_micros,
        )

        return self._start_frame_id_inst

    def replayed_frame(
        self,
        source: SeriesFrame,
        target_frame_id: int,
    ) -> FlatBufferWrapper[SeriesFrame]:
        builder = flatbuffers.Builder(0)
        target = fb_builder.CreateSeriesFrame(
            builder=builder,
            frame_id=target_frame_id,
            ingested_at_micros=-1,  # source.IngestedAtMicros(),
            sampling_sets_offsets=[
                replayed_sampling_set(
                    builder=builder,
                    sampling_set_data=sampling_set_data,
                    source_frame_id=source.FrameId(),
                    target_frame_id=target_frame_id,
                    frame_size_micros=self.sft_info.frame_size_micros,
                    is_timestamped=self.timestamped_map[sampling_set_data.Id()],
                )
                for sampling_set_data in [
                    source.SamplingSets(idx)
                    for idx in range(source.SamplingSetsLength())
                ]
            ],
            submitter_id=-1,  # source.SubmitterId(),
            sensor_id=-1,  # source.SensorId(),
        )
        builder.Finish(target)
        target_frame_bytes = builder.Output()

        return FlatBufferWrapper(
            _bytes=target_frame_bytes,
            schema=SeriesFrame,
        )

    def frames(self) -> Iterator[SeriesFramePackage]:
        """
        Each time this property is accessed, it will access the injected
        reader anew, but continue to increment the `frame_idx` so that
        the next frame will replay where the last one left off.

        To reset the `frame_idx`, client code can just set the `frame_idx`
        manually (or create a new instance).
        """

        for pkg in self.reader.frames():
            if self._first_frame_id is None:
                self._first_frame_id = pkg.frame.fb.FrameId()

            yield SeriesFramePackage(
                frame=self.replayed_frame(
                    source=pkg.frame.fb,
                    target_frame_id=self.start_frame_id + self.frame_idx,
                ),
                sft_info=pkg.sft_info,
            )
            # `int` in python 3 is unbounded, can continue indefinitely. Also,
            # the `FrameId` is effectively a timestamp, so continuously
            # `inc`ing this value is by design.
            self.frame_idx += 1


def replayed_sampling_set(
    builder: flatbuffers.Builder,
    sampling_set_data: SamplingSetData,
    source_frame_id: int,
    target_frame_id: int,
    frame_size_micros: int,
    is_timestamped: bool,
):
    source_channels = [
        sampling_set_data.Channels(idx)
        for idx in range(sampling_set_data.ChannelsLength())
    ]

    target_channels: List[int] = list()
    for source_channel_data in source_channels:
        is_timestamp_channel = is_timestamped and source_channel_data.Id() == 0

        if not is_timestamp_channel or (source_frame_id == target_frame_id):
            # Just copy verbatim
            target_channel_data = replayed_channel_data(
                builder,
                source_channel_data,
            )
        else:
            # Update the timestamps relative to the target frameId
            source_frame_start_micros = source_frame_id * frame_size_micros
            target_frame_start_micros = target_frame_id * frame_size_micros

            # np.copy is used because the returned array is read-only
            timestamps = np.copy(from_channel_data(data=source_channel_data))
            timestamps -= source_frame_start_micros  # relativize
            timestamps += target_frame_start_micros  # frame shift

            target_channel_data = fb_builder.CreateChannelData(
                builder=builder,
                id=0,  # `source_channel_data.Id()` is constantly 0 here
                readings_offset=fb_builder.CreateInt64Channel(
                    builder=builder,
                    data=timestamps.tolist(),  # List[int]
                ),
                readings_type=ReadingType.INT64,
            )

        target_channels.append(target_channel_data)

    np_relative_timestamps = sampling_set_data.RelativeTimestampsAsNumpy()
    relative_timestamps = (
        np_relative_timestamps.tobytes()
        if np_relative_timestamps not in {0, -1}  # error conditions/no data
        else None
    )

    return fb_builder.CreateSamplingSetData(
        builder=builder,
        channel_offsets=target_channels,
        id=sampling_set_data.Id(),
        start_offset_micros=sampling_set_data.StartOffsetMicros(),
        relative_timestamps_type=sampling_set_data.RelativeTimestampsType(),
        relative_timestamps_unit_in_micros=sampling_set_data.RelativeTimestampsUnitInMicros(),
        relative_timestamps=relative_timestamps,
        relative_timestamps_sample_count=sampling_set_data.RelativeTimestampsSampleCount(),
    )


def replayed_channel_data(
    builder: flatbuffers.Builder,
    channel_data: ChannelData,
) -> int:
    readings_type = channel_data.ReadingsType()

    return fb_builder.CreateChannelData(
        builder=builder,
        id=channel_data.Id(),
        readings_offset=channel(
            builder=builder,
            readings_type=readings_type,
            readings=from_channel_data(data=channel_data).tolist(),
        ),
        readings_type=readings_type,
    )


reading_type_map = {
    ReadingType.BINARY32: fb_builder.CreateBinary32Channel,
    ReadingType.BINARY64: fb_builder.CreateBinary64Channel,
    ReadingType.INT8: fb_builder.CreateInt8Channel,
    ReadingType.INT16: fb_builder.CreateInt16Channel,
    ReadingType.INT32: fb_builder.CreateInt32Channel,
    ReadingType.INT64: fb_builder.CreateInt64Channel,
    ReadingType.NAMED_READING: fb_builder.CreateNamedReadingChannel,
    ReadingType.STRING: fb_builder.CreateStringChannel,
}


def channel(
    builder: flatbuffers.Builder,
    readings_type: int,
    readings: List[Any],
) -> int:
    return reading_type_map[readings_type](builder=builder, data=readings)
