import os
import time
import warnings
from datetime import datetime
from typing import Dict, Optional, Union

# noinspection PyUnresolvedReferences
# jwt package depends on cryptography
import cryptography  # noqa: F401; pylint: disable=unused-import
import jwt
import requests

from yandex.cloud.iam.v1.iam_token_service_pb2 import CreateIamTokenRequest

_MDS_ADDR = "169.254.169.254"
_MDS_URL = "http://{}/computeMetadata/v1/instance/service-accounts/default/token"
_MDS_HEADERS = {"Metadata-Flavor": "Google"}
_MDS_TIMEOUT = (1.0, 1.0)  # 1sec connect, 1sec read

YC_API_ENDPOINT = "api.cloud.yandex.net"


def set_up_yc_api_endpoint(endpoint: str) -> str:
    # pylint: disable-next=global-statement
    global YC_API_ENDPOINT
    YC_API_ENDPOINT = endpoint
    return YC_API_ENDPOINT


def __validate_service_account_key(sa_key: Optional[dict]) -> bool:
    if not isinstance(sa_key, dict):
        raise RuntimeError(f"Invalid Service Account Key: expecting dictionary, actually got {type(sa_key)}")

    obj_id = sa_key.get("id")
    sa_id = sa_key.get("service_account_id")
    private_key = sa_key.get("private_key")

    if not obj_id:
        raise RuntimeError("Invalid Service Account Key: missing key object id.")

    if not sa_id:
        raise RuntimeError("Invalid Service Account Key: missing service account id.")

    if not private_key:
        raise RuntimeError("Invalid Service Account Key: missing private key.")

    private_key_prefix = "-----BEGIN PRIVATE KEY-----"
    if not isinstance(private_key, str) or private_key_prefix not in private_key:
        error_message = (
            "Invalid Service Account Key: private key is in incorrect format."
            f"Should start with {private_key_prefix}.\n"
            "To obtain one you can use YC CLI: yc iam key create --output sa.json --service-account-id <id>"
        )
        raise RuntimeError(error_message)
    return True


class MetadataAuth:
    def __init__(self, metadata_addr: str):
        self.__metadata_addr = metadata_addr

    def url(self) -> str:
        return _MDS_URL.format(self.__metadata_addr)

    def get_token(self) -> str:
        r = requests.get(self.url(), headers=_MDS_HEADERS, timeout=_MDS_TIMEOUT)
        r.raise_for_status()
        response = r.json()
        return response["access_token"]


class TokenAuth:
    def __init__(self, token: str):
        self.__oauth_token = token

    def get_token_request(self) -> "CreateIamTokenRequest":
        return CreateIamTokenRequest(yandex_passport_oauth_token=self.__oauth_token)


class ServiceAccountAuth:
    __SECONDS_IN_HOUR = 60.0 * 60.0

    def __init__(self, sa_key: Dict[str, str], endpoint: Optional[str] = None):
        self.__sa_key = sa_key
        self._endpoint = endpoint if endpoint is not None else YC_API_ENDPOINT

    def get_token_request(self) -> "CreateIamTokenRequest":
        return CreateIamTokenRequest(jwt=self.__prepare_request(self._endpoint))

    def __prepare_request(self, endpoint: str) -> str:
        now = time.time()
        now_utc = datetime.utcfromtimestamp(now)
        exp_utc = datetime.utcfromtimestamp(now + self.__SECONDS_IN_HOUR)
        url = f"https://iam.{endpoint}/iam/v1/tokens"
        payload = {
            "iss": self.__sa_key["service_account_id"],
            "aud": url,
            "iat": now_utc,
            "exp": exp_utc,
        }

        headers = {
            "typ": "JWT",
            "alg": "PS256",
            "kid": self.__sa_key["id"],
        }

        return jwt.encode(payload, self.__sa_key["private_key"], algorithm="PS256", headers=headers)


class IamTokenAuth:
    def __init__(self, iam_token: str):
        self.__iam_token = iam_token

    def get_token(self) -> str:
        return self.__iam_token


def get_auth_token_requester(
    token: Optional[str] = None,
    service_account_key: Optional[dict] = None,
    iam_token: Optional[str] = None,
    metadata_addr: Optional[str] = None,
    endpoint: Optional[str] = None,
) -> Union["MetadataAuth", "TokenAuth", "IamTokenAuth", "ServiceAccountAuth"]:
    if endpoint is None:
        endpoint = YC_API_ENDPOINT
    auth_methods = [("token", token), ("service_account_key", service_account_key), ("iam_token", iam_token)]
    auth_methods = [(auth_type, value) for auth_type, value in auth_methods if value is not None]

    if len(auth_methods) == 0:
        metadata_addr = metadata_addr if metadata_addr is not None else os.environ.get("YC_METADATA_ADDR", _MDS_ADDR)
        return MetadataAuth(metadata_addr)

    if len(auth_methods) > 1:
        raise RuntimeError(f"Conflicting API credentials properties are set: {[auth[0] for auth in auth_methods]}.")

    if token is not None:
        warnings.warn(
            """The 'token' argument corresponds to OAuthToken credential provider, which is deprecated at Yandex Cloud.
            By the end of 2026 will be fully discontinued. Please consider to use another credetials provider.""",
            DeprecationWarning,
            stacklevel=2,
        )
        return TokenAuth(token=token)
    if iam_token is not None:
        return IamTokenAuth(iam_token)
    if service_account_key is not None and __validate_service_account_key(service_account_key):
        return ServiceAccountAuth(service_account_key, endpoint)

    raise RuntimeError("Unknown auth method")
