"""
karapace -
Custom middleware system on top of `aiohttp` implementing HTTP server and
client components for use in Aiven's REST applications.

Copyright (c) 2023 Aiven Ltd
See LICENSE for details
"""

from accept_types import get_best_match
from collections.abc import Callable
from http import HTTPStatus
from karapace.core.config import Config, create_server_ssl_context
from karapace.statsd import StatsClient
from karapace.core.utils import json_decode, json_encode
from karapace.version import __version__
from typing import NoReturn, overload

import aiohttp
import aiohttp.web
import aiohttp.web_exceptions
import asyncio
import email
import hashlib
import logging
import re
import time

SERVER_NAME = f"Karapace/{__version__}"
JSON_CONTENT_TYPE = "application/json"

SCHEMA_CONTENT_TYPES = [
    "application/vnd.schemaregistry.v1+json",
    "application/vnd.schemaregistry+json",
    JSON_CONTENT_TYPE,
    "application/octet-stream",
]
SCHEMA_ACCEPT_VALUES = [
    "application/vnd.schemaregistry.v1+json",
    "application/vnd.schemaregistry+json",
    JSON_CONTENT_TYPE,
]

# TODO -> accept more general values as well
REST_CONTENT_TYPE_RE = re.compile(
    r"application/((vnd\.kafka(\.(?P<embedded_format>avro|json|protobuf|binary|jsonschema))?(\.(?P<api_version>v[12]))?"
    r"\+(?P<serialization_format>json))|(?P<general_format>json|octet-stream))"
)

REST_ACCEPT_RE = re.compile(
    r"(application|\*)/((vnd\.kafka(\.(?P<embedded_format>avro|json|"
    r"protobuf|binary|jsonschema))?(\.(?P<api_version>v[12]))?\+"
    r"(?P<serialization_format>json))|(?P<general_format>json|\*))"
)


def is_success(http_status: HTTPStatus) -> bool:
    """True if response has a 2xx status_code"""
    return http_status.value >= 200 and http_status.value < 300


def parse_header(header_value: str) -> tuple[str, dict[str, str]]:
    msg = email.message.Message()
    msg["Content-Type"] = header_value
    main_value = msg.get_content_type()
    params = {k: v for k, v in msg.get_params()[1:]}
    return main_value, params


class HTTPRequest:
    def __init__(
        self,
        *,
        url: str,
        query,
        headers: dict[str, str],
        path_for_stats: str,
        method: str,
        content_type: str | None = None,
        accepts: str | None = None,
    ):
        self.url = url
        self.headers = headers
        self._header_cache: dict[str, str | None] = {}
        self.query = query
        self.content_type = content_type
        self.accepts = accepts
        self.path_for_stats = path_for_stats
        self.method = method
        self.json: dict | None = None

    @overload
    def get_header(self, header: str) -> str | None: ...

    @overload
    def get_header(self, header: str, default_value: str) -> str: ...

    def get_header(self, header, default_value=None):
        upper_cased = header.upper()
        if upper_cased in self._header_cache:
            return self._header_cache[upper_cased]
        for h in self.headers.keys():
            if h.upper() == upper_cased:
                value = self.headers[h]
                self._header_cache[upper_cased] = value
                return value
        if upper_cased == "CONTENT-TYPE":
            # sensible default
            self._header_cache[upper_cased] = JSON_CONTENT_TYPE
        else:
            self._header_cache[upper_cased] = default_value
        return self._header_cache[upper_cased]

    def __repr__(self):
        return f"HTTPRequest(url={self.url} query={self.query} method={self.method} json={self.json!r})"


class HTTPResponse(Exception):
    """A custom Response object derived from Exception so it can be raised
    in response handler callbacks."""

    status: HTTPStatus
    json: None | list | dict

    def __init__(
        self,
        body,
        *,
        status: HTTPStatus = HTTPStatus.OK,
        content_type: str | None = None,
        headers: dict[str, str] | None = None,
    ) -> None:
        self.body = body
        self.status = status
        self.headers = dict(headers) if headers else {}

        if isinstance(body, (dict, list)):
            self.headers["Content-Type"] = JSON_CONTENT_TYPE
            self.json = body
        else:
            self.json = None
        if content_type:
            self.headers["Content-Type"] = content_type
        super().__init__(f"HTTPResponse {status.value}")

    def ok(self) -> bool:
        """True if resposne has a 2xx status_code"""
        return is_success(self.status)

    def __repr__(self) -> str:
        return f"HTTPResponse(status={self.status} body={self.body})"


