#  Copyright (C) 2024
#  ABM, Moscow
#
#  UNPUBLISHED PROPRIETARY MATERIAL.
#  ALL RIGHTS RESERVED.
#
#  Authors: Mikhail Mamrov <m.mamrov@abm-jsc.ru>
#
import logging
from dataclasses import dataclass, field
from enum import Enum, unique
from typing import TypeAlias

from aiohttp import hdrs
from http_tools import AbmServiceConnector
from init_helpers.dict_to_dataclass import dict_to_dataclass
from yarl import URL

logger = logging.getLogger(__name__)


def construct_origin_from_referer(referer: str) -> str:
    return str(URL(referer).origin())


TokenId: TypeAlias = str
JWT: TypeAlias = str


@unique
class TokenType(str, Enum):
    ACCESS_TOKEN = "access_token"
    REFRESH_TOKEN = "refresh_token"


@dataclass(frozen=True)
class TokenPair:
    refresh_token: JWT
    access_token: JWT
    type: str = field(init=False, default="Bearer")


@dataclass(frozen=True)
class JWTPortalPayload:
    id: int
    origin: str


@dataclass(frozen=True)
class JWTPortalMemberPayload:
    id: int
    role_id: int
    registered_at: int


@dataclass(frozen=True)
class JWTPayload:
    portal: JWTPortalPayload
    portal_member: JWTPortalMemberPayload
    id: TokenId
    type: TokenType
    iat: int
    exp: int


class AuthServerConnector:
    _SERVER_NAME_HEADER = "server_name"

    Config = AbmServiceConnector.Config

    @dataclass
    class Context(AbmServiceConnector.Context):
        server_name: str

    def __init__(self, config: Config, context: Context) -> None:
        self.config = config
        self.context = context
        self._connector = AbmServiceConnector(config, context)
        logger.info(f"{type(self).__name__} inited")

    async def get_public_key(self) -> str:
        return await self._connector.get(
            path="v1/public_key/get",
            headers={self._SERVER_NAME_HEADER: self.context.server_name},
        )

    async def login(self, login: str, password: str, origin: str) -> TokenPair:
        answer = await self._connector.post(
            path="/v1/login",
            payload={
                "login": login,
                "password": password,
            },
            headers={hdrs.ORIGIN: origin},
        )
        return dict_to_dataclass(answer, TokenPair)

    async def login_as_content_creator_from_admin(self, token: str, portal_id: int, origin: str) -> TokenPair:
        answer = await self._connector.post(
            path="/v1/admin/content_creator/login",
            payload={
                "token": token,
                "portal_id": portal_id,
            },
            headers={hdrs.ORIGIN: origin},
        )
        return dict_to_dataclass(answer, TokenPair)

    async def validate_access_token(self, token: str, origin: str) -> JWTPayload:
        answer = await self._connector.post(
            path="v1/access_token/validate",
            payload=None,
            headers={
                hdrs.ORIGIN: origin,
                hdrs.AUTHORIZATION: f"Bearer {token}",
            },
        )
        return dict_to_dataclass(answer, JWTPayload)

    async def refresh(
        self,
        refresh_token: str,
        origin: str,
        replace_refresh_token: bool = True,
        server_name: str | None = None,
    ) -> TokenPair:
        headers = {hdrs.ORIGIN: origin}
        if server_name is not None:
            headers["server_name"] = server_name

        answer = await self._connector.post(
            path="/v1/refresh",
            payload={
                "refresh_token": refresh_token,
                "replace_refresh_token": replace_refresh_token,
            },
            headers=headers,
        )
        return dict_to_dataclass(answer, TokenPair)

    async def logout(self, refresh_token: str, origin: str) -> None:
        await self._connector.post(
            path="/v1/logout",
            payload={"refresh_token": refresh_token},
            headers={hdrs.ORIGIN: origin},
        )
