# Copyright (c) 2016 Uber Technologies, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from __future__ import absolute_import

import inspect
import logging
import os
import socket
import sys
from functools import partial

import tornado.gen
import tornado.ioloop
import tornado.iostream
import tornado.tcpserver
from tornado.netutil import bind_sockets

from . import hyperbahn
from ..context import RequestContextProvider
from ..deprecate import deprecate
from ..enum import enum
from ..errors import AlreadyListeningError
from ..event import EventEmitter
from ..event import EventRegistrar
from ..glossary import (
    TCHANNEL_LANGUAGE,
    TCHANNEL_LANGUAGE_VERSION,
    TCHANNEL_VERSION,
)
from ..net import local_ip
from ..schemes import DEFAULT_NAMES
from ..schemes import JSON
from ..serializer.json import JsonSerializer
from ..serializer.raw import RawSerializer
from .connection import StreamConnection
from .connection import INCOMING
from .dispatch import RequestDispatcher
from .peer import PeerGroup

log = logging.getLogger('tchannel')


State = enum(
    'State',
    ready=0,
    closing=1,
    closed=2,
)


class TChannel(object):
    """Manages inbound and outbound connections to various hosts."""

    # TODO deprecate in favor of top-level tchannel.TChannel

    FALLBACK = RequestDispatcher.FALLBACK

    def __init__(self, name, hostport=None, process_name=None,
                 known_peers=None, trace=False, dispatcher=None,
                 reuse_port=False, context_provider_fn=None,
                 _from_new_api=False):
        """Build or re-use a TChannel.

        :param name:
            Name is used to identify client or service itself.

        :param hostport:
            The host-port at which the service behind this TChannel is
            reachable. The port specified in the ``hostport`` is what the
            server will listen on. If unspecified, the system will attempt to
            determine the local network IP for this host and use an
            OS-assigned port.

        :param process_name:
            Name of this process. This is used for logging only. If
            unspecified, this will default to ``$processName[$processId]``.

        :param known_peers:
            A list of host-ports at which already known peers can be reached.
            Defaults to an empty list.

        :param trace:
            Flag to turn on/off zipkin trace. It can be a bool variable or
            a function that return true or false.

        :param context_provider_fn:
            A getter function to retrieve an instance of
            context.RequestContextProvider used to manage
            thread-local request context.
        """

        self._state = State.ready
        if context_provider_fn:
            self.context_provider_fn = context_provider_fn
        else:
            context_provider = RequestContextProvider()
            self.context_provider_fn = lambda: context_provider

        if not dispatcher:
            self._handler = RequestDispatcher(
                context_provider_fn=lambda: self.context_provider_fn())
        else:
            self._handler = dispatcher

        self.peers = PeerGroup(self)

        self._port = 0
        self._host = None
        if hostport:
            self._host, port = hostport.rsplit(':', 1)
            self._port = int(port)

        if not self._host:
            # TChannel(":4040") => determine IP automatically but use port
            # 4040
            self._host = local_ip()

        self.process_name = process_name or "%s[%s]" % (
            sys.argv[0], os.getpid()
        )

        self.name = name
        self._trace = trace

        # register event hooks
        self.event_emitter = EventEmitter()
        self.hooks = EventRegistrar(self.event_emitter)

        if known_peers:
            for peer_hostport in known_peers:
                self.peers.add(peer_hostport)

        # server created from calling listen()
        self._server = None

        # allow SO_REUSEPORT
        self._reuse_port = reuse_port

        # warn if customers are still using this old and soon to be deleted api
        if _from_new_api is False:
            deprecate(
                "tchannel.tornado.TChannel is deprecated and will be removed" +
                " in a a future version - please switch usage to " +
                "tchannel.TChannel object. Thank you."
            )

    @property
    def trace(self):
        if callable(self._trace):
            return self._trace()
        else:
            return self._trace

    def advertise(
        self,
        routers=None,
        name=None,
        timeout=None,
        router_file=None
    ):
        """Make a service available on the Hyperbahn routing mesh.

        This will make contact with a Hyperbahn host from a list of known
        Hyperbahn routers. Additional Hyperbahn connections will be established
        once contact has been made with the network.

        :param router:
            A seed list of addresses of Hyperbahn routers, e.g.,
            ``["127.0.0.1:23000"]``.

        :param name:
            The identity of this service on the Hyperbahn.

            This is usually unnecessary, as it defaults to the name given when
            initializing the :py:class:`TChannel` (which is used as your
            identity as a caller).

        :returns:
            A future that resolves to the remote server's response after
            the first advertise finishes.

            Advertisement will continue to happen periodically.
        """
        name = name or self.name

        if not self.is_listening():
            self.listen()

        return hyperbahn.advertise(self, name, routers, timeout, router_file)

    @property
    def closed(self):
        return self._state == State.closed

    def close(self):
        if self._state in [State.closed, State.closing]:
            return

        self._state = State.closing
        try:
            self.peers.clear()
            if self._server:
                self._server.stop()
        finally:
            self._state = State.closed

    @property
    def host(self):
        return self._host

    @property
    def hostport(self):
        return "%s:%d" % (self._host, self._port)

    @property
    def port(self):
        return self._port

    def request(self,
                hostport=None,
                service=None,
                arg_scheme=None,
                retry=None,
                **kwargs):
        """Initiate a new request through this TChannel.

        :param hostport:
            Host to which the request will be made. If unspecified, a random
            known peer will be picked. This is not necessary if using
            Hyperbahn.

        :param service:
            The name of a service available on Hyperbahn. Defaults to an empty
            string.

        :param arg_scheme:
            Determines the serialization scheme for the request. One of 'raw',
            'json', or 'thrift'. Defaults to 'raw'.

        :param rety:
            One of 'n' (never retry), 'c' (retry on connection errors), 't'
            (retry on timeout), 'ct' (retry on connection errors and timeouts).

            Defaults to 'c'.
        """
        # TODO disallow certain parameters or don't propagate them backwards.
        # For example, blacklist and rank threshold aren't really
        # user-configurable right now.
        return self.peers.request(hostport=hostport,
                                  service=service,
                                  arg_scheme=arg_scheme,
                                  retry=retry,
                                  **kwargs)

    def listen(self, port=None):
        """Start listening for incoming connections.

        A request handler must have already been specified with
        ``TChannel.host``.

        :param port:
            An explicit port to listen on. This is unnecessary when advertising
            on Hyperbahn.

        :returns:
            Returns immediately.

        :raises AlreadyListeningError:
            If listen was already called.
        """

        if self.is_listening():
            raise AlreadyListeningError(
                "listen has already been called"
            )

        if port:
            assert not self._port, "Port has already been set."
            self._port = int(port)

        assert self._handler, "Call .host with a RequestHandler first"
        server = TChannelServer(self)

        bind_sockets_kwargs = {
            'port': self._port,
            # ipv6 causes random address already in use (socket.error w errno
            # == 98) when getaddrinfo() returns multiple values
            # @see https://github.com/uber/tchannel-python/issues/256
            'family': socket.AF_INET,
        }
        if self._reuse_port is True:
            # allow multiple processes to share the same port,
            # this is really useful in a world where services launch N
            # processes per container/os-space, where N is
            # the amount of cpus for example
            bind_sockets_kwargs['reuse_port'] = True

        sockets = bind_sockets(**bind_sockets_kwargs)
        assert sockets, "No sockets bound for port %d" % self._port

        # If port was 0, the OS probably assigned something better.
        self._port = sockets[0].getsockname()[1]

        server.add_sockets(sockets)

        # assign server so we don't listen twice
        self._server = server

    def is_listening(self):

        if self._server:
            return True

        return False

    def receive_call(self, message, connection):

        if not self._handler:
            log.warn(
                "Received %s but a handler has not been defined.", str(message)
            )
            return
        return self._handler.handle(message, connection)

    def _register_simple(self, endpoint, scheme, f):
        """Register a simple endpoint with this TChannel.

        :param endpoint:
            Name of the endpoint being registered.
        :param scheme:
            Name of the arg scheme under which the endpoint will be
            registered.
        :param f:
            Callable handler for the endpoint.
        """
        assert scheme in DEFAULT_NAMES, ("Unsupported arg scheme %s" % scheme)
        if scheme == JSON:
            req_serializer = JsonSerializer()
            resp_serializer = JsonSerializer()
        else:
            req_serializer = RawSerializer()
            resp_serializer = RawSerializer()
        self._handler.register(endpoint, f, req_serializer, resp_serializer)
        return f

    def _register_thrift(self, service_module, handler, **kwargs):
        """Register a Thrift endpoint on this TChannel.

        :param service_module:
            Reference to the Thrift-generated module for the service being
            registered.
        :param handler:
            Handler for the endpoint
        :param method:
            Name of the Thrift method being registered. If omitted, ``f``'s
            name is assumed to be the method name.
        :param service:
            Name of the Thrift service. By default this is determined
            automatically from the module name.
        """
        import tchannel.thrift as thrift
        # Imported inside the function so that we don't have a hard dependency
        # on the Thrift library. This function is usable only if the Thrift
        # library is installed.
        thrift.register(self._handler, service_module, handler, **kwargs)
        return handler

    def register(self, endpoint, scheme=None, handler=None, **kwargs):
        """Register a handler with this TChannel.

        This may be used as a decorator:

        .. code-block:: python

            app = TChannel(name='bar')

            @app.register("hello", "json")
            def hello_handler(request, response):
                params = yield request.get_body()

        Or as a function:

        .. code-block:: python

            # Here we have a Thrift handler for `Foo::hello`
            app.register(Foo, "hello", hello_thrift_handler)

        :param endpoint:
            Name of the endpoint being registered. This should be a reference
            to the Thrift-generated module if this is a Thrift endpoint. It
            may also be ``TChannel.FALLBACK`` if it's intended to be a
            catch-all endpoint.
        :param scheme:
            Name of the scheme under which the endpoint is being registered.
            One of "raw", "json", and "thrift". Defaults to "raw", except if
            "endpoint" was a module, in which case this defaults to "thrift".

        :param handler:
            If specified, this is the handler function. If ignored, this
            function returns a decorator that can be used to register the
            handler function.

        :returns:
            If ``handler`` was specified, this returns ``handler``. Otherwise,
            it returns a decorator that can be applied to a function to
            register it as the handler.
        """
        assert endpoint is not None, "endpoint is required"

        if endpoint is TChannel.FALLBACK:
            decorator = partial(self._handler.register, TChannel.FALLBACK)
            if handler is not None:
                return decorator(handler)
            else:
                return decorator

        if not scheme:
            # scheme defaults to raw, unless the endpoint is a service module.
            if inspect.ismodule(endpoint):
                scheme = "thrift"
            else:
                scheme = "raw"
        scheme = scheme.lower()
        if scheme == 'thrift':
            decorator = partial(self._register_thrift, endpoint, **kwargs)
        else:
            decorator = partial(
                self._register_simple, endpoint, scheme, **kwargs
            )

        if handler is not None:
            return decorator(handler)
        else:
            return decorator


class TChannelServer(tornado.tcpserver.TCPServer):
    __slots__ = ('tchannel',)

    def __init__(self, tchannel):
        super(TChannelServer, self).__init__()
        self.tchannel = tchannel

    @tornado.gen.coroutine
    def handle_stream(self, stream, address):
        log.debug("New incoming connection from %s:%d" % address)

        conn = StreamConnection(
            connection=stream,
            tchannel=self.tchannel,
            direction=INCOMING,
        )

        yield conn.expect_handshake(headers={
            'host_port': self.tchannel.hostport,
            'process_name': self.tchannel.process_name,
            'tchannel_language': TCHANNEL_LANGUAGE,
            'tchannel_language_version': TCHANNEL_LANGUAGE_VERSION,
            'tchannel_version': TCHANNEL_VERSION,
        })

        log.debug(
            "Successfully completed handshake with %s:%s (%s)",
            conn.remote_host,
            conn.remote_host_port,
            conn.remote_process_name)

        self.tchannel.peers.get(
            "%s:%s" % (conn.remote_host,
                       conn.remote_host_port)
        ).register_incoming_conn(conn)

        yield conn.serve(handler=self._handle)

    def _handle(self, message, connection):
        self.tchannel.receive_call(message, connection)
