# Copyright 2010 New Relic, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from newrelic.api.datastore_trace import DatastoreTrace, wrap_datastore_trace
from newrelic.api.transaction import current_transaction
from newrelic.common.object_wrapper import (
    FunctionWrapper,
    wrap_function_wrapper,
    wrap_object,
)


def _instance_info(memcache_host):
    try:
        host = memcache_host.ip
        port_path_or_id = str(memcache_host.port)
    except AttributeError:
        host = "localhost"
        port_path_or_id = str(memcache_host.address)

    return (host, port_path_or_id, None)


def _nr_get_server_wrapper(wrapped, instance, args, kwargs):
    transaction = current_transaction()

    if transaction is None:
        return wrapped(*args, **kwargs)

    result = wrapped(*args, **kwargs)

    instance_info = (None, None, None)

    try:
        tracer_settings = transaction.settings.datastore_tracer
        host = result[0]

        # memcached does not set the db name attribute so there's
        # no need to check the db name reporting setting
        if tracer_settings.instance_reporting.enabled and host is not None:
            instance_info = _instance_info(host)
    except:
        instance_info = ("unknown", "unknown", None)

    transaction._nr_datastore_instance_info = instance_info

    return result


def MemcacheSingleWrapper(wrapped, product, target, operation, module):
    def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
        transaction = current_transaction()

        if transaction is None:
            return wrapped(*args, **kwargs)

        transaction._nr_datastore_instance_info = (None, None, None)

        dt = DatastoreTrace(product, target, operation)

        with dt:
            result = wrapped(*args, **kwargs)

            instance_info = transaction._nr_datastore_instance_info
            (host, port_path_or_id, db) = instance_info
            dt.host = host
            dt.port_path_or_id = port_path_or_id

            return result

    return FunctionWrapper(wrapped, _nr_datastore_trace_wrapper_)


def wrap_memcache_single(module, object_path, product, target, operation):
    wrap_object(module, "Client.%s" % object_path, MemcacheSingleWrapper, (product, target, operation, module))


_memcache_client_methods = (
    "delete",
    "incr",
    "decr",
    "add",
    "append",
    "prepend",
    "replace",
    "set",
    "cas",
    "get",
    "gets",
)

_memcache_multi_methods = ("delete_multi", "get_multi", "set_multi", "get_stats", "get_slabs", "flush_all")


def instrument_memcache(module):
    wrap_function_wrapper(module, "Client._get_server", _nr_get_server_wrapper)

    for name in _memcache_client_methods:
        if hasattr(module.Client, name):
            wrap_memcache_single(module, name, product="Memcached", target=None, operation=name)

    for name in _memcache_multi_methods:
        if hasattr(module.Client, name):
            wrap_datastore_trace(module, "Client.%s" % name, product="Memcached", target=None, operation=name)
