# 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.

import sys
from inspect import isawaitable

from newrelic.api.function_trace import FunctionTrace, function_trace
from newrelic.api.time_trace import notice_error
from newrelic.api.transaction import current_transaction
from newrelic.api.web_transaction import web_transaction
from newrelic.common.object_names import callable_name
from newrelic.common.object_wrapper import function_wrapper, wrap_function_wrapper


def _bind_add(uri, methods, handler, *args, **kwargs):
    return uri, methods, handler, args, kwargs


@function_wrapper
def _nr_wrapper_handler_(wrapped, instance, args, kwargs):
    transaction = current_transaction()

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

    name = callable_name(wrapped)
    view_class = getattr(wrapped, "view_class", None)
    view = view_class or wrapped
    if view_class:
        try:
            method = args[0].method.lower()
            name = callable_name(view_class) + "." + method
            view = getattr(view_class, method)
        except:
            pass

    transaction.set_transaction_name(name, priority=3)
    import sanic

    transaction.add_framework_info(name="Sanic", version=sanic.__version__)

    with FunctionTrace(name=name, source=view):
        return wrapped(*args, **kwargs)


@function_wrapper
def _nr_sanic_router_add(wrapped, instance, args, kwargs):
    uri, methods, handler, args, kwargs = _bind_add(*args, **kwargs)

    # Cache the callable_name on the handler object
    callable_name(handler)
    if hasattr(wrapped, "view_class"):
        callable_name(wrapped.view_class)
    wrapped_handler = _nr_wrapper_handler_(handler)

    return wrapped(uri, methods, wrapped_handler, *args, **kwargs)


@function_wrapper
def _nr_sanic_router_get(wrapped, instance, args, kwargs):
    # Rename all transactions that generate an exception in the router get
    try:
        return wrapped(*args, **kwargs)
    except Exception:
        transaction = current_transaction()
        if transaction:
            name = callable_name(wrapped)
            transaction.set_transaction_name(name, priority=2)
        raise


@function_wrapper
def _nr_wrapper_error_handler_(wrapped, instance, args, kwargs):
    transaction = current_transaction()

    if not transaction:
        return wrapped(*args, **kwargs)

    name = callable_name(wrapped)
    transaction.set_transaction_name(name, priority=1)
    try:
        response = function_trace(name=name)(wrapped)(*args, **kwargs)
    except:
        notice_error()
        raise

    return response


def _bind_error_add(exception, handler, *args, **kwargs):
    return exception, handler


@function_wrapper
def _nr_sanic_error_handlers(wrapped, instance, args, kwargs):
    exception, handler = _bind_error_add(*args, **kwargs)

    # Cache the callable_name on the handler object
    callable_name(handler)
    wrapped_handler = _nr_wrapper_error_handler_(handler)

    return wrapped(exception, wrapped_handler)


@function_wrapper
def error_response(wrapped, instance, args, kwargs):
    transaction = current_transaction()

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

    exc_info = sys.exc_info()
    try:
        response = wrapped(*args, **kwargs)
    except:
        notice_error(exc_info)
        # We record the original exception and the exception generated by the
        # error handler
        notice_error()
        raise
    else:
        # response can be a response object or a coroutine
        if hasattr(response, "status"):
            notice_error(error=exc_info, status_code=response.status)
        else:
            notice_error(exc_info)
    finally:
        exc_info = None

    return response


def _sanic_app_init(wrapped, instance, args, kwargs):
    result = wrapped(*args, **kwargs)

    error_handler = getattr(instance, "error_handler")
    if hasattr(error_handler, "response"):
        instance.error_handler.response = error_response(error_handler.response)
    if hasattr(error_handler, "add"):
        error_handler.add = _nr_sanic_error_handlers(error_handler.add)

    router = getattr(instance, "router")
    if hasattr(router, "add"):
        router.add = _nr_sanic_router_add(router.add)
    if hasattr(router, "get"):
        # Cache the callable_name on the router.get
        callable_name(router.get)
        router.get = _nr_sanic_router_get(router.get)

    return result


def _nr_sanic_response_get_headers(wrapped, instance, args, kwargs):
    # We call the wrapped function twice to allow the function to mutate the headers
    result = wrapped(*args, **kwargs)
    transaction = current_transaction()

    if transaction is None:
        return result

    # instance is the response object
    cat_headers = transaction.process_response(str(instance.status), instance.headers.items())

    for header_name, header_value in cat_headers:
        if header_name not in instance.headers:
            instance.headers[header_name] = header_value

    return wrapped(*args, **kwargs)


