# 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-server/-/raw/main/LICENSE

from __future__ import annotations

import logging
import re
from collections import defaultdict
from collections.abc import Iterable, Mapping
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any, Protocol

import numpy
import toml
from arrakis import Channel

from .partition import PartitionMetadata, partition_channels

logger = logging.getLogger("arrakis")

PARTITION_METADATA_KEY = "__internal_partition_metadata"
PUBLISHER_KEY = "publisher"


@dataclass(frozen=True)
class PublisherInfo:
    """Publisher-level configuration.

    Parameters
    ----------
    publisher_id : str
        Unique identifier for the publisher.
    stride : int
        Default stride (nanoseconds) for channels from this publisher.
    max_latency : int
        Default max latency (nanoseconds) for channels from this publisher.
    dynamic : bool
        Whether this publisher can dynamically partition channels
        at runtime via the partition endpoint.

    """

    publisher_id: str
    stride: int
    max_latency: int
    dynamic: bool = False


class ChannelMetadataBackend(Protocol):
    def update(
        self,
        channels: Iterable[Channel],
        partition_metadata: Mapping[str, PartitionMetadata],
    ) -> None:
        """Update channel metadata.

        Parameters
        ----------
        channels : Iterable[Channel]
            Channels for which to update metadata with.
        partition_metadata : Mapping[str, PartitionMetadata]
            Per partition metadata.
        """
        ...

    def load(cls, *args, **kwargs) -> list[Channel]:
        """Load channel metadata"""
        ...

    def find(
        self,
        *,
        pattern: str,
        data_type: list[str],
        min_rate: int,
        max_rate: int,
        publisher: list[str],
    ) -> Iterable[Channel]:
        """Find channels matching a set of conditions.

        Parameters
        ----------
        pattern : str
            Channel pattern to match channels with, using regular expressions.
        data_type : list[str]
            Data types to match.
        min_rate : int
            Minimum sampling rate for channels.
        max_rate : int
            Maximum sampling rate for channels.
        publisher : list[str]
            Publishers to match.

        Returns
        -------
        Iterable[Channel]
            Channel objects for all channels matching query.

        """
        ...

    def get_publisher_info(self, publisher_id: str) -> PublisherInfo | None:
        """Get publisher info by publisher ID.

        Parameters
        ----------
        publisher_id : str
            The publisher ID to look up.

        Returns
        -------
        PublisherInfo or None
            The publisher info if found, None otherwise.

        """
        ...

    def __getitem__(self, name):
        """Get channel metadata for individual channel"""
        ...

    def describe(self, *, channels: Iterable[str]) -> Iterable[Channel]:
        """Get channel metadata for channels requested.

        Parameters
        ----------
        channels : Iterable[str]
            Channels to request.

        Returns
        -------
        Channel
            Channel objects, one per channel requested.

        """
        return [self[channel] for channel in channels]


