#  Copyright 2025 Synnax Labs, Inc.
#
#  Use of this software is governed by the Business Source License included in the file
#  licenses/BSL.txt.
#
#  As of the Change Date specified in that file, in accordance with the Business Source
#  License, use of this software will be governed by the Apache License, Version 2.0,
#  included in the file licenses/APL.txt.

from typing import overload

import pandas as pd
from alamos import NOOP, Instrumentation
from freighter import AsyncStreamClient, StreamClient, UnaryClient, WebsocketClient

from synnax.channel.payload import (
    ChannelKey,
    ChannelKeys,
    ChannelName,
    ChannelNames,
    ChannelParams,
    ChannelPayload,
    normalize_channel_params,
)
from synnax.channel.retrieve import ChannelRetriever
from synnax.exceptions import QueryError
from synnax.framer.adapter import ReadFrameAdapter, WriteFrameAdapter
from synnax.framer.deleter import Deleter
from synnax.framer.frame import CrudeFrame, Frame
from synnax.framer.iterator import AUTO_SPAN, Iterator
from synnax.framer.streamer import AsyncStreamer, Streamer
from synnax.framer.writer import CrudeWriterMode, Writer, WriterMode
from synnax.ontology import ID
from synnax.telem import (
    CrudeSeries,
    CrudeTimeStamp,
    MultiSeries,
    TimeRange,
    TimeSpan,
)
from synnax.telem.control import Authority, CrudeAuthority

ontology_type = ID(type="framer")


