# Copyright (c) 2022, California Institute of Technology and contributors
#
# You should have received a copy of the licensing terms for this
# software included in the file "LICENSE" located in the top-level
# directory of this package. If you did not, you can view a copy at
# https://git.ligo.org/ngdd/arrakis-python/-/raw/main/LICENSE

"""Channel information."""

from __future__ import annotations

import json
import re
from dataclasses import asdict, dataclass, fields
from functools import cached_property
from typing import TYPE_CHECKING, Any

import numpy

if TYPE_CHECKING:
    import pyarrow


CHANNEL_NAME_RE = re.compile(
    r"^(?P<domain>.*):(?P<subsystem>[^_-]*)((?P<delimiter>[_-])(?P<rest>.*))?$"
)


@dataclass(frozen=True)
class ParsedChannelName:
    """A parsed channel name

    Channel names have the following structure:

        <domain>:<subsystem>[-_]<rest>

    """

    domain: str
    subsystem: str
    subsystem_delimiter: str
    rest: str

    @classmethod
    def parse(cls, name: str) -> ParsedChannelName:
        """Parse a channel name into it's constituent parts"""
        if rem := CHANNEL_NAME_RE.match(name):
            return cls(
                rem["domain"],
                rem["subsystem"],
                rem["delimiter"],
                rem["rest"],
            )
        msg = f"Invalid channel name format: '{name}'"
        raise ValueError(msg)


@dataclass(frozen=True)
class Channel:
    """Metadata associated with a channel.

    Channels have the form {domain}:*.

    Parameters
    ----------
    name : str
        The name associated with this channel.
    data_type : str
        The data type associated with this channel.
    sample_rate : float
        The sampling rate associated with this channel.
    time : int, optional
        The timestamp when this metadata became active.
    publisher : str, optional
        The publisher associated with this channel.
    partition_id : str, optional
        The Kafka partition ID associated with this channel.
    partition_index : int, optional
        Partition index for the channel.  It is unique within the partition and
        allows the use of an integer value to identify the channel instead of a
        string.
    stride : int, optional
        Time duration in individual blocks of data for this channel,
        in nanoseconds.
    max_latency : int, optional
        Maximum expected publication latency for this channel's data,
        in seconds.

    """

    name: str
    data_type: str | numpy.dtype
    sample_rate: float
    time: int | None = None
    publisher: str | None = None
    partition_id: str | None = None
    partition_index: int | None = None
    stride: int | None = None
    max_latency: int | None = None
    # See schemas.METADATA_FIELDS for how these map to arrow fields

    @staticmethod
    def fields() -> tuple[str, ...]:
        """Channel field names"""
        return tuple(field.name for field in fields(Channel))

    def __post_init__(self) -> None:
        # evaluate the parsed name to validate, will throw ValueError
        # if invalid
        _ = self.parsed_name
        # handle numpy.dtype arguments
        assert self.data_type is not None, "data_type=None is forbidden."
        if isinstance(self.data_type, numpy.dtype):
            object.__setattr__(self, "data_type", self.data_type.name)

    @cached_property
    def parsed_name(self) -> ParsedChannelName:
        return ParsedChannelName.parse(self.name)

    @property
    def domain(self) -> str:
        return self.parsed_name.domain

    @property
    def subsystem(self) -> str:
        return self.parsed_name.subsystem

    def __repr__(self) -> str:
        s = f"<{self.name}, {self.sample_rate} Hz, {self.data_type!r}"
        for key, value in self.as_dict().items():
            if key in ["name", "sample_rate", "data_type"]:
                continue
            if value is None:
                continue
            s += f", {key}="
            if isinstance(value, int):
                s += f"{value:_}"
            else:
                s += f"{value!r}"
        s += ">"
        return s

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

    def __eq__(self, other) -> bool:
        # name, data type and sample rate are always required to match
        is_equal = (
            self.name == other.name
            and self.data_type == other.data_type
            and self.sample_rate == other.sample_rate
        )

        # optional fields match only if both are defined
        if self.time is not None and other.time is not None:
            is_equal &= self.time == other.time
        if self.publisher and other.publisher:
            is_equal &= self.publisher == other.publisher
        if self.partition_id and other.partition_id:
            is_equal &= self.partition_id == other.partition_id
        if self.partition_index and other.partition_index:
            is_equal &= self.partition_index == other.partition_index
        if self.stride and other.stride:
            is_equal &= self.stride == other.stride
        if self.max_latency and other.max_latency:
            is_equal &= self.max_latency == other.max_latency

        return is_equal

    def as_dict(self) -> dict[str, Any]:
        """Return metadata as "serialized" dict"""
        return asdict(self)

    @classmethod
    def from_json(cls, payload: str) -> Channel:
        """Create a Channel from its JSON representation.

        Parameters
        ----------
        payload : str
            The JSON-serialized channel.

        Returns
        -------
        Channel
            The newly created channel.

        """
        obj = json.loads(payload)
        return cls(**obj)

    @classmethod
    def from_field(cls, field: pyarrow.field) -> Channel:
        """Create a Channel from Arrow Flight field metadata.

        Parameters
        ----------
        field : pyarrow.field
            The channel field containing relevant metadata.

        Returns
        -------
        Channel
            The newly created channel.

        """
        data_type = _list_dtype_to_str(field.type)
        sample_rate = float(field.metadata[b"rate"].decode())
        return cls(field.name, data_type, sample_rate)


def _list_dtype_to_str(dtype: pyarrow.ListType) -> str:
    """Return a string representation of the list's inner data type.

    Note that this does not always match the string representation
    of Arrow's internal data types, to match the behavior across
    different languages.

    Parameters
    ----------
    dtype : pyarrow.ListType
        The list data type to inspect.

    Returns
    -------
    str
        A string representation of the list's inner data type.

    """
    inner_dtype = str(dtype.value_type)
    if inner_dtype == "float":
        return "float32"
    if inner_dtype == "double":
        return "float64"
    return inner_dtype
