#  Copyright (C) 2021
#  ABM, Moscow
#
#  UNPUBLISHED PROPRIETARY MATERIAL.
#  ALL RIGHTS RESERVED.
#
#  Authors: Mike Orlov <m.orlov@abm-jsc.ru>
#

from typing import Mapping, Any, Type

from aiohttp import BasicAuth, FormData
from aiohttp.web_exceptions import HTTPServiceUnavailable
from init_helpers import Jsonable
from multidict import CIMultiDict

from http_tools.http_server_connector import HttpServerConnector
from http_tools.request_tracing import CONTEXT_VAR_REQUEST_ID, X_REQUEST_ID


class AbmServiceConnector:
    Config = HttpServerConnector.Config
    Context = HttpServerConnector.Context

    def __init__(
            self, config: Config, context: Context, connector: Type[HttpServerConnector] = HttpServerConnector
    ) -> None:
        self._connector = connector(config, context)

    async def get(self,
                  path: str,
                  url_query: Mapping[str, Any] | None = None,
                  headers: Mapping[str, str] | None = None,
                  auth: BasicAuth | None = None,
                  timeout_sec: float | None = None) -> Any:
        headers = self._prepare_headers(headers)
        answer = await self._connector.get(path, url_query, headers, auth, timeout_sec)
        return self._parse_answer(answer)

    async def head(self,
                   path: str,
                   url_query: Mapping[str, Any] | None = None,
                   headers: Mapping[str, str] | None = None,
                   auth: BasicAuth | None = None,
                   timeout_sec: float | None = None) -> Any:
        headers = self._prepare_headers(headers)
        answer = await self._connector.head(path, url_query, headers, auth, timeout_sec)
        return self._parse_answer(answer)

    async def put(self,
                  path: str,
                  payload: bytes | Jsonable | FormData,
                  url_query: Mapping[str, Any] | None = None,
                  headers: Mapping[str, str] | None = None,
                  auth: BasicAuth | None = None,
                  timeout_sec: float | None = None) -> Any:
        headers = self._prepare_headers(headers)
        if isinstance(payload, (bytes, FormData)):
            answer = await self._connector.put(path, payload, url_query, headers, auth, timeout_sec)
        else:
            answer = await self._connector.put_json(path, payload, url_query, headers, auth, timeout_sec)
        return self._parse_answer(answer)

    async def post(self,
                   path: str,
                   payload: bytes | Jsonable | FormData,
                   url_query: Mapping[str, Any] | None = None,
                   headers: Mapping[str, str] | None = None,
                   auth: BasicAuth | None = None,
                   timeout_sec: float | None = None) -> Any:
        headers = self._prepare_headers(headers)
        if isinstance(payload, (bytes, FormData)):
            answer = await self._connector.post(path, payload, url_query, headers, auth, timeout_sec)
        else:
            answer = await self._connector.post_json(path, payload, url_query, headers, auth, timeout_sec)
        return self._parse_answer(answer)

    async def patch(self,
                    path: str,
                    payload: bytes | Jsonable | FormData,
                    url_query: Mapping[str, Any] | None = None,
                    headers: Mapping[str, str] | None = None,
                    auth: BasicAuth | None = None,
                    timeout_sec: float | None = None) -> Any:
        headers = self._prepare_headers(headers)
        if isinstance(payload, (bytes, FormData)):
            answer = await self._connector.patch(path, payload, url_query, headers, auth, timeout_sec)
        else:
            answer = await self._connector.patch_json(path, payload, url_query, headers, auth, timeout_sec)
        return self._parse_answer(answer)

    async def delete(self,
                     path: str,
                     payload: bytes | Jsonable | FormData,
                     url_query: Mapping[str, Any] | None = None,
                     headers: Mapping[str, str] | None = None,
                     auth: BasicAuth | None = None,
                     timeout_sec: float | None = None) -> Any:
        headers = self._prepare_headers(headers)
        if isinstance(payload, (bytes, FormData)):
            answer = await self._connector.delete(path, payload, url_query, headers, auth, timeout_sec)
        else:
            answer = await self._connector.delete_json(path, payload, url_query, headers, auth, timeout_sec)
        return self._parse_answer(answer)

    def _parse_answer(self, answer: dict[str, Any] | bytes) -> Any:
        if not isinstance(answer, dict):
            return answer
        try:
            if not answer['done']:
                raise HTTPServiceUnavailable(reason=answer['error'])
            return answer.get('result')
        except KeyError as e:
            raise HTTPServiceUnavailable(reason=f'Incorrect answer from {self._connector.config.url}: {answer}') from e

    @staticmethod
    def _prepare_headers(required_headers: Mapping[str, str] | None) -> CIMultiDict:
        headers = CIMultiDict()
        headers.update(required_headers or {})
        if request_id := CONTEXT_VAR_REQUEST_ID.get():
            headers[X_REQUEST_ID] = request_id
        return headers