class ChannelConfigBackend(ChannelMetadataBackend):
    """A channel metadata backend backed by configuration.

    Channel metadata is stored as a configuration file in the TOML format.

    """

    def __init__(
        self,
        cache_file: Path | None = None,
        enforce: Iterable[str] | None = None,
    ):
        """initialize backend

        Parameters
        ----------
        cache_file : Path, optional
            Path to file that will hold channel metadata for
            publishers that are publishing their own channel lists.
            Or None.
        enforce : list[str], optional
            A list of Channel properties to enforce the presence of
            (beyond the default of ["sample_rate", "data_type"]).

        """
        self.cache_file = cache_file
        # always enforced parameters
        self.enforce = {"sample_rate", "data_type"}
        if enforce:
            self.enforce |= set(enforce)

        self.metadata: dict[str, Channel] = {}
        self.partition_metadata: dict[str, PartitionMetadata] = {}
        self.extra: dict[str, dict[str, Any]] = {}
        self.publisher_info: dict[str, PublisherInfo] = {}
        # used for tracking if channels have been updated
        self._updated: dict[str, bool] = {}

        if self.cache_file is not None and self.cache_file.exists():
            self.load(self.cache_file)
            # reset the load tracking after loading the cache
            self._updated = {}

    def __getitem__(self, name):
        """channel metadata by channel name"""
        return self.metadata[name]

    def get_publisher_info(self, publisher_id: str) -> PublisherInfo | None:
        """Get publisher info by publisher ID."""
        return self.publisher_info.get(publisher_id)

    def validate(self, channel: Channel) -> None:
        """return True if channel contains all enforced properties

        Raises a ValueError if the enforced properties are not present
        or None.

        """
        for prop in self.enforce:
            if not hasattr(channel, prop) or not getattr(channel, prop):
                raise ValueError(f"channel '{channel.name}' missing property {prop}")

    def update(
        self,
        channels: Iterable[Channel],
        partition_metadata: Mapping[str, PartitionMetadata],
        overwrite: bool = False,
    ):
        """Update channel metadata and sync to cache

        Channels will be validated according to the enforced property
        list.

        Parameters
        ----------
        channels : list[Channel]
            List of channels to update.
        partition_metadata : Mapping[str, PartitionMetadata]
            Per partition metadata.
        overwrite : bool
            Whether to allow overwriting existing channels or not
            (default: False)

        """
        # update in-memory channel map
        for channel in channels:
            if not overwrite and self._updated.get(channel.name, False):
                raise ValueError(
                    f"attempt to overwrite existing channel: {channel.name}"
                )
            self.validate(channel)
            self.metadata[channel.name] = channel
            self._updated[channel.name] = True
        for partition in partition_metadata:
            if partition in self.partition_metadata:
                self.partition_metadata[partition].merge_in(
                    partition_metadata[partition]
                )
            else:
                self.partition_metadata[partition] = partition_metadata[partition]

        if not self.cache_file:
            return

        # write updated channel map to disk
        metadata = {}
        for name, channel in self.metadata.items():
            metadata[name] = channel.as_dict()
            del metadata[name]["name"]
            if self.extra and name in self.extra:
                metadata[name].update(self.extra[name])
        metadata[PARTITION_METADATA_KEY] = {
            k: partition_metadata[k].to_dict() for k in partition_metadata
        }

        if self.cache_file is not None:
            logger.info("saving metadata to cache %s", self.cache_file)
            with self.cache_file.open("w") as f:
                toml.dump(metadata, f)

    def load(
        self,
        path: Path,
        overwrite: bool = False,
    ) -> list[Channel]:
        """Load channel metadata from TOML file.

        If a ``[publisher]`` section is present, the file is treated as
        a publisher configuration: publisher-level fields are applied to
        all channels, partitions are assigned, and stale cached channels
        for that publisher are purged (static) or reconciled (dynamic).

        If no ``[publisher]`` section is present, the file is treated as
        a cache file and channels are loaded as-is.

        Parameters
        ----------
        path : Path
            Path to channel metadata toml file.
        overwrite : bool
            Whether to allow overwriting existing channels or not
            (default: False)

        Returns
        -------
        List[Channel]
            List of channels loaded from file.

        """
        logger.info("loading channel description file: %s", path)
        metadata = {}
        with path.open("r") as f:
            metadata = toml.load(f)

        # common metadata for all channels defined in a "common" block
        common = metadata.pop("common", {})
        publisher_section = metadata.pop(PUBLISHER_KEY, {})
        tmp = metadata.pop(PARTITION_METADATA_KEY, {})
        partition_metadata = {k: PartitionMetadata.from_dict(tmp[k]) for k in tmp}
        del tmp

        # store publisher info if present
        pub_info: PublisherInfo | None = None
        if publisher_section:
            pub_info = PublisherInfo(
                publisher_id=publisher_section["publisher_id"],
                stride=publisher_section["stride"],
                max_latency=publisher_section["max_latency"],
                dynamic=publisher_section.get("dynamic", False),
            )
            self.publisher_info[pub_info.publisher_id] = pub_info

        channels = []
        extra = {}
        for name, meta in metadata.items():
            cmeta = {
                field: meta.pop(field, common.get(field, None))
                for field in Channel.fields()
            }
            # publisher-level fields override channel/common values
            if publisher_section:
                cmeta["stride"] = publisher_section["stride"]
                cmeta["max_latency"] = publisher_section["max_latency"]
                cmeta["publisher"] = publisher_section["publisher_id"]
            cmeta["name"] = name
            channel = Channel(**cmeta)
            channels.append(channel)
            extra[name] = meta

        if pub_info is not None:
            publisher_id = pub_info.publisher_id

            channels = partition_channels(
                channels,
                metadata=self.metadata,
                partition_metadata=partition_metadata,
                publisher=publisher_id,
            )

            if pub_info.dynamic:
                # reconcile any cached channels for this publisher
                # with current publisher info
                for name, channel in self.metadata.items():
                    if channel.publisher != publisher_id:
                        continue
                    if (
                        channel.stride == pub_info.stride
                        and channel.max_latency == pub_info.max_latency
                    ):
                        continue
                    self.metadata[name] = replace(
                        channel,
                        stride=pub_info.stride,
                        max_latency=pub_info.max_latency,
                    )
            else:
                # purge any old channels for this publisher that are
                # no longer active
                channel_names = {channel.name for channel in channels}
                stale_channels = set()
                for name, channel in self.metadata.items():
                    if channel.publisher == publisher_id and name not in channel_names:
                        stale_channels.add(name)
                for name in stale_channels:
                    del self.metadata[name]

        self.update(channels, partition_metadata, overwrite=overwrite)
        self.extra.update(extra)

        return channels

    @property
    def scopes(self) -> dict[str, list[dict[str, Any]]]:
        """The scopes that the set of channels span."""
        scopes: dict[str, list[dict[str, Any]]] = defaultdict(list)
        for channel in self.metadata.values():
            scopes[channel.domain].append({"subsystem": channel.subsystem})
        return scopes

    def find(
        self,
        *,
        pattern: str,
        data_type: list[str],
        min_rate: int,
        max_rate: int,
        publisher: list[str],
    ) -> Iterable[Channel]:
        expr = re.compile(pattern)
        channels = []
        dtypes = set(data_type)
        publishers = set(publisher)
        for channel in self.metadata.values():
            if expr.match(channel.name):
                rate = channel.sample_rate
                if not (rate >= min_rate and rate <= max_rate):
                    continue
                if dtypes and channel.data_type not in dtypes:
                    continue
                if publishers and channel.publisher not in publishers:
                    continue
                channels.append(channel)

        return channels