class Client:
    """FramerClient provides interfaces for reading and writing segmented
    telemetry from a Synnax Cluster. SegmentClient should not be instantiated
    directly, but rather used through the synnax.Synnax class.
    """

    __stream_client: WebsocketClient
    __async_client: AsyncStreamClient
    __unary_client: UnaryClient
    __channels: ChannelRetriever
    __deleter: Deleter
    instrumentation: Instrumentation

    def __init__(
        self,
        stream_client: WebsocketClient,
        async_client: AsyncStreamClient,
        unary_client: UnaryClient,
        retriever: ChannelRetriever,
        deleter: Deleter,
        instrumentation: Instrumentation = NOOP,
    ):
        self.__stream_client = stream_client
        self.__async_client = async_client
        self.__unary_client = unary_client
        self.__channels = retriever
        self.__deleter = deleter
        self.instrumentation = instrumentation

    def open_writer(
        self,
        start: CrudeTimeStamp,
        channels: ChannelParams,
        authorities: CrudeAuthority | list[CrudeAuthority] = Authority.ABSOLUTE,
        *,
        name: str = "",
        strict: bool = False,
        suppress_warnings: bool = True,
        mode: CrudeWriterMode = WriterMode.PERSIST_STREAM,
        err_on_unauthorized: bool = False,
        enable_auto_commit: bool = True,
        auto_index_persist_interval: TimeSpan = 1 * TimeSpan.SECOND,
        err_on_extra_chans: bool = True,
        use_experimental_codec: bool = True,
    ) -> Writer:
        """Opens a new writer on the given channels.

        :param start: Sets the starting timestamp for the first sample in the writer. If
        this timestamp overlaps with existing data for ANY of the provided channels,
        the writer will fail to open.
        :param channels: The channels to write to. This can be a single channel name,
        a list of channel names, a single channel key, or a list of channel keys.
        :param authorities: The control authority to set for each channel on the writer.
        Defaults to absolute authority. If not working with concurrent control,
        it's best to leave this as the default.
        :param name: The name of the writer used in control subject.
        :param strict: Sets whether the writer will fail to write if the data for a
        particular channel does not exactly match this data type. When False,
        the default, the writer will automatically convert the data to the correct
        type if possible.
        :param suppress_warnings: Suppress various print warnings that may be emitted
        by the writer.
        :param mode: sets the persistence and streaming mode of the writer. The default
        mode is WriterModePersistStream. See the WriterMode documentation for more.
        :param err_on_unauthorized: sets whether the writer should return an error if
        it attempts to write to a channel it does not have control over.
        :param enable_auto_commit: determines whether the writer will automatically
        commit. If enable_auto_commit is true, then the writer will commit after each
        write, and will flush that commit to index after the specified
        auto_index_persist_interval.
        :param auto_index_persist_interval: interval at which commits to the index will
        be persisted. To persist every commit to guarantee minimal loss of data, set
        auto_index_persist_interval to AlwaysAutoIndexPersist.
        """
        adapter = WriteFrameAdapter(
            retriever=self.__channels,
            err_on_extra_chans=err_on_extra_chans,
            strict_data_types=strict,
            suppress_warnings=suppress_warnings,
        )
        adapter.update(channels)
        return Writer(
            start=start,
            adapter=adapter,
            client=self.__stream_client,
            authorities=authorities,
            name=name,
            mode=mode,
            err_on_unauthorized=err_on_unauthorized,
            enable_auto_commit=enable_auto_commit,
            auto_index_persist_interval=auto_index_persist_interval,
            use_experimental_codec=use_experimental_codec,
        )

    def open_iterator(
        self,
        tr: TimeRange,
        channels: ChannelParams,
        chunk_size: int = 1e5,
    ) -> Iterator:
        """Opens a new iterator over the given channels within the provided time range.

        :param channels: A list of channel keys to iterator over.
        :param tr: A time range to iterate over.
        :param chunk_size: The number of samples to read in a chunk with AutoSpan. Defaults to 100000
        :returns: An Iterator over the given channels within the provided time
        range. See the Iterator documentation for more.
        """
        adapter = ReadFrameAdapter(self.__channels)
        adapter.update(channels)
        return Iterator(
            tr=tr,
            adapter=adapter,
            client=self.__stream_client,
            chunk_size=chunk_size,
            instrumentation=self.instrumentation,
        )

    @overload
    def write(
        self,
        start: CrudeTimeStamp,
        frame: CrudeFrame,
        strict: bool = False,
    ): ...

    @overload
    def write(
        self,
        start: CrudeTimeStamp,
        channel: ChannelKey | ChannelName | ChannelPayload,
        data: CrudeSeries,
        strict: bool = False,
    ):
        """Writes telemetry to the given channel starting at the given timestamp.

        :param channel: The key of the channel to write to.
        :param start: The starting timestamp of the first sample in data.
        :param data: The telemetry to write to the channel.
        :returns: None.
        """
        ...

    @overload
    def write(
        self,
        start: CrudeTimeStamp,
        channel: ChannelKeys | ChannelNames | list[ChannelPayload],
        series: list[CrudeSeries],
        strict: bool = False,
    ): ...

    def write(
        self,
        start: CrudeTimeStamp,
        channels: ChannelParams | ChannelPayload | list[ChannelPayload] | CrudeFrame,
        series: CrudeSeries | list[CrudeSeries] | None = None,
        strict: bool = False,
    ):
        parsed_channels = list()
        if isinstance(channels, (list, ChannelKey, ChannelPayload, ChannelName)):
            parsed_channels = channels
        elif isinstance(channels, dict):
            parsed_channels = list(channels.keys())
        elif isinstance(channels, Frame):
            parsed_channels = channels.channels
        elif isinstance(channels, pd.DataFrame):
            parsed_channels = list(channels.columns)
        with self.open_writer(
            start=start,
            channels=parsed_channels,
            strict=strict,
            mode=WriterMode.PERSIST,
            err_on_unauthorized=True,
            auto_index_persist_interval=TimeSpan.MAX,
        ) as w:
            w.write(channels, series)

    @overload
    def read(
        self,
        tr: TimeRange,
        channels: ChannelKeys | ChannelNames,
    ) -> Frame: ...

    @overload
    def read(
        self,
        tr: TimeRange,
        channels: ChannelKey | ChannelName,
    ) -> MultiSeries: ...

    def read(
        self,
        tr: TimeRange,
        channels: ChannelParams,
    ) -> MultiSeries | Frame:
        """
        Reads telemetry from the channel between the two timestamps.

        :param tr: The time range to read from.
        :param channels: The key or name of the channel to read from.

        :returns: A tuple where the first item is a numpy array containing the telemetry
        and the second item is the time range occupied by that array.
        """
        normal = normalize_channel_params(channels)
        frame = self._read_frame(tr, channels)
        if len(normal.channels) > 1:
            return frame
        series = frame.get(normal.channels[0], None)
        if series is None:
            raise QueryError(
                f"""No data found for channel {normal.channels[0]} between {tr}"""
            )
        return series

    def read_latest(
        self,
        channels: ChannelKey | ChannelName,
        n: int = 1,
    ) -> MultiSeries: ...

    def read_latest(
        self,
        channels: ChannelKeys | ChannelNames,
        n: int = 1,
    ) -> Frame: ...

    def read_latest(
        self,
        channels: ChannelParams,
        n: int = 1,
    ) -> Frame:
        """
        Reads the latest n samples from time_channel and data_channel.

        Args:
            n: The number of samples to read.

        Returns:
            A frame containing the latest n samples from time_channel and data_channel
        """
        normal = normalize_channel_params(channels)
        aggregate = Frame()
        if n > 0:
            with self.open_iterator(
                tr=TimeRange.MAX,
                channels=channels,
                chunk_size=n,
            ) as i:
                i.seek_last()
                i.prev(AUTO_SPAN)
                aggregate.append(i.value)
        if len(normal.channels) > 1:
            return aggregate
        return aggregate.get(normal.channels[0], MultiSeries([]))

    def open_streamer(
        self,
        channels: ChannelParams,
        downsample_factor: int = 1,
        use_experimental_codec: bool = True,
    ) -> Streamer:
        """Opens a new streamer on the given channels. The streamer will immediately
        being receiving frames of data from the given channels.

        :param channels: The channels to stream from. This can be a single channel name,
        a list of channel names, a single channel key, or a list of channel keys.

        :param downsample_factor: The downsample factor to use for the streamer.
        """
        adapter = ReadFrameAdapter(self.__channels)
        adapter.update(channels)
        return Streamer(
            adapter=adapter,
            client=self.__stream_client,
            downsample_factor=downsample_factor,
            use_experimental_codec=use_experimental_codec,
        )

    async def open_async_streamer(
        self, channels: ChannelParams, downsample_factor: int = 1
    ) -> AsyncStreamer:
        adapter = ReadFrameAdapter(self.__channels)
        adapter.update(channels)
        s = AsyncStreamer(
            adapter=adapter,
            client=self.__async_client,
            downsample_factor=downsample_factor,
        )
        await s._open()
        return s

    def delete(self, channels: ChannelParams, tr: TimeRange) -> None:
        """
        delete deletes data in the specified channels in the specified time range.
        Note that the time range is start-inclusive and end-exclusive.
        Also note that deleting all data in a channel does not delete the channel; to
        delete a channel, use client.channels.delete().
        :param channels: channels to delete data from.
        :param tr: time range to delete data from.
        """
        self.__deleter.delete(channels, tr)

    def _read_frame(
        self,
        tr: TimeRange,
        channels: ChannelParams,
    ) -> Frame:
        aggregate = Frame()
        with self.open_iterator(tr, channels) as it:
            for fr in it:
                aggregate.append(fr)
        return aggregate
