# -*- coding: utf-8 -*-
"""
SeedLink request client for ObsPy.

:copyright:
    The ObsPy Development Team (devs@obspy.org)
:license:
    GNU Lesser General Public License, Version 3
    (https://www.gnu.org/copyleft/lesser.html)
"""
import fnmatch
import warnings

import numpy as np
from lxml import etree

from obspy import Stream
from .slclient import SLClient, SLPacket
from .client.seedlinkconnection import SeedLinkConnection
from obspy.io.mseed.core import _read_mseed


class Client(object):
    """
    SeedLink request client.

    This client is intended for requests of specific, finite time windows.
    To work with continuous realtime data streams please see
    :class:`~obspy.clients.seedlink.slclient.SLClient` and
    :class:`~obspy.clients.seedlink.easyseedlink.EasySeedLinkClient`.

    :type server: str
    :param server: Server name or IP address to connect to (e.g.
        "localhost", "rtserver.ipgp.fr")
    :type port: int
    :param port: Port at which the seedlink server is operating (default is
        `18000`).
    :type timeout: float
    :param timeout: Network timeout for low-level network connection in
        seconds.
    :type debug: bool
    :param debug: Switches on debugging output.
    """
    def __init__(self, server, port=18000, timeout=20, debug=False):
        """
        Initializes the SeedLink request client.
        """
        self.timeout = timeout
        self.debug = debug
        self._server_url = "%s:%i" % (server, port)
        self._station_cache = None
        self._station_cache_level = None

    def _init_client(self):
        """
        Make fresh connection to seedlink server

        Should be done before any request to server, since SLClient keeps
        things like multiselect etc for subsequent requests
        """
        self._slclient = SLClient(timeout=self.timeout)

    def _connect(self):
        """
        Open new connection to seedlink server.
        """
        self._slclient.slconn = SeedLinkConnection(timeout=self.timeout)
        self._slclient.slconn.set_sl_address(self._server_url)
        self._slclient.slconn.netto = self.timeout

    def get_waveforms(self, network, station, location, channel, starttime,
                      endtime):
        """
        Request waveform data from the seedlink server.

        >>> from obspy import UTCDateTime
        >>> client = Client('rtserver.ipgp.fr')
        >>> t = UTCDateTime() - 1500
        >>> st = client.get_waveforms("G", "FDFM", "10", "BHZ", t, t + 5)
        >>> print(st)  # doctest: +ELLIPSIS
        1 Trace(s) in Stream:
        G.FDFM.10.BHZ | 20... | 20.0 Hz, ... samples

        Most servers support '?' single-character wildcard in location and
        channel code fields:

        >>> st = client.get_waveforms("G", "FDFM", "1?", "B??", t, t + 5)
        >>> st = st.sort(reverse=True)
        >>> print(st)  # doctest: +ELLIPSIS
        3 Trace(s) in Stream:
        G.FDFM.10.BHZ | 20... | 20.0 Hz, ... samples
        G.FDFM.10.BHN | 20... | 20.0 Hz, ... samples
        G.FDFM.10.BHE | 20... | 20.0 Hz, ... samples

        Depending on server capabilities, '*' multi-character wildcards might
        work in any parameter:

        >>> st = client.get_waveforms("*", "FDFM", "1*", "B*", t, t + 5)
        >>> st = st.sort(reverse=True)
        >>> print(st)  # doctest: +ELLIPSIS
        3 Trace(s) in Stream:
        G.FDFM.10.BHZ | 20... | 20.0 Hz, ... samples
        G.FDFM.10.BHN | 20... | 20.0 Hz, ... samples
        G.FDFM.10.BHE | 20... | 20.0 Hz, ... samples

        .. note::

            Support of wildcards strongly depends on the queried seedlink
            server. In general, '?' as single character wildcard seems to work
            well in location code and channel code fields for most servers.
            Usage of '*' relies on the server supporting info requests on
            station or even channel level, see :meth:`Client.get_info()`.

        :type network: str
        :param network: Network code. See note on wildcards above.
        :type station: str
        :param station: Station code. See note on wildcards above.
        :type location: str
        :param location: Location code. See note on wildcards above.
        :type channel: str
        :param channel: Channel code. See note on wildcards above.
        :type starttime: :class:`~obspy.core.utcdatetime.UTCDateTime`
        :param starttime: Start time of requested time window.
        :type endtime: :class:`~obspy.core.utcdatetime.UTCDateTime`
        :param endtime: End time of requested time window.
        """
        # need to do an info request?
        if any('*' in x for x in (network, station, location, channel)) \
                or any('?' in x for x in (network, station)):
            # need to do an info request on channel level?
            if any('*' in x for x in (location, channel)):
                info = self.get_info(network=network, station=station,
                                     location=location, channel=channel,
                                     level='channel', cache=True)
                multiselect = ["%s_%s:%s%s" % (net, sta, loc, cha)
                               for net, sta, loc, cha in info]
            # otherwise keep location/channel wildcards and do request on
            # station level only
            else:
                info = self.get_info(network=network, station=station,
                                     level='station', cache=True)
                multiselect = ["%s_%s:%s%s" % (net, sta, location, channel)
                               for net, sta in info]
            multiselect = ','.join(multiselect)
            return self._multiselect_request(multiselect, starttime, endtime)

        # if no info request is needed, we just work with the given input
        # (might have some '?' wildcards in loc/cha)
        if len(location) > 2:
            msg = ("Location code ('%s') only supports a maximum of 2 "
                   "characters.") % location
            raise ValueError(msg)
        elif len(location) == 1:
            msg = ("Single character location codes that are not an '*' are "
                   "untested.")
            warnings.warn(msg)
        if location:
            loccha = "%2s%3s" % (location, channel)
        else:
            loccha = channel
        seedlink_id = "%s_%s:%s" % (network, station, loccha)
        return self._multiselect_request(seedlink_id, starttime, endtime)

    def _multiselect_request(self, multiselect, starttime, endtime):
        """
        Make a multiselect request to underlying seedlink client

        Multiselect string is one or more comma separated
        network/station/location/channel combinations as defined by seedlink
        standard, e.g.
        "NETWORK_STATION:LOCATIONCHANNEL,NETWORK_STATION:LOCATIONCHANNEL"
        where location+channel may contain '?' characters but should be exactly
        5 characters long.

        :rtype: :class:`~obspy.core.stream.Stream`
        """
        self._init_client()
        self._slclient.multiselect = multiselect
        self._slclient.begin_time = starttime
        self._slclient.end_time = endtime
        self._connect()
        self._slclient.initialize()
        self.stream = Stream()

        # holds all data received by the client
        self.bulk_data_container = list()

        self._slclient.run(packet_handler=self._fast_packet_handler)

        # process the data all in one go, much more efficient
        input = np.concatenate(self.bulk_data_container)
        self.bulk_data_container = list()
        stream = _read_mseed(input)

        stream.trim(starttime, endtime)
        self.stream = None
        return stream

    def get_info(self, network=None, station=None, location=None, channel=None,
                 level='station', cache=True, warn_on_excluded_stations=False):
        """
        Request available stations information from the seedlink server.

        Supports ``fnmatch`` wildcards, e.g. ``*`` and ``?``, in ``network``,
        ``station``, ``location`` and ``channel``.

        >>> client = Client('rtserver.ipgp.fr')
        >>> info = client.get_info(station="FDFM")
        >>> print(info)
        [('G', 'FDFM')]
        >>> info = client.get_info(
        ...     station="FD?M", channel='[LBH]HZ', location='10',
        ...     level='channel')
        >>> print(info)  # doctest: +NORMALIZE_WHITESPACE
        [('G', 'FDFM', '10', 'BHZ'), ('G', 'FDFM', '10', 'HHZ'),
         ('G', 'FDFM', '10', 'LHZ')]

        Available station information is cached after the first request to the
        server, so use ``cache=False`` on subsequent requests if there is a
        need to force fetching new information from the server (should only
        concern programs running in background for a very long time).

        .. note::
            Stations/channels are excluded from the results for which the
            server indicates it is serving them in general but it also states
            no data are in ring buffer currently.
            If interested in these "no data" stations/channels, either set
            ``warn_on_excluded_stations=True`` which will show a warning
            message with excluded stations or use ``debug=True`` when
            initializing the client which will print the raw server ``seedlink
            INFO`` xml response which will show these stations listed with
            ``begin_seq`` and ``end_seq`` both with value ``'000000'``.

        :type network: str
        :param network: Network code. Supports ``fnmatch`` wildcards, e.g.
            ``*`` and ``?``.
        :type station: str
        :param station: Station code. Supports ``fnmatch`` wildcards, e.g.
            ``*`` and ``?``.
        :type location: str
        :param location: Location code. Supports ``fnmatch`` wildcards, e.g.
            ``*`` and ``?``.
        :type channel: str
        :param channel: Channel code. Supports ``fnmatch`` wildcards, e.g.
            ``*`` and ``?``.
        :type cache: bool
        :param cache: Subsequent function calls are cached, use ``cache=False``
            to force fetching station metadata again from the server.
        :type warn_on_excluded_stations: bool
        :param warn_on_excluded_stations: Whether to show a warning for
            stations that are excluded from the results because the server
            indicates there is no data currently available.
        :rtype: list
        :returns: list of 2-tuples (or 4-tuples with ``level='channel'``) with
            network/station (network/station/location/channel, respectively)
            code combinations for which data is served by the server.
        """
        if level not in ('station', 'channel'):
            msg = "Invalid option for 'level': '%s'" % str(level)
            raise ValueError(msg)
        if level == 'station' and \
                any(x is not None for x in (location, channel)):
            msg = ("location and channel options are ignored in get_info() if "
                   "level='station'.")
            warnings.warn(msg)
        # deteremine if we have a usable cache and check if it is at least the
        # requested level of detail
        if cache and self._station_cache is not None \
                and level in ('station', self._station_cache_level):
            if level == 'station':
                if self._station_cache_level == 'station':
                    info = [(net, sta) for net, sta in self._station_cache
                            if fnmatch.fnmatch(net, network or '*') and
                            fnmatch.fnmatch(sta, station or '*')]
                    return sorted(info)
                else:
                    info = [(net, sta) for net, sta, loc, cha
                            in self._station_cache
                            if fnmatch.fnmatch(net, network or '*') and
                            fnmatch.fnmatch(sta, station or '*')]
                    return sorted(set(info))
            info = [(net, sta, loc, cha) for net, sta, loc, cha in
                    self._station_cache if
                    fnmatch.fnmatch(net, network or '*') and
                    fnmatch.fnmatch(sta, station or '*') and
                    fnmatch.fnmatch(loc, location or '*') and
                    fnmatch.fnmatch(cha, channel or '*')]
            return sorted(info)

        self._init_client()
        if level == 'station':
            self._slclient.infolevel = "STATIONS"
        elif level == 'channel':
            self._slclient.infolevel = "STREAMS"
        self._slclient.verbose = 1
        self._connect()
        self._slclient.initialize()
        # self._slclient.run()
        self._slclient.run(packet_handler=self._packet_handler)
        info = self._slclient.slconn.info_string
        try:
            xml = etree.fromstring(info)
        except ValueError as e:
            msg = 'Unicode strings with encoding declaration are not supported'
            if msg not in str(e):
                raise
            parser = etree.XMLParser(encoding='utf-8')
            xml = etree.fromstring(info.encode('utf-8'), parser=parser)
        station_cache = set()
        excluded_stations = set()
        for tag in xml.xpath('./station'):
            net = tag.attrib['network']
            sta = tag.attrib['name']
            item = (net, sta)
            if level == 'channel':
                subtags = tag.xpath('./stream')
                # If no data is in ring buffer (e.g. station outage?) then it
                # seems the seedlink server replies with no subtags for the
                # channels
                if not subtags:
                    excluded_stations.add(item)
                    continue
                for subtag in subtags:
                    loc = subtag.attrib['location']
                    cha = subtag.attrib['seedname']
                    station_cache.add(item + (loc, cha))
            elif level == 'station':
                # remove stations that seem to have no data
                if all(tag.attrib[key] == '000000'
                       for key in ('begin_seq', 'end_seq')):
                    excluded_stations.add(item)
                    continue
                station_cache.add(item)
            else:
                raise NotImplementedError()
        # change results to an Inventory object
        self._station_cache = station_cache
        self._station_cache_level = level
        if warn_on_excluded_stations and excluded_stations:
            msg = ('Some stations were excluded from results because server '
                   'indicates no data available (use debug=True in Client '
                   'initialization for details, suppress warning with '
                   'warn_on_excluded_stations=False): ')
            msg += ', '.join('.'.join(item)
                             for item in sorted(excluded_stations))
            warnings.warn(msg)
        return self.get_info(
            network=network, station=station, location=location,
            channel=channel, cache=True, level=level)

    def _packet_handler(self, count, slpack):
        """
        Custom packet handler that accumulates all waveform packets in a
        stream.
        """
        # check if not a complete packet
        if slpack is None or (slpack == SLPacket.SLNOPACKET) or \
                (slpack == SLPacket.SLERROR):
            return False

        # get basic packet info
        type_ = slpack.get_type()
        if self.debug:
            print(type_)

        # process INFO packets here
        if type_ == SLPacket.TYPE_SLINF:
            if self.debug:
                print(SLPacket.TYPE_SLINF)
            return False
        elif type_ == SLPacket.TYPE_SLINFT:
            if self.debug:
                print("Complete INFO:",
                      self._slclient.slconn.get_info_string())
            return True

        # process packet data
        trace = slpack.get_trace()
        if trace is None:
            if self.debug:
                print("Blockette contains no trace")
            return False

        # new samples add to the main stream which is then trimmed
        self.stream += trace
        self.stream.merge(-1)
        return False

    def _fast_packet_handler(self, count, slpack):
        """
        Custom packet handler that accumulates all received data in a
        list of numpy arrays. These can then be converted in one go,
        reducing the processing time up to 90% compared to
        _packet_handler().
        """
        # check if not a complete packet
        if slpack is None or (slpack == SLPacket.SLNOPACKET) or \
                (slpack == SLPacket.SLERROR):
            return False

        # get basic packet info
        type_ = slpack.get_type()
        if self.debug:
            print(type_)

        # process INFO packets here
        if type_ == SLPacket.TYPE_SLINF:
            if self.debug:
                print(SLPacket.TYPE_SLINF)
            return False
        elif type_ == SLPacket.TYPE_SLINFT:
            if self.debug:
                print("Complete INFO:",
                      self._slclient.slconn.get_info_string())
            return True

        # process packet data
        msr_record = slpack.msrecord
        if msr_record is None:
            if self.debug:
                print("Blockette contains no trace")
            return False

        # add to bulk_data_container
        self.bulk_data_container.append(np.frombuffer(msr_record,
                                        dtype=np.int8))
        return False


if __name__ == '__main__':
    import doctest
    doctest.testmod(exclude_empty=True)
