"""HTTP server

:copyright: Copyright (c) 2024 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""

from pykern.pkcollections import PKDict
from pykern.pkdebug import pkdc, pkdlog, pkdp, pkdexc, pkdformat
import inspect
import msgpack
import pykern.pkasyncio
import pykern.pkcollections
import pykern.pkconfig
import pykern.quest
import re
import tornado.web


#: Http auth header name
_AUTH_HEADER = "Authorization"

#: http auth header scheme bearer
_AUTH_HEADER_SCHEME_BEARER = "Bearer"

#: POSIT: Matches anything generated by `unique_key`
_UNIQUE_KEY_CHARS_RE = r"\w+"

#: Regex to test format of auth header and extract token
_AUTH_HEADER_RE = re.compile(
    _AUTH_HEADER_SCHEME_BEARER + r"\s+(" + _UNIQUE_KEY_CHARS_RE + ")",
    re.IGNORECASE,
)

_CONTENT_TYPE_HEADER = "Content-Type"
_CONTENT_TYPE = "application/msgpack"

_VERSION_HEADER = "X-PyKern-HTTP-Version"

_VERSION_HEADER_VALUE = "1"

_API_NAME_RE = re.compile(rf"^{pykern.quest.API.METHOD_PREFIX}(\w+)")


def server_start(api_classes, attr_classes, http_config, coros=()):
    """Start the `_HTTPServer` in asyncio

    Args:
        api_classes (Iterable): `pykern.quest.API` subclasses to be dispatched
        attr_classes (Iterable): `pykern.quest.Attr` subclasses to create API instance
        http_config (PKDict): auth_secret and `pkasyncio.Loop.http_server` arg
        coros (Iterable): list of coroutines to be passed to `pkasyncio.Loop.run`
    """
    l = pykern.pkasyncio.Loop()
    _HTTPServer(l, api_classes, attr_classes, http_config)
    if coros:
        l.run(*coros)
    l.start()


class Reply:

    def __init__(self, result=None, exc=None, api_error=None):
        def _exception(exc):
            if exc is None:
                pkdlog("ERROR: no reply and no exception")
                return 500
            if isinstance(exc, APINotFound):
                return 404
            if isinstance(exc, APIForbidden):
                return 403
            pkdlog("untranslated exception={}", exc)
            return 500

        if isinstance(result, Reply):
            self.http_status = result.http_status
            self.content = result.content
        elif result is not None or api_error is not None:
            self.http_status = 200
            self.content = PKDict(
                api_error=api_error,
                api_result=result,
            )
        else:
            self.http_status = _exception(exc)
            self.content = None


class ReplyExc(Exception):
    """Raised to end the request.

    Args:
        pk_args (dict): exception args that specific to this module
        log_fmt (str): server side log data
    """

    def __init__(self, *args, **kwargs):
        super().__init__()
        if "pk_args" in kwargs:
            self.pk_args = kwargs["pk_args"]
            del kwargs["pk_args"]
        else:
            self.pk_args = PKDict()
        if args or kwargs:
            kwargs["pkdebug_frame"] = inspect.currentframe().f_back.f_back
            pkdlog(*args, **kwargs)

    def __repr__(self):
        a = self.pk_args
        return "{}({})".format(
            self.__class__.__name__,
            ",".join(
                ("{}={}".format(k, a[k]) for k in sorted(a.keys())),
            ),
        )

    def __str__(self):
        return self.__repr__()


class APIForbidden(ReplyExc):
    """Raised for forbidden or protocol error"""

    pass


class APINotFound(ReplyExc):
    """Raised for an object not found"""

    pass


class HTTPClient:
    """Wrapper for `tornado.httpclient.AsyncHTTPClient`

    Args:
        http_config (PKDict): tcp_ip, tcp_port, api_uri, auth_secret, request_config

    """

    def __init__(self, http_config):
        self._uri = (
            f"http://{http_config.tcp_ip}:{http_config.tcp_port}{http_config.api_uri}"
        )
        self._headers = PKDict(
            {
                _AUTH_HEADER: f"{_AUTH_HEADER_SCHEME_BEARER} {_auth_secret(http_config.auth_secret)}",
                _CONTENT_TYPE_HEADER: _CONTENT_TYPE,
                _VERSION_HEADER: _VERSION_HEADER_VALUE,
            }
        )
        self._request_config = http_config.get("request_config") or PKDict()

    async def call_api(self, api_name, api_args):
        """Make a request to the API server

        `http_config.request_config` (see `__init__` and if it exists) is passed verbatim to `AsyncHTTPClient.fetch`.

        Args:
            api_name (str): what to call on the server
            api_args (PKDict): passed verbatim to the API on the server.
        Returns:
            str: value of `api_result`.
        Raises:
           APIError: if there was an raise in the API or on a server protocol violation
           Exception: other exceptions that `AsyncHTTPClient.fetch` may raise, e.g. NotFound
        """
        # Need to be careful with the lifecycle of AsyncHTTPClient
        # https://github.com/radiasoft/pykern/issues/529
        r = await tornado.httpclient.AsyncHTTPClient(force_instance=True).fetch(
            self._uri,
            body=_pack_msg(PKDict(api_name=api_name, api_args=api_args)),
            headers=self._headers,
            method="POST",
            **self._request_config,
        )
        rv, e = _unpack_msg(r)
        if e:
            raise pykern.quest.APIError(*e)
        if rv.api_error:
            raise pykern.quest.APIError(
                "api_error={} api_name={} api_args={}", rv.api_error, api_name, api_args
            )
        return rv.api_result


class _HTTPRequestHandler(tornado.web.RequestHandler):
    def initialize(self, server):
        self.pykern_server = server
        self.pykern_context = PKDict()

    async def get(self):
        await self.pykern_server.dispatch(self)

    async def post(self):
        await self.pykern_server.dispatch(self)


class _HTTPServer:

    def __init__(self, loop, api_classes, attr_classes, http_config):
        def _api_class_funcs():
            for c in api_classes:
                for n, o in inspect.getmembers(c, predicate=inspect.isfunction):
                    yield PKDict(api_class=c, api_func=o, api_func_name=n)

        def _api_map():
            rv = PKDict()
            for a in _api_class_funcs():
                n = a.api_func_name
                if not ((m := _API_NAME_RE.search(n)) and n in a.api_class.__dict__):
                    continue
                a.api_name = m.group(1)
                if a.api_name in rv:
                    raise AssertionError(
                        "duplicate api={a.api_name} class={a.api_class.__name__}"
                    )
                if not inspect.iscoroutinefunction(a.pkdel("api_func")):
                    raise AssertionError(
                        "api_func={n} is not async class={a.api_class.__name__}"
                    )
                rv[a.api_name] = a
            return rv

        h = http_config.copy()
        self.loop = loop
        self.api_map = _api_map()
        self.attr_classes = attr_classes
        self.auth_secret = _auth_secret(h.pkdel("auth_secret"))
        h.uri_map = ((h.api_uri, _HTTPRequestHandler, PKDict(server=self)),)
        self.api_uri = h.pkdel("api_uri")
        h.log_function = self._log_end
        self.req_id = 0
        loop.http_server(h)

    async def dispatch(self, handler):
        async def _call(api, api_args):
            with pykern.quest.start(api.api_class, self.attr_classes) as qcall:
                return await getattr(qcall, api.api_func_name)(api_args)

        m = None
        r = None
        try:
            self.req_id += 1
            handler.pykern_context.req_id = self.req_id
            handler.pykern_context.api_name = None
            handler.pykern_context.req_msg = None
            try:
                self._log(handler, "start")
                self._authenticate(handler)
                m, e = _unpack_msg(handler.request)
                if e:
                    raise APIForbidden(*e)
                handler.pykern_context.req_msg = m
                self._log(handler, "call")
                if not (a := self.api_map.get(m.api_name)):
                    raise APINotFound()
                r = Reply(result=await _call(a, m.api_args))
            except pykern.quest.APIError as e:
                self._log(handler, "api_error", e)
                r = Reply(api_error=str(e))
            except Exception as e:
                self._log(handler, "error", e)
                r = Reply(exc=e)
            self._send_reply(handler, r)
        except Exception as e:
            self._log(
                handler,
                "reply_error",
                e,
                getattr(r, "content", None),
            )
            raise

    def _authenticate(self, handler):
        def _token(headers):
            if not (h := headers.get(_AUTH_HEADER)):
                return None
            if m := _AUTH_HEADER_RE.search(h):
                return m.group(1)
            return None

        if handler.request.method != "POST":
            raise APIForbidden()
        if t := _token(handler.request.headers):
            if t == self.auth_secret:
                return
            raise APIForbidden("token mismatch")
        raise APIForbidden("no token")

    def _log(self, handler, which, exc=None, reply=None):
        def _add(key, value):
            nonlocal f, a
            if value is not None:
                f += (" " if f else "") + key + "={}"
                a.append(value)

        f = ""
        a = []
        _add("req_id", handler.pykern_context.get("req_id"))
        _add("api", handler.pykern_context.get("api_name"))
        if exc:
            _add("exception", exc)
            _add("req", handler.pykern_context.get("req_msg"))
            if which != "api_error":
                _add("reply", reply)
                _add("stack", pkdexc())
        self.loop.http_log(handler, which, f, a)

    def _log_end(self, handler):
        self._log(handler, "end")

    def _send_reply(self, handler, reply):
        if (c := reply.content) is None:
            m = b""
        else:
            m = _pack_msg(c)
        handler.set_header(_CONTENT_TYPE_HEADER, _CONTENT_TYPE)
        handler.set_header(_VERSION_HEADER, _VERSION_HEADER_VALUE)
        handler.set_header("Content-Length", str(len(m)))
        handler.set_status(reply.http_status)
        handler.write(m)


def _auth_secret(value):
    if value:
        if len(value) < 16:
            raise AssertionError("secret too short len={len(value)} (<16)")
        return value
    if pykern.pkconfig.in_dev_mode():
        return "default_dev_secret"
    raise AssertionError("must supply http_config.auth_secret")


def _pack_msg(content):
    import datetime

    def _datetime(obj):
        if isinstance(obj, datetime.datetime):
            return int(obj.timestamp())
        return obj

    p = msgpack.Packer(autoreset=False, default=_datetime)
    p.pack(content)
    # TODO(robnagler) getbuffer() would be better
    return p.bytes()


def _unpack_msg(request):
    def _header(name, value):
        if not (v := request.headers.get(name)):
            return ("missing header={}", name)
        if v != value:
            return ("unexpected {}={}", name, c)
        return None

    if e := (
        _header(_VERSION_HEADER, _VERSION_HEADER_VALUE)
        or _header(_CONTENT_TYPE_HEADER, _CONTENT_TYPE)
    ):
        return None, e
    u = msgpack.Unpacker(
        object_pairs_hook=pykern.pkcollections.object_pairs_hook,
    )
    u.feed(request.body)
    return u.unpack(), None
