#  Copyright 2026 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 __future__ import annotations

from collections.abc import Iterator, Sequence
from typing import TypeAlias, cast, overload

from pandas import DataFrame
from pydantic import BaseModel, Field

import synnax.channel.payload as channel
from synnax.exceptions import ValidationError
from synnax.telem import CrudeSeries, MultiSeries, Series, TypedCrudeSeries


class FramePayload(BaseModel):
    keys: list[channel.Key]
    series: list[Series]

    def __init__(
        self,
        keys: list[int] | None = None,
        series: list[Series] | None = None,
    ):
        if series is None:
            series = list()
        if keys is None:
            keys = list()
        super().__init__(series=series, keys=keys)


class Frame:
    """
    A frame is a collection of telemetry series mapped to particular channels. Frames
    can be keyed by channel name or channel key, but not both.
    """

    channels: list[channel.Key] | list[str]
    series: list[Series] = Field(default_factory=list)

    def __init__(
        self,
        channels: (
            list[channel.Key]
            | tuple[channel.Key]
            | list[str]
            | tuple[str]
            | DataFrame
            | Frame
            | FramePayload
            | dict[channel.Key, TypedCrudeSeries]
            | dict[str, TypedCrudeSeries]
            | None
        ) = None,
        series: Sequence[TypedCrudeSeries] | None = None,
    ):
        if isinstance(channels, Frame):
            self.channels = channels.channels
            self.series = channels.series
        elif isinstance(channels, FramePayload):
            self.channels = list(channels.keys)
            self.series = channels.series
        elif isinstance(channels, DataFrame):
            self.channels = channels.columns.to_list()
            self.series = [Series(data=channels[k]) for k in self.channels]
        elif isinstance(channels, dict):
            self.channels = cast(list[channel.Key] | list[str], list(channels.keys()))
            self.series = [Series(d) for d in channels.values()]
        elif (series is None or isinstance(series, list)) and (
            channels is None or isinstance(channels, list)
        ):
            self.series = list() if series is None else [Series(d) for d in series]
            self.channels = (
                cast(list[channel.Key] | list[str], list(channels))
                if channels
                else list()
            )
        else:
            raise ValueError(f"""
                [Frame] - invalid construction arguments. Received {channels}
                and {series}.
            """)

    def __str__(self) -> str:
        return self.to_df().__str__()

    @overload
    def append(self, col_or_frame: channel.Key | str, series: Series) -> None:
        """Adds a new series to the frame for the given channel.
        :param col_or_frame: The channel key or name to append the series to.
        :param series: The series to append to the frame
        """
        ...

    @overload
    def append(self, col_or_frame: Frame) -> None:
        """Appends the given frame to the current frame, modifying the current frame in place.
        :param col_or_frame: The frame to append to the current frame.
        """
        ...

    def append(
        self,
        col_or_frame: channel.Key | str | Frame,
        series: Series | None = None,
    ) -> None:
        """Appends a frame or series to the current frame. If a frame is provided, the
        series and channels from the frame are appended to the current frame. If a series
        is provided, the series is appended to the frame for the given channel.
        :param col_or_frame: The channel key or name to append the series to, or the frame
            to append to the current frame.
        :param series: The series to append to the frame, if col_or_frame is a channel key
            or name.
        """
        if isinstance(col_or_frame, Frame):
            self.series.extend(col_or_frame.series)
            self.channels.extend(col_or_frame.channels)  # type: ignore[arg-type]
        else:
            if series is None:
                raise ValueError("series must be provided when appending a channel")
            self.series.append(series)
            self.channels.append(col_or_frame)  # type: ignore[arg-type]

    def items(
        self,
    ) -> Iterator[tuple[channel.Key | str, Series]]:
        """
        Returns a generator of tuples containing the channel and series for each channel
        in the frame.
        """
        return zip(self.channels, self.series)

    def __getitem__(self, key: channel.Key | str) -> MultiSeries:
        if not isinstance(key, (channel.Key, str)):
            return self.to_df()[key]
        indexes = [i for i, k in enumerate(self.channels) if k == key]
        return MultiSeries([self.series[i] for i in indexes])

    def get(
        self, key: channel.Key | str, default: MultiSeries | None = None
    ) -> MultiSeries | None:
        """Gets the series for the given channel key or name. If the channel does not
        exist in the frame, returns the default value or None if no default is provided.
        :param key: The channel key or name to get the series for.
        :param default: The default value to return if the channel does not exist in the
        """
        try:
            return self[key]
        except ValueError:
            return default

    def to_payload(self) -> FramePayload:
        """Converts the frame to its payload representation for transport over the
        network. This method should typically only be used internally.
        :raises: ValidationError if the frame is keyed by channel name instead of key.
        """
        if not all(isinstance(k, channel.Key) for k in self.channels):
            diff = [k for k in self.channels if not isinstance(k, channel.Key)]
            raise ValidationError(f"""
            Cannot convert a frame that contains channel names to a payload.
            The following channels are invalid: {diff}
            """)
        keys: list[int] = [k for k in self.channels if isinstance(k, int)]
        return FramePayload(keys=keys, series=self.series)

    def to_df(self) -> DataFrame:
        """Converts the frame to a pandas DataFrame. Each column in the DataFrame
        corresponds to a channel in the frame.
        """
        base: dict[channel.Key | str, object] = dict()
        channels = cast(list[channel.Key | str], self.channels)
        for k in set(channels):
            v = self.get(k)
            if v is None:
                continue
            try:
                base[k] = v.__array__()
            except TypeError:
                base[k] = list(v)
        return DataFrame(base)

    def __contains__(self, key: channel.Key | str) -> bool:
        return key in self.channels


CrudeFrame: TypeAlias = (
    Frame | FramePayload | dict[channel.Key, CrudeSeries] | DataFrame
)