def http_error(message, content_type: str, code: HTTPStatus) -> NoReturn:
    raise HTTPResponse(
        body=json_encode(
            {
                "error_code": code,
                "message": message,
            },
            binary=True,
        ),
        headers={"Content-Type": content_type},
        status=code,
    )


class RestApp:
    def __init__(
        self, *, app_name: str, config: Config, not_ready_handler: Callable[[HTTPRequest], None] | None = None
    ) -> None:
        self.app_name = app_name
        self.config = config
        self.app_request_metric = f"{app_name}_request"
        self.app_shutdown_count_metric = f"{app_name}_shutdown_count"
        self.app = self._create_aiohttp_application(config=config)
        self.log = logging.getLogger(self.app_name)
        self.stats = StatsClient(config=config)
        self.app.on_shutdown.append(self.close_by_app)
        self.not_ready_handler = not_ready_handler

    def _create_aiohttp_application(self, *, config: Config) -> aiohttp.web.Application:
        return aiohttp.web.Application(client_max_size=config.get_max_request_size())

    async def close_by_app(self, app: aiohttp.web.Application) -> None:
        await self.close()

    async def close(self) -> None:
        """Method used to free all the resources allocated by the applicaiton.

        This will be called as a callback by the aiohttp server. It needs to be
        set as hook because the awaitables have to run inside the event loop
        created by the aiohttp library.
        """
        self.log.warning("=======> Received shutdown signal, closing Application <=======")
        self.stats.increase(self.app_shutdown_count_metric)
        self.stats.close()

    @staticmethod
    def cors_and_server_headers_for_request(*, request, origin="*"):
        return {
            "Access-Control-Allow-Origin": origin,
            "Access-Control-Allow-Methods": "DELETE, GET, OPTIONS, POST, PUT",
            "Access-Control-Allow-Headers": "Authorization, Content-Type",
            "Server": SERVER_NAME,
        }

    def check_rest_headers(self, request: HTTPRequest) -> dict:
        method = request.method
        default_content = "application/vnd.kafka.json.v2+json"
        default_accept = "*/*"
        result = {"content_type": default_content}
        content_matcher = REST_CONTENT_TYPE_RE.search(parse_header(request.get_header("Content-Type", default_content))[0])
        accept_matcher = REST_ACCEPT_RE.search(request.get_header("Accept", default_accept))
        if method in {"POST", "PUT"}:
            if not content_matcher:
                http_error(
                    message="HTTP 415 Unsupported Media Type",
                    content_type=result["content_type"],
                    code=HTTPStatus.UNSUPPORTED_MEDIA_TYPE,
                )
        if content_matcher and accept_matcher:
            header_info = content_matcher.groupdict()
            header_info["embedded_format"] = header_info.get("embedded_format") or "binary"
            result["requests"] = header_info
            result["accepts"] = accept_matcher.groupdict()
            return result
        http_error(
            message="HTTP 406 Not Acceptable",
            content_type=result["content_type"],
            code=HTTPStatus.NOT_ACCEPTABLE,
        )

    def check_schema_headers(self, request: HTTPRequest):
        method = request.method
        response_default_content_type = "application/vnd.schemaregistry.v1+json"
        content_type = request.get_header("Content-Type", JSON_CONTENT_TYPE)

        if method in {"POST", "PUT"} and parse_header(content_type)[0] not in SCHEMA_CONTENT_TYPES:
            http_error(
                message="HTTP 415 Unsupported Media Type",
                content_type=response_default_content_type,
                code=HTTPStatus.UNSUPPORTED_MEDIA_TYPE,
            )
        accept_val = request.get_header("Accept")
        if accept_val:
            if accept_val in ("*/*", "*") or accept_val.startswith("*/"):
                return response_default_content_type
            content_type_match = get_best_match(accept_val, SCHEMA_ACCEPT_VALUES)
            if not content_type_match:
                self.log.debug("Unexpected Accept value: %r", accept_val)
                http_error(
                    message="HTTP 406 Not Acceptable",
                    content_type=response_default_content_type,
                    code=HTTPStatus.NOT_ACCEPTABLE,
                )
            return content_type_match
        return response_default_content_type

    async def _handle_request(
        self,
        *,
        request,
        path_for_stats,
        callback,
        schema_request=False,
        callback_with_request=False,
        json_request=False,
        rest_request=False,
        user=None,
    ):
        start_time = time.monotonic()
        resp = None
        rapu_request = HTTPRequest(
            headers=request.headers,
            query=request.query,
            method=request.method,
            url=request.url,
            path_for_stats=path_for_stats,
        )
        try:
            if request.method == "OPTIONS":
                origin = request.headers.get("Origin")
                if not origin:
                    raise HTTPResponse(body="OPTIONS missing Origin", status=HTTPStatus.BAD_REQUEST)
                headers = self.cors_and_server_headers_for_request(request=rapu_request, origin=origin)
                raise HTTPResponse(body=b"", status=HTTPStatus.OK, headers=headers)

            body = await request.read()
            if json_request:
                if not body:
                    raise HTTPResponse(body="Missing request JSON body", status=HTTPStatus.BAD_REQUEST)
                try:
                    _, options = parse_header(rapu_request.get_header("Content-Type"))
                    charset = email.utils.collapse_rfc2231_value(options.get("charset", "utf-8"))
                    body_string = body.decode(charset)
                    rapu_request.json = json_decode(body_string)
                except UnicodeDecodeError:
                    raise HTTPResponse(body=f"Request body is not valid {charset}", status=HTTPStatus.BAD_REQUEST)
                except LookupError:
                    raise HTTPResponse(body=f"Unknown charset {charset}", status=HTTPStatus.BAD_REQUEST)
                except ValueError:
                    raise HTTPResponse(body="Invalid request JSON body", status=HTTPStatus.BAD_REQUEST)

                # Prevent string, int etc. going further from here
                if not isinstance(rapu_request.json, dict):
                    http_error(
                        message="Malformed request",
                        content_type=JSON_CONTENT_TYPE,
                        code=HTTPStatus.BAD_REQUEST,
                    )
            else:
                if body not in {b"", b"{}"}:
                    raise HTTPResponse(body="No request body allowed for this operation", status=HTTPStatus.BAD_REQUEST)

            callback_kwargs = dict(request.match_info)
            if callback_with_request:
                callback_kwargs["request"] = rapu_request

            if rest_request:
                params = self.check_rest_headers(rapu_request)
                if "requests" in params:
                    rapu_request.content_type = params["requests"]
                    params.pop("requests")
                if "accepts" in params:
                    rapu_request.accepts = params["accepts"]
                    params.pop("accepts")
                callback_kwargs.update(params)

            if schema_request:
                content_type = self.check_schema_headers(rapu_request)
                callback_kwargs["content_type"] = content_type

            if user is not None:
                callback_kwargs["user"] = user

            try:
                if self.not_ready_handler is not None:
                    await self.not_ready_handler(rapu_request)
                data = await callback(**callback_kwargs)
                status = HTTPStatus.OK
                headers = {}
            except HTTPResponse as ex:
                data = ex.body
                status = ex.status
                headers = ex.headers
            except asyncio.CancelledError:
                # Re-raise if aiohttp cancelled the task (e.g. client disconnected) without internal server error
                raise
            except Exception:
                self.log.exception("Internal server error")
                headers = {"Content-Type": "application/json"}
                data = {"error_code": HTTPStatus.INTERNAL_SERVER_ERROR.value, "message": "Internal server error"}
                status = HTTPStatus.INTERNAL_SERVER_ERROR
            headers.update(self.cors_and_server_headers_for_request(request=rapu_request))

            if isinstance(data, (dict, list)):
                resp_bytes = json_encode(data, sort_keys=True, binary=True)
            elif isinstance(data, str):
                if "Content-Type" not in headers:
                    headers["Content-Type"] = "text/plain; charset=utf-8"
                resp_bytes = data.encode("utf-8")
            else:
                resp_bytes = data

            # On 204 - NO CONTENT there is no point of calculating cache headers
            if is_success(status):
                if resp_bytes:
                    etag = f'"{hashlib.md5(resp_bytes).hexdigest()}"'
                else:
                    etag = '""'
                if_none_match = request.headers.get("if-none-match")
                if if_none_match and if_none_match.replace("W/", "") == etag:
                    status = HTTPStatus.NOT_MODIFIED
                    resp_bytes = b""

                headers["access-control-expose-headers"] = "etag"
                headers["etag"] = etag

            resp = aiohttp.web.Response(body=resp_bytes, status=status.value, headers=headers)
        except HTTPResponse as ex:
            if isinstance(ex.body, str):
                resp = aiohttp.web.Response(text=ex.body, status=ex.status.value, headers=ex.headers)
            else:
                resp = aiohttp.web.Response(body=ex.body, status=ex.status.value, headers=ex.headers)
        except aiohttp.web_exceptions.HTTPRequestEntityTooLarge:
            # This exception is not our usual http response, so to keep a consistent error interface
            # we construct http response manually here
            status = HTTPStatus.REQUEST_ENTITY_TOO_LARGE
            body = json_encode(
                {
                    "error_code": status,
                    "message": "HTTP Request Entity Too Large",
                },
                binary=True,
            )
            headers = {"Content-Type": "application/json"}
            resp = aiohttp.web.Response(body=body, status=status.value, headers=headers)
        except (ConnectionError, aiohttp.ClientError, asyncio.CancelledError, asyncio.TimeoutError) as exc:
            # TCP level connection errors and timeouts, e.g. TCP reset, client closes connection, task takes too long.
            error_msg = "Unexpected connection or timeout error"
            self.log.debug(error_msg, exc_info=exc)
            # No response can be returned and written to client, aiohttp expects some response here.
            resp = aiohttp.web.Response(text=error_msg, status=HTTPStatus.SERVICE_UNAVAILABLE.value)
        except Exception as ex:
            self.stats.unexpected_exception(ex=ex, where="rapu_wrapped_callback")
            self.log.exception("Unexpected error handling user request: %s %s", request.method, request.url)
            resp = aiohttp.web.Response(text="Internal Server Error", status=HTTPStatus.INTERNAL_SERVER_ERROR.value)
        finally:
            self.stats.timing(
                self.app_request_metric,
                time.monotonic() - start_time,
                tags={
                    "path": path_for_stats,
                    # no `resp` means that we had a failure in exception handler
                    "result": resp.status if resp else 0,
                    "method": request.method,
                },
            )

        return resp

    def route(
        self,
        path,
        *,
        callback,
        method,
        schema_request=False,
        with_request=None,
        json_body=None,
        rest_request=False,
        auth=None,
    ):
        # pretty path for statsd reporting
        path_for_stats = re.sub(r"<[\w:]+>", "x", path)

        # bottle compatible routing
        aio_route = path
        aio_route = re.sub(r"<(\w+):path>", r"{\1:.+}", aio_route)
        aio_route = re.sub(r"<(\w+)>", r"{\1}", aio_route)

        if (method in {"POST", "PUT"}) and with_request is None:
            with_request = True

        if with_request and json_body is None:
            json_body = True

        async def wrapped_callback(request):
            if auth is not None:
                user = auth.authenticate(request)
            else:
                user = None

            return await self._handle_request(
                request=request,
                path_for_stats=path_for_stats,
                callback=callback,
                schema_request=schema_request,
                callback_with_request=with_request,
                json_request=json_body,
                rest_request=rest_request,
                user=user,
            )

        async def wrapped_cors(request):
            return await self._handle_request(
                request=request,
                path_for_stats=path_for_stats,
                callback=None,
            )

        if not aio_route.endswith("/"):
            self.app.router.add_route(method, aio_route + "/", wrapped_callback)
            self.app.router.add_route(method, aio_route, wrapped_callback)
        else:
            self.app.router.add_route(method, aio_route, wrapped_callback)
            self.app.router.add_route(method, aio_route[:-1], wrapped_callback)
        try:
            self.app.router.add_route("OPTIONS", aio_route, wrapped_cors)
        except RuntimeError as ex:
            if "Added route will never be executed, method OPTIONS is already registered" not in str(ex):
                raise

    def run(self) -> None:
        ssl_context = create_server_ssl_context(self.config)

        aiohttp.web.run_app(
            app=self.app,
            host=self.config.host,
            port=self.config.port,
            ssl_context=ssl_context,
            access_log_class=self.config.access_log_class,
            access_log_format='%Tfs %{x-client-ip}i "%r" %s "%{user-agent}i" response=%bb request_body=%{content-length}ib',
        )
