from __future__ import absolute_import

try:
    # TODO: Find better typing for functions OR change signature
    from typing import Any, Callable, Dict, List, Union

    from werkzeug import exceptions

    # TODO: is there really a need for this anoying import?
    from . import flask_typing as typing
except ImportError:
    pass

import flask
import wrapt

from supertenant import consts
from supertenant.supermeter.data import http_data, http_requests_utils
from supertenant.supermeter.logger import (
    log_instrumentation_success,
    log_integration_module_debug,
    log_integration_module_exception,
)
from supertenant.supermeter.managers import http_manager
from supertenant.supermeter.managers.actions import SyncActions
from supertenant.supermeter.scope_manager import Span


def st_before_request(*argv, **kwargs):
    # type: (Any, Any) -> None
    action = None
    rc = None
    try:
        request = flask.request
        env = request.environ
        before_data = http_data.HTTPServerData(consts.INTEGRATION_MODULE_PYTHON_FLASK)

        before_data.set_method(str(request.method))
        if "PATH_INFO" in env:
            before_data.set_path(env["PATH_INFO"])
        if "QUERY_STRING" in env and len(env["QUERY_STRING"]):
            before_data.set_params(env["QUERY_STRING"])
        if "HTTP_HOST" in env:
            before_data.set_host(env["HTTP_HOST"])

        if "HTTP_USER_AGENT" in env:
            before_data.set_user_agent(env["HTTP_USER_AGENT"])

        before_data.set_headers_from_wsgi_env(env)

        # try to determine the resource ID so we can put it into consts.INTEGRATION_MODULE_RESOURCE_ID
        flask_app_name = flask.current_app.import_name
        before_data.set_tag(consts.LABEL_FLASK_IMPORT_NAME, flask_app_name)
        if flask_app_name and flask_app_name != "__main__":
            before_data.set_integration_module_resource_id(flask_app_name)
        else:
            host = before_data.get_host()
            if host:
                before_data.set_integration_module_resource_id(host)

        span_id, act, poll_key = http_manager.HTTPManager.open_span(before_data)
        if span_id is not None:
            span = Span(span_id, http_data.HTTPServerData(consts.INTEGRATION_MODULE_PYTHON_FLASK))
            action, action_desc = SyncActions.get_action(span_id, act, poll_key)
            if action == consts.ACTION_REJECT:
                assert isinstance(span.finish_data, http_data.HTTPData)
                rc = http_requests_utils.reject_http_request(action, action_desc, span.finish_data)
            # we are going to raise exception and should close the span in teardown request
            flask.g.span = span

    except Exception as exc:
        log_integration_module_exception(consts.INTEGRATION_MODULE_PYTHON_FLASK, "before_request", exc)

    if action == consts.ACTION_REJECT:
        if rc is not None:
            raise flask.abort(rc)  # TODO: oof
        else:
            raise flask.abort(429)
    return None


def st_after_request(response):
    # type: (flask.Response) -> flask.Response
    span = None
    try:
        # If we're not tracing, just return
        if not hasattr(flask.g, "span"):
            return response

        span = flask.g.span
        if span is not None:
            data = span.finish_data
            data.set_status(response.status_code)
    except Exception as exc:
        log_integration_module_exception(consts.INTEGRATION_MODULE_PYTHON_FLASK, "after_request", exc)
    finally:
        if span is not None:
            span.finish()
            flask.g.span = None
    return response


def st_teardown_request(*argv, **kwargs):
    # type: (Any, Any) -> None
    # TODO: this is the type in flask# type: (Union[BaseException, None]) -> Awaitable[None]
    """
    In the case of exceptions, after_request isn't called
    so we capture those cases here.
    """
    if hasattr(flask.g, "span") and flask.g.span is not None:
        if len(argv) > 0 and argv[0] is not None:
            span = flask.g.span
            span.finish_data.mark_error()
            if consts.LABEL_HTTP_STATUS not in span.finish_data.tags:
                span.finish_data.set_status(500)
        flask.g.span.finish()
        flask.g.span = None


@wrapt.patch_function_wrapper("flask", "Flask.handle_user_exception")
def handle_user_exception(wrapped, instance, argv, kwargs):
    # type: (Callable[..., Union[exceptions.HTTPException, typing.ResponseReturnValue]], Any, List[Any], Dict[str, Any]) -> Union[exceptions.HTTPException, typing.ResponseReturnValue] # noqa: E501
    # Call original and then try to do post processing
    response = wrapped(*argv, **kwargs)
    try:
        if hasattr(flask.g, "span") and flask.g.span is not None:
            span = flask.g.span
            data = span.finish_data
            data.mark_error()
            if response is not None:
                if isinstance(response, tuple):
                    status_code = response[1]
                else:
                    if hasattr(response, "code"):
                        status_code = getattr(response, "code")
                    else:
                        status_code = getattr(response, "status_code")

                data.set_status(status_code)

            span.finish()
            flask.g.span = None
    except Exception as exc:
        log_integration_module_exception(consts.INTEGRATION_MODULE_PYTHON_FLASK, "handle_user_exception", exc)

    return response


@wrapt.patch_function_wrapper("flask", "Flask.__init__")
def flask_init(wrapped, instance, argv, kwargs):
    # type: (Callable[..., None], Any, List[Any], Dict[str, Any]) -> None
    res = wrapped(*argv, **kwargs)
    try:
        if not hasattr(instance, "_super_instrumented"):
            log_integration_module_debug(
                consts.INTEGRATION_MODULE_PYTHON_FLASK, "applying before/after instrumentation funcs (vanilla)"
            )
            setattr(instance, "_super_instrumented", True)
            instance.after_request(st_after_request)
            instance.before_request(st_before_request)
            instance.teardown_request(st_teardown_request)
    except Exception as exc:
        log_integration_module_exception(consts.INTEGRATION_MODULE_PYTHON_FLASK, "__init__", exc)
    return res


log_instrumentation_success(consts.INTEGRATION_MODULE_PYTHON_FLASK, getattr(flask, "__version__"))