async def _nr_sanic_response_send(wrapped, instance, args, kwargs):
    transaction = current_transaction()
    result = wrapped(*args, **kwargs)
    if isawaitable(result):
        await result

    if transaction is None:
        return result

    # instance is the response object
    cat_headers = transaction.process_response(str(instance.status), instance.headers.items())

    for header_name, header_value in cat_headers:
        if header_name not in instance.headers:
            instance.headers[header_name] = header_value

    return result


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

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

    # instance is the response object
    cat_headers = transaction.process_response(str(instance.status), instance.headers.items())

    for header_name, header_value in cat_headers:
        if header_name not in instance.headers:
            instance.headers[header_name] = header_value

    return wrapped(*args, **kwargs)


def _nr_wrapper_middleware_(attach_to):
    is_request_middleware = attach_to == "request"

    @function_wrapper
    def _wrapper(wrapped, instance, args, kwargs):
        transaction = current_transaction()

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

        name = callable_name(wrapped)
        if is_request_middleware:
            transaction.set_transaction_name(name, priority=2)
        response = function_trace(name=name)(wrapped)(*args, **kwargs)

        return response

    return _wrapper


def _bind_middleware(middleware, attach_to="request", *args, **kwargs):
    return middleware, attach_to


def _nr_sanic_register_middleware_(wrapped, instance, args, kwargs):
    middleware, attach_to = _bind_middleware(*args, **kwargs)

    # Cache the callable_name on the middleware object
    callable_name(middleware)
    middleware_func = middleware
    if hasattr(middleware, "func"):
        name = callable_name(middleware.func)
        middleware_func = middleware.func

    wrapped_middleware = _nr_wrapper_middleware_(attach_to)(middleware_func)
    wrapped(wrapped_middleware, attach_to)
    return middleware


def _bind_request(request, *args, **kwargs):
    return request


def _nr_sanic_transaction_wrapper_(wrapped, instance, args, kwargs):
    request = _bind_request(*args, **kwargs)
    # If the request is a websocket request do not wrap it
    if request.headers.get("upgrade", "").lower() == "websocket":
        return wrapped(*args, **kwargs)

    return web_transaction(
        request_method=request.method,
        request_path=request.path,
        query_string=request.query_string,
        headers=request.headers,
    )(wrapped)(*args, **kwargs)


def _nr_wrap_touchup_run(wrapped, instance, args, kwargs):
    # TouchUp uses metaprogramming to rewrite methods of classes on startup.
    # To properly wrap them we need to catch the call to TouchUp.run and
    # reinstrument any methods that were replaced with uninstrumented versions.

    result = wrapped(*args, **kwargs)

    if "sanic.app" in sys.modules:
        module = sys.modules["sanic.app"]
        target = args[0]

        if isinstance(target, module.Sanic):
            # Reinstrument class after metaclass "TouchUp" has finished rewriting methods on the class.
            target_cls = module.Sanic
            if hasattr(target_cls, "handle_request") and not hasattr(target_cls.handle_request, "__wrapped__"):
                wrap_function_wrapper(module, "Sanic.handle_request", _nr_sanic_transaction_wrapper_)

    return result


def instrument_sanic_app(module):
    wrap_function_wrapper(module, "Sanic.handle_request", _nr_sanic_transaction_wrapper_)
    wrap_function_wrapper(module, "Sanic.__init__", _sanic_app_init)
    wrap_function_wrapper(module, "Sanic.register_middleware", _nr_sanic_register_middleware_)
    if hasattr(module.Sanic, "register_named_middleware"):
        wrap_function_wrapper(module, "Sanic.register_named_middleware", _nr_sanic_register_middleware_)


def instrument_sanic_response(module):
    if hasattr(module.BaseHTTPResponse, "send"):
        wrap_function_wrapper(module, "BaseHTTPResponse.send", _nr_sanic_response_send)
    else:
        if hasattr(module.BaseHTTPResponse, "get_headers"):
            wrap_function_wrapper(module, "BaseHTTPResponse.get_headers", _nr_sanic_response_get_headers)
        if hasattr(module.BaseHTTPResponse, "_parse_headers"):
            wrap_function_wrapper(module, "BaseHTTPResponse._parse_headers", _nr_sanic_response_parse_headers)


def instrument_sanic_touchup_service(module):
    if hasattr(module, "TouchUp") and hasattr(module.TouchUp, "run"):
        wrap_function_wrapper(module, "TouchUp.run", _nr_wrap_touchup_run)